#include "world.h"
#include "datatypes/vector.h"
#include "enum/part.h"
#include "logger.h"
#include "objects/part/basepart.h"
#include "objects/part/part.h"
#include "objects/service/workspace.h"
#include "physics/convert.h"
#include "timeutil.h"

#include <Jolt/Jolt.h>
#include <Jolt/Core/JobSystemThreadPool.h>
#include <Jolt/Core/TempAllocator.h>
#include <Jolt/Physics/Collision/BroadPhase/BroadPhaseLayer.h>
#include <Jolt/Physics/Collision/ObjectLayer.h>
#include <Jolt/Core/Factory.h>
#include <Jolt/Core/Memory.h>
#include <Jolt/Physics/Body/BodyCreationSettings.h>
#include <Jolt/Physics/Body/BodyInterface.h>
#include <Jolt/Physics/Body/MotionType.h>
#include <Jolt/Physics/Collision/RayCast.h>
#include <Jolt/Physics/Collision/Shape/BoxShape.h>
#include <Jolt/Physics/Collision/Shape/SphereShape.h>
#include <Jolt/Physics/EActivation.h>
#include <Jolt/Physics/PhysicsSettings.h>
#include <Jolt/RegisterTypes.h>
#include <Jolt/Physics/Collision/CollisionCollectorImpl.h>
#include <Jolt/Physics/Collision/CastResult.h>
#include <Jolt/Physics/Collision/Shape/SubShapeID.h>
#include <Jolt/Physics/Body/BodyFilter.h>
#include <Jolt/Physics/Body/BodyLockInterface.h>
#include <Jolt/Physics/Collision/NarrowPhaseQuery.h>
#include <Jolt/Physics/Constraints/FixedConstraint.h>
#include <memory>

static JPH::TempAllocator* allocator;
static JPH::JobSystem* jobSystem;


namespace Layers
{
	static constexpr JPH::ObjectLayer DYNAMIC = 0;
	static constexpr JPH::ObjectLayer ANCHORED = 1;
	static constexpr JPH::ObjectLayer NUM_LAYERS = 2;
};

namespace BPLayers
{
	static constexpr JPH::BroadPhaseLayer ANCHORED(0);
	static constexpr JPH::BroadPhaseLayer DYNAMIC(1);
	static constexpr uint NUM_LAYERS(2);
};

PhysWorld::PhysWorld() {
    worldImpl.Init(4096, 0, 4096, 4096, broadPhaseLayerInterface, objectBroadPhasefilter, objectLayerPairFilter);
    worldImpl.SetGravity(JPH::Vec3(0, -196, 0));
	JPH::PhysicsSettings settings = worldImpl.GetPhysicsSettings();
	// settings.mPointVelocitySleepThreshold = 0.04f; // Fix parts not sleeping
    // settings.mNumVelocitySteps *= 20;
    // settings.mNumPositionSteps *= 20;
	worldImpl.SetPhysicsSettings(settings);
}

PhysWorld::~PhysWorld() {
}

void PhysWorld::addBody(std::shared_ptr<BasePart> part) {
    syncBodyProperties(part);
}

void PhysWorld::removeBody(std::shared_ptr<BasePart> part) {
    // TODO:
}

JPH::Shape* makeShape(std::shared_ptr<BasePart> basePart) {
    if (std::shared_ptr<Part> part = std::dynamic_pointer_cast<Part>(basePart)) {
        switch (part->shape) {
        case PartType::Block:
            return new JPH::BoxShape(convert<JPH::Vec3>(part->size / 2.f), JPH::cDefaultConvexRadius);
        case PartType::Ball:
            return new JPH::SphereShape(glm::min(part->size.X(), part->size.Y(), part->size.Z()) / 2.f);
            break;
        }
    }
    return nullptr;
}

void PhysWorld::syncBodyProperties(std::shared_ptr<BasePart> part) {
    JPH::BodyInterface& interface = worldImpl.GetBodyInterface();

    JPH::EMotionType motionType = part->anchored ? JPH::EMotionType::Static : JPH::EMotionType::Dynamic;
    JPH::EActivation activationMode = part->anchored ? JPH::EActivation::DontActivate : JPH::EActivation::Activate;
    JPH::ObjectLayer objectLayer = part->anchored ? Layers::ANCHORED : Layers::DYNAMIC;

    JPH::Body* body = part->rigidBody.bodyImpl;

    // Generate a new rigidBody
    if (body == nullptr) {
        JPH::Shape* shape = makeShape(part);
        JPH::BodyCreationSettings settings(shape, convert<JPH::Vec3>(part->position()), convert<JPH::Quat>((glm::quat)part->cframe.RotMatrix()), motionType, objectLayer);
        settings.mRestitution = 0.5;

        body = interface.CreateBody(settings);
        body->SetUserData((JPH::uint64)part.get());
        part->rigidBody.bodyImpl = body;

        interface.AddBody(body->GetID(), activationMode);
        interface.SetLinearVelocity(body->GetID(), convert<JPH::Vec3>(part->velocity));
    } else {
        std::shared_ptr<Part> part2 = std::dynamic_pointer_cast<Part>(part);
        bool shouldUpdateShape = (part2 != nullptr && part->rigidBody._lastShape != part2->shape) || part->rigidBody._lastSize == part->size;

        if (shouldUpdateShape) {
            // const JPH::Shape* oldShape = body->GetShape();
            JPH::Shape* newShape = makeShape(part);

            interface.SetShape(body->GetID(), newShape, true, activationMode);
            // Seems like Jolt manages its memory for us, so we don't need the below
            // delete oldShape;
        }

        interface.SetObjectLayer(body->GetID(), objectLayer);        
        interface.SetMotionType(body->GetID(), motionType, activationMode);
        interface.SetPositionRotationAndVelocity(body->GetID(), convert<JPH::Vec3>(part->position()), convert<JPH::Quat>((glm::quat)part->cframe.RotMatrix()), convert<JPH::Vec3>(part->velocity), /* Angular velocity is NYI: */ body->GetAngularVelocity());
    }
}

void physicsInit() {
    JPH::RegisterDefaultAllocator();
    JPH::Factory::sInstance = new JPH::Factory();
    JPH::RegisterTypes();

    allocator = new JPH::TempAllocatorImpl(10 * 1024 * 1024);
    jobSystem = new JPH::JobSystemThreadPool(JPH::cMaxPhysicsJobs, JPH::cMaxPhysicsBarriers, std::thread::hardware_concurrency() - 1);
}

void physicsDeinit() {
    JPH::UnregisterTypes();
    delete JPH::Factory::sInstance;
    JPH::Factory::sInstance = nullptr;
}

tu_time_t physTime;
void PhysWorld::step(float deltaTime) {
    tu_time_t startTime = tu_clock_micros();
    // Depending on the load, it may be necessary to call this with a differing collision step count
    // 5 seems to be a good number supporting the high gravity
    worldImpl.Update(deltaTime, 5, allocator, jobSystem);

    JPH::BodyInterface& interface = worldImpl.GetBodyInterface();
    JPH::BodyIDVector bodyIDs;
    worldImpl.GetBodies(bodyIDs);
    for (JPH::BodyID bodyID : bodyIDs) {
        std::shared_ptr<BasePart> part = ((Instance*)interface.GetUserData(bodyID))->shared<BasePart>();
        part->cframe = CFrame(convert<Vector3>(interface.GetPosition(bodyID)), convert<glm::quat>(interface.GetRotation(bodyID)));
    }

    physTime = tu_clock_micros() - startTime;
}

PhysJoint PhysWorld::createJoint(PhysJointInfo& type, std::shared_ptr<BasePart> part0, std::shared_ptr<BasePart> part1) {
    if (part0->rigidBody.bodyImpl == nullptr
        || part1->rigidBody.bodyImpl == nullptr
        || !part0->workspace()
        || !part1->workspace()
        || part0->workspace()->physicsWorld != shared_from_this()
        || part1->workspace()->physicsWorld != shared_from_this()
    ) { Logger::fatalError("Failed to create joint between two parts due to the call being invalid"); panic(); };

    JPH::TwoBodyConstraint* constraint;
    if (PhysJointGlueInfo* info = dynamic_cast<PhysJointGlueInfo*>(&type)) {
        JPH::FixedConstraintSettings settings;
        settings.mAutoDetectPoint = true; // TODO: Replace this with anchor point
        constraint = settings.Create(*part0->rigidBody.bodyImpl, *part1->rigidBody.bodyImpl);
    } else if (PhysJointWeldInfo* info = dynamic_cast<PhysJointWeldInfo*>(&type)) {
        JPH::FixedConstraintSettings settings;
        settings.mAutoDetectPoint = true; // TODO: Replace this with anchor point
        constraint = settings.Create(*part0->rigidBody.bodyImpl, *part1->rigidBody.bodyImpl);
    } else if (PhysJointSnapInfo* info = dynamic_cast<PhysJointSnapInfo*>(&type)) {
        JPH::FixedConstraintSettings settings;
        settings.mAutoDetectPoint = true; // TODO: Replace this with anchor point
        constraint = settings.Create(*part0->rigidBody.bodyImpl, *part1->rigidBody.bodyImpl);
    } else {
        panic();
    }

    worldImpl.AddConstraint(constraint);
    return { constraint };
}

void PhysWorld::destroyJoint(PhysJoint joint) {
    worldImpl.RemoveConstraint(joint.jointImpl);
}

class PhysRayCastBodyFilter : public JPH::BodyFilter {
    bool ShouldCollideLocked(const JPH::Body &inBody) const override {
        std::shared_ptr<BasePart> part = ((Instance*)inBody.GetUserData())->shared<BasePart>();

        // Ignore specifically "hidden" parts from raycast
        // TODO: Replace this with a better system... Please.
        if (!part->rigidBody.isCollisionsEnabled()) return false;

        return true;
    }
};

std::optional<const RaycastResult> PhysWorld::castRay(Vector3 point, Vector3 rotation, float maxLength, std::optional<RaycastFilter> filter, unsigned short categoryMaskBits) {
    if (filter != std::nullopt) { Logger::fatalError("The filter property of PhysWorld::castRay is not yet implemented"); panic(); };

    const JPH::BodyLockInterface& lockInterface = worldImpl.GetBodyLockInterfaceNoLock();
    const JPH::BodyInterface& interface = worldImpl.GetBodyInterface();
    const JPH::NarrowPhaseQuery& query = worldImpl.GetNarrowPhaseQuery();

    // First we cast a ray to find a matching part
    Vector3 end = point + rotation.Unit() * maxLength;
    JPH::RRayCast ray { convert<JPH::Vec3>(point), convert<JPH::Vec3>(end) };
    JPH::RayCastResult result;
    PhysRayCastBodyFilter bodyFilter;
    bool hitFound = query.CastRay(ray, result, {}, {}, bodyFilter);
    
    // No matches found, return empty
    if (!hitFound) return std::nullopt;

    // Next we cast a ray to find the hit surface and its world position and normal
    JPH::BodyID hitBodyId = result.mBodyID;
    std::shared_ptr<BasePart> part = ((Instance*)interface.GetUserData(hitBodyId))->shared<BasePart>();
    const JPH::Shape* shape = interface.GetShape(hitBodyId);

    // Find the hit position and hence the surface normal of the shape at that specific point
    Vector3 hitPosition = point + rotation.Unit() * (maxLength * result.mFraction);
    JPH::Vec3 surfaceNormal = shape->GetSurfaceNormal(result.mSubShapeID2, convert<JPH::Vec3>(part->cframe.Inverse() * hitPosition));
    Vector3 worldNormal = part->cframe.Rotation() * convert<Vector3>(surfaceNormal);

    return RaycastResult {
        .worldPoint = hitPosition,
        .worldNormal = worldNormal,
        .body = lockInterface.TryGetBody(hitBodyId),
        .hitPart = part,
    };
}

uint BroadPhaseLayerInterface::GetNumBroadPhaseLayers() const {
    return BPLayers::NUM_LAYERS;
}

JPH::BroadPhaseLayer BroadPhaseLayerInterface::GetBroadPhaseLayer(JPH::ObjectLayer inLayer) const {
    switch (inLayer) {
        case Layers::DYNAMIC: return BPLayers::DYNAMIC;
        case Layers::ANCHORED: return BPLayers::ANCHORED;
        default: panic();
    }
}

const char * BroadPhaseLayerInterface::GetBroadPhaseLayerName(JPH::BroadPhaseLayer inLayer) const {
    using T = JPH::BroadPhaseLayer::Type;
    switch ((T)inLayer) {
        case (T)BPLayers::DYNAMIC: return "DYNAMIC";
        case (T)BPLayers::ANCHORED: return "ANCHORED";
        default: panic();
    }
}

bool ObjectBroadPhaseFilter::ShouldCollide(JPH::ObjectLayer inLayer1, JPH::BroadPhaseLayer inLayer2) const {
    return true;
}

bool ObjectLayerPairFilter::ShouldCollide(JPH::ObjectLayer inLayer1, JPH::ObjectLayer inLayer2) const {
    switch (inLayer1) {
    case Layers::DYNAMIC:
        return true;
    case Layers::ANCHORED:
        return inLayer2 == Layers::DYNAMIC;
    default:
        panic();
    }
}