#include "pch.h"
#include "PhysicsSystem.h"
#include <algorithm>
#include <vector>
#include <limits>
#include <cmath>
#include <spdlog/spdlog.h>

#include <glm/vec2.hpp>
#include "../Model2D.h"
#include "../ModelPart.h"

namespace Vivid2D::util {

    // ==========================================================
    // PhysicsSystem Implementation
    // ==========================================================
    PhysicsSystem::PhysicsSystem() :
        m_gravity(0.0f, -98.0f),
        m_airResistance(0.1f),
        m_timeScale(1.0f),
        m_enabled(true),
        m_initialized(false),
        m_accumulatedTime(0.0f),
        m_maxSubSteps(5),
        m_fixedTimeStep(1.0f / 60.0f),
        m_updateCount(0),
        m_averageUpdateTime(0.0f),
        m_windForce(0.0f, 0.0f),
        m_windEnabled(false) {
    }

    void PhysicsSystem::initialize() {
        if (m_initialized) return;
        reset();
        m_initialized = true;
    }

    void PhysicsSystem::reset() {
        m_particles.clear();
        m_springs.clear();
        m_constraints.clear();
        m_colliders.clear();
        m_lastUpdateTime = std::chrono::steady_clock::now();
        m_accumulatedTime = 0.0f;
        m_updateCount = 0;
        m_averageUpdateTime = 0.0f;
    }

    std::shared_ptr<PhysicsParticle> PhysicsSystem::addParticle(const std::string& id, glm::vec2 position, float mass) {
        auto particle = std::make_shared<PhysicsParticle>(id, position, mass);
        m_particles[id] = particle;
        return particle;
    }

    std::shared_ptr<PhysicsParticle> PhysicsSystem::addParticleFromModelPart(const std::string& id, std::shared_ptr<ModelPart> part, float mass) {
        glm::vec2 position = part->getPosition();
        auto particle = addParticle(id, position, mass);
        particle->setUserData(part.get());
        return particle;
    }

    bool PhysicsSystem::removeParticle(const std::string& id) {
        // Remove related springs and constraints
        m_springs.erase(std::remove_if(m_springs.begin(), m_springs.end(),
            [&](const auto& spring) {
                return spring->getParticleA()->getId() == id || spring->getParticleB()->getId() == id;
            }), m_springs.end());

        m_constraints.erase(std::remove_if(m_constraints.begin(), m_constraints.end(),
            [&](const auto& constraint) {
                return constraint->getParticle()->getId() == id;
            }), m_constraints.end());

        return m_particles.erase(id) > 0;
    }

    std::shared_ptr<PhysicsParticle> PhysicsSystem::getParticle(const std::string& id) {
        auto it = m_particles.find(id);
        return (it != m_particles.end()) ? it->second : nullptr;
    }

    std::shared_ptr<PhysicsSpring> PhysicsSystem::addSpring(const std::string& id, std::shared_ptr<PhysicsParticle> a, std::shared_ptr<PhysicsParticle> b, float restLength, float stiffness, float damping) {
        auto spring = std::make_shared<PhysicsSpring>(id, a, b, restLength, stiffness, damping);
        m_springs.push_back(spring);
        return spring;
    }

    std::shared_ptr<PhysicsSpring> PhysicsSystem::addSpring(const std::string& id, std::shared_ptr<PhysicsParticle> a, std::shared_ptr<PhysicsParticle> b, float stiffness, float damping) {
        float restLength = glm::distance(a->getPosition(), b->getPosition());
        return addSpring(id, a, b, restLength, stiffness, damping);
    }

    bool PhysicsSystem::removeSpring(const std::shared_ptr<PhysicsSpring>& spring) {
        auto it = std::remove(m_springs.begin(), m_springs.end(), spring);
        if (it != m_springs.end()) {
            m_springs.erase(it, m_springs.end());
            return true;
        }
        return false;
    }

    std::shared_ptr<PositionConstraint> PhysicsSystem::addPositionConstraint(std::shared_ptr<PhysicsParticle> particle, glm::vec2 targetPosition) {
        auto constraint = std::make_shared<PositionConstraint>(particle, targetPosition);
        m_constraints.push_back(constraint);
        return constraint;
    }

    std::shared_ptr<DistanceConstraint> PhysicsSystem::addDistanceConstraint(std::shared_ptr<PhysicsParticle> particle, std::shared_ptr<PhysicsParticle> target, float maxDistance) {
        auto constraint = std::make_shared<DistanceConstraint>(particle, target, maxDistance);
        m_constraints.push_back(constraint);
        return constraint;
    }

    bool PhysicsSystem::removeConstraint(const std::shared_ptr<PhysicsConstraint>& constraint) {
        auto it = std::remove(m_constraints.begin(), m_constraints.end(), constraint);
        if (it != m_constraints.end()) {
            m_constraints.erase(it, m_constraints.end());
            return true;
        }
        return false;
    }

    std::shared_ptr<CircleCollider> PhysicsSystem::addCircleCollider(const std::string& id, glm::vec2 center, float radius) {
        auto collider = std::make_shared<CircleCollider>(id, center, radius);
        m_colliders.push_back(collider);
        return collider;
    }

    std::shared_ptr<RectangleCollider> PhysicsSystem::addRectangleCollider(const std::string& id, glm::vec2 center, float width, float height) {
        auto collider = std::make_shared<RectangleCollider>(id, center, width, height);
        m_colliders.push_back(collider);
        return collider;
    }

    bool PhysicsSystem::removeCollider(const std::shared_ptr<PhysicsCollider>& collider) {
        auto it = std::remove(m_colliders.begin(), m_colliders.end(), collider);
        if (it != m_colliders.end()) {
            m_colliders.erase(it, m_colliders.end());
            return true;
        }
        return false;
    }

    void PhysicsSystem::setWindForce(float x, float y) {
        m_windForce = { x, y };
        m_windEnabled = (x != 0.0f || y != 0.0f);
    }

    void PhysicsSystem::setWindForce(const glm::vec2& windForce) {
        setWindForce(windForce.x, windForce.y);
    }

    glm::vec2 PhysicsSystem::getWindForce() const {
        return m_windForce;
    }

    void PhysicsSystem::setWindEnabled(bool enabled) {
        m_windEnabled = enabled;
    }

    bool PhysicsSystem::isWindEnabled() const {
        return m_windEnabled;
    }

    void PhysicsSystem::update(float deltaTime, Model2D& model) {
        if (!m_enabled || !m_initialized) return;

        auto startTime = std::chrono::steady_clock::now();

        float scaledDeltaTime = deltaTime * m_timeScale;
        m_accumulatedTime += scaledDeltaTime;

        int numSubSteps = 0;
        while (m_accumulatedTime >= m_fixedTimeStep && numSubSteps < m_maxSubSteps) {
            updatePhysics(m_fixedTimeStep);
            m_accumulatedTime -= m_fixedTimeStep;
            numSubSteps++;
        }

        applyToModel(model);

        auto endTime = std::chrono::steady_clock::now();
        updatePerformanceStats(std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime).count());
    }

    void PhysicsSystem::updatePhysics(float deltaTime) {
        for (auto const& [id, particle] : m_particles) {
            particle->clearForces();
        }

        applyGravity();
        applyWind();

        for (const auto& spring : m_springs) {
            spring->applyForce(deltaTime);
        }

        for (auto const& [id, particle] : m_particles) {
            if (particle->isMovable()) {
                particle->update(deltaTime);
            }
        }

        int constraintIterations = 3;
        for (int i = 0; i < constraintIterations; ++i) {
            for (const auto& constraint : m_constraints) {
                if (constraint->isEnabled()) {
                    constraint->apply(deltaTime);
                }
            }

            for (const auto& collider : m_colliders) {
                if (!collider->isEnabled()) continue;
                for (auto const& [id, particle] : m_particles) {
                    if (!particle->isMovable()) continue;
                    if (collider->collidesWith(*particle)) {
                        collider->resolveCollision(*particle, deltaTime);
                    }
                }
            }
        }

        syncVelocitiesFromPositions(deltaTime);
        handleParticleCollisions(deltaTime);
        applyAirResistance(deltaTime);
    }


    void PhysicsSystem::syncVelocitiesFromPositions(float deltaTime) {
        if (deltaTime <= 0.0f) return;
        for (auto const& [id, particle] : m_particles) {
            if (!particle->isMovable()) continue;
            glm::vec2 pos = particle->m_position;
            glm::vec2 prev = particle->m_previousPosition;
            particle->m_velocity = (pos - prev) / deltaTime;
        }
    }


    void PhysicsSystem::applyGravity() {
        for (auto const& [id, particle] : m_particles) {
            if (particle->isMovable() && particle->isAffectedByGravity()) {
                glm::vec2 gravityForce = m_gravity * particle->getMass();
                particle->addForce(gravityForce);
            }
        }
    }

    void PhysicsSystem::applyWind() {
        if (!m_windEnabled || glm::length2(m_windForce) == 0.0f) {
            return;
        }
        for (auto const& [id, particle] : m_particles) {
            if (particle->isMovable() && particle->isAffectedByWind()) {
                particle->addForce(m_windForce);
            }
        }
    }

    void PhysicsSystem::applyAirResistance(float deltaTime) {
        if (m_airResistance <= 0.0f) return;
        float factor = 1.0f / (1.0f + m_airResistance * deltaTime);
        for (auto const& [id, particle] : m_particles) {
            if (!particle->isMovable()) continue;
            particle->m_velocity *= factor;
        }
    }

    void PhysicsSystem::handleParticleCollisions(float deltaTime) {
        std::vector<std::shared_ptr<PhysicsParticle>> particleList;
        particleList.reserve(m_particles.size());
        for (auto const& [id, particle] : m_particles) {
            particleList.push_back(particle);
        }

        for (size_t i = 0; i < particleList.size(); ++i) {
            auto p1 = particleList[i];
            if (!p1->isMovable()) continue;

            for (size_t j = i + 1; j < particleList.size(); ++j) {
                auto p2 = particleList[j];
                if (!p2->isMovable()) continue;

                glm::vec2 pos1 = p1->getPosition();
                glm::vec2 pos2 = p2->getPosition();
                glm::vec2 delta = pos2 - pos1;
                float dist2 = glm::length2(delta);
                float minDist = p1->getRadius() + p2->getRadius();

                if (dist2 < minDist * minDist && dist2 > 1e-12f) {
                    float dist = glm::sqrt(dist2);
                    glm::vec2 normal = delta / dist;
                    float penetration = minDist - dist;

                    float invM1 = p1->getInverseMass();
                    float invM2 = p2->getInverseMass();
                    float invSum = invM1 + invM2;
                    if (invSum <= 0.0f) continue;

                    glm::vec2 correction = normal * (penetration / invSum);
                    p1->translatePosition(-correction * invM1);
                    p2->translatePosition(correction * invM2);

                    glm::vec2 v1 = p1->getVelocity();
                    glm::vec2 v2 = p2->getVelocity();
                    glm::vec2 relVel = v2 - v1;
                    float velAlongNormal = glm::dot(relVel, normal);

                    if (velAlongNormal < 0.0f) {
                        float restitution = 0.5f;
                        float j2 = -(1.0f + restitution) * velAlongNormal;
                        j2 /= invSum;

                        glm::vec2 impulse = normal * j2;
                        p1->setVelocity(v1 - impulse * invM1);
                        p2->setVelocity(v2 + impulse * invM2);
                    }
                }
            }
        }
    }

    void PhysicsSystem::applyToModel(Model2D& model) {
        for (auto const& [id, particle] : m_particles) {
            void* userData = particle->getUserData();
            if (userData) {
                auto* part = static_cast<ModelPart*>(userData);
                part->setPosition(particle->getPosition());

                if (glm::length2(particle->getVelocity()) > 0.1f) {
                    float angle = atan2(particle->getVelocity().y, particle->getVelocity().x);
                    part->setRotation(angle);
                }
            }
        }
    }


    void PhysicsSystem::updatePerformanceStats(long long nanoTime) {
        float millis = nanoTime / 1'000'000.0f;
        if (m_updateCount == 0) {
            m_averageUpdateTime = millis;
        }
        else {
            m_averageUpdateTime = m_averageUpdateTime * 0.95f + millis * 0.05f;
        }
        m_updateCount++;
    }

    PhysicsPerformanceReport PhysicsSystem::getPerformanceReport() const {
        return PhysicsPerformanceReport(
            m_particles.size(), m_springs.size(), m_constraints.size(), m_colliders.size(),
            m_averageUpdateTime, m_updateCount
        );
    }

    // Getters and Setters Implementation
    glm::vec2 PhysicsSystem::getGravity() const { return m_gravity; }
    void PhysicsSystem::setGravity(float x, float y) { m_gravity = { x, y }; }
    void PhysicsSystem::setGravity(const glm::vec2& gravity) { m_gravity = gravity; }
    float PhysicsSystem::getAirResistance() const { return m_airResistance; }
    void PhysicsSystem::setAirResistance(float airResistance) { m_airResistance = std::max(0.0f, airResistance); }
    float PhysicsSystem::getTimeScale() const { return m_timeScale; }
    void PhysicsSystem::setTimeScale(float timeScale) { m_timeScale = std::max(0.0f, timeScale); }
    bool PhysicsSystem::isEnabled() const { return m_enabled; }
    void PhysicsSystem::setEnabled(bool enabled) { m_enabled = enabled; }
    bool PhysicsSystem::isInitialized() const { return m_initialized; }
    const std::unordered_map<std::string, std::shared_ptr<PhysicsParticle>>& PhysicsSystem::getParticles() const { return m_particles; }
    const std::vector<std::shared_ptr<PhysicsSpring>>& PhysicsSystem::getSprings() const { return m_springs; }
    const std::vector<std::shared_ptr<PhysicsConstraint>>& PhysicsSystem::getConstraints() const { return m_constraints; }
    const std::vector<std::shared_ptr<PhysicsCollider>>& PhysicsSystem::getColliders() const { return m_colliders; }


    bool PhysicsSystem::hasActivePhysics() {
        if (!m_enabled || !m_initialized) {
            return false;
        }

        bool hasMovableParticles = false;
        for (const auto& pair : m_particles) {
            if (pair.second->isMovable() && pair.second->isAffectedByGravity()) {
                hasMovableParticles = true;
                break;
            }
        }
        if (!hasMovableParticles) {
            return false;
        }

        bool hasActiveSprings = std::any_of(m_springs.begin(), m_springs.end(), [](const auto& spring) {
            return spring->isEnabled() && (spring->getParticleA()->isMovable() || spring->getParticleB()->isMovable());
            });

        bool hasSignificantMotion = false;
        for (const auto& pair : m_particles) {
            const auto& p = pair.second;
            if (!p->isMovable()) continue;

            if (glm::length2(p->getVelocity()) > 0.1f) {
                hasSignificantMotion = true;
                break;
            }
            glm::vec2 positionDelta = p->getPosition() - p->getPreviousPosition();
            if (glm::length2(positionDelta) > 0.001f) {
                hasSignificantMotion = true;
                break;
            }
        }

        bool hasActiveConstraints = std::any_of(m_constraints.begin(), m_constraints.end(), [](const auto& c) {
            return c->isEnabled() && c->getParticle()->isMovable();
            });

        return hasActiveSprings || hasSignificantMotion || hasActiveConstraints;
    }

    // ==========================================================
    // PhysicsParticle Implementation
    // ==========================================================
    PhysicsParticle::PhysicsParticle(std::string id, glm::vec2 position, float mass)
        : m_id(std::move(id)), m_position(position), m_previousPosition(position),
        m_velocity(0.0f), m_acceleration(0.0f), m_forceAccumulator(0.0f),
        m_radius(2.0f), m_movable(true), m_affectedByGravity(true), m_affectedByWind(true), m_userData(nullptr) {
        if (mass <= 0.0f) {
            m_mass = std::numeric_limits<float>::infinity();
            m_inverseMass = 0.0f;
        }
        else {
            m_mass = mass;
            m_inverseMass = 1.0f / mass;
        }
    }

    void PhysicsParticle::update(float deltaTime) {
        if (!m_movable) return;
        m_acceleration = m_forceAccumulator * m_inverseMass;
        m_velocity += m_acceleration * deltaTime;
        m_previousPosition = m_position;
        m_position += m_velocity * deltaTime;
        m_forceAccumulator = { 0.0f, 0.0f };
    }

    void PhysicsParticle::addForce(const glm::vec2& force) {
        if (m_inverseMass == 0.0f) return;
        m_forceAccumulator += force;
    }

    void PhysicsParticle::clearForces() {
        m_forceAccumulator = { 0.0f, 0.0f };
    }

    void PhysicsParticle::translatePosition(const glm::vec2& delta) {
        m_position += delta;
    }

    void PhysicsParticle::setPreviousPosition(const glm::vec2& prev) {
        m_previousPosition = prev;
    }

    void PhysicsParticle::applyNewVelocity(const glm::vec2& newVelocity) {
        m_velocity = newVelocity;
    }

    // ==========================================================
    // PhysicsSpring Implementation
    // ==========================================================
    PhysicsSpring::PhysicsSpring(std::string id, std::shared_ptr<PhysicsParticle> a, std::shared_ptr<PhysicsParticle> b, float restLength, float stiffness, float damping)
        : m_id(std::move(id)), m_particleA(a), m_particleB(b), m_restLength(restLength), m_stiffness(stiffness), m_damping(damping), m_enabled(true) {}

    void PhysicsSpring::applyForce(float deltaTime) {
        if (!m_enabled) return;

        glm::vec2 delta = m_particleB->getPosition() - m_particleA->getPosition();
        float currentLength = glm::length(delta);
        if (currentLength < 0.0001f) return;

        glm::vec2 dir = delta / currentLength;
        float stretch = currentLength - m_restLength;
        glm::vec2 springForce = dir * (m_stiffness * stretch);

        glm::vec2 relativeVelocity = m_particleB->getVelocity() - m_particleA->getVelocity();
        float velocityAlongSpring = glm::dot(relativeVelocity, dir);
        glm::vec2 dampingForce = dir * (m_damping * velocityAlongSpring);

        glm::vec2 totalForce = springForce - dampingForce;

        if (m_particleA->isMovable()) {
            m_particleA->addForce(-totalForce);
        }
        if (m_particleB->isMovable()) {
            m_particleB->addForce(totalForce);
        }
    }

    // ==========================================================
    // Constraints Implementation
    // ==========================================================
    PositionConstraint::PositionConstraint(std::shared_ptr<PhysicsParticle> particle, glm::vec2 targetPosition)
        : m_particle(particle), m_targetPosition(targetPosition), m_strength(0.5f), m_enabled(true) {}

    void PositionConstraint::apply(float deltaTime) {
        if (!m_enabled || !m_particle->isMovable()) return;
        glm::vec2 currentPos = m_particle->getPosition();
        glm::vec2 delta = m_targetPosition - currentPos;
        glm::vec2 correction = delta * m_strength;
        m_particle->setPosition(currentPos + correction);
    }

    DistanceConstraint::DistanceConstraint(std::shared_ptr<PhysicsParticle> particle, std::shared_ptr<PhysicsParticle> target, float maxDistance)
        : m_particle(particle), m_target(target), m_maxDistance(maxDistance), m_enabled(true) {}

    void DistanceConstraint::apply(float deltaTime) {
        if (!m_enabled || !m_particle->isMovable()) return;
        glm::vec2 delta = m_particle->getPosition() - m_target->getPosition();
        float distance = glm::length(delta);

        if (distance > m_maxDistance) {
            glm::vec2 correction = glm::normalize(delta) * (distance - m_maxDistance);
            m_particle->setPosition(m_particle->getPosition() - correction);
        }
    }

    // ==========================================================
    // Colliders Implementation
    // ==========================================================
    CircleCollider::CircleCollider(std::string id, glm::vec2 center, float radius)
        : m_id(std::move(id)), m_center(center), m_radius(radius), m_enabled(true) {}

    bool CircleCollider::collidesWith(const PhysicsParticle& particle) const {
        return glm::distance2(particle.getPosition(), m_center) < (m_radius + particle.getRadius()) * (m_radius + particle.getRadius());
    }

    void CircleCollider::resolveCollision(PhysicsParticle& particle, float deltaTime) {
        glm::vec2 toParticle = particle.getPosition() - m_center;
        float distance = glm::length(toParticle);
        float combined = m_radius + particle.getRadius();
        float overlap = combined - distance;

        if (overlap > 0.0001f) {
            glm::vec2 normal = (distance < 0.0001f) ? glm::vec2(0.0f, 1.0f) : toParticle / distance;

            particle.translatePosition(normal * (overlap + 0.001f));

            glm::vec2 v = particle.getVelocity();
            float dot = glm::dot(v, normal);
            if (dot < 0) {
                float restitution = 0.6f;
                glm::vec2 vPrime = v - normal * ((1.0f + restitution) * dot);
                particle.setVelocity(vPrime);
            }
        }
    }

    RectangleCollider::RectangleCollider(std::string id, glm::vec2 center, float width, float height)
        : m_id(std::move(id)), m_center(center), m_width(width), m_height(height), m_enabled(true) {}

    bool RectangleCollider::collidesWith(const PhysicsParticle& particle) const {
        glm::vec2 particlePos = particle.getPosition();
        float left = m_center.x - m_width / 2.0f - particle.getRadius();
        float right = m_center.x + m_width / 2.0f + particle.getRadius();
        float bottom = m_center.y - m_height / 2.0f - particle.getRadius();
        float top = m_center.y + m_height / 2.0f + particle.getRadius();

        return particlePos.x >= left && particlePos.x <= right &&
            particlePos.y >= bottom && particlePos.y <= top;
    }

    void RectangleCollider::resolveCollision(PhysicsParticle& particle, float deltaTime) {
        glm::vec2 particlePos = particle.getPosition();
        float left = m_center.x - m_width / 2.0f;
        float right = m_center.x + m_width / 2.0f;
        float bottom = m_center.y - m_height / 2.0f;
        float top = m_center.y + m_height / 2.0f;

        float closestX = glm::clamp(particlePos.x, left, right);
        float closestY = glm::clamp(particlePos.y, bottom, top);

        glm::vec2 closestPoint(closestX, closestY);
        glm::vec2 normal = particlePos - closestPoint;

        float dist = glm::length(normal);
        if (dist < 0.0001f) { // Inside rectangle
            // Find minimum exit distance
            float dLeft = particlePos.x - left;
            float dRight = right - particlePos.x;
            float dBottom = particlePos.y - bottom;
            float dTop = top - particlePos.y;

            float min_d = dLeft;
            normal = glm::vec2(-1.0, 0.0);

            if (dRight < min_d) { min_d = dRight; normal = glm::vec2(1.0, 0.0); }
            if (dBottom < min_d) { min_d = dBottom; normal = glm::vec2(0.0, -1.0); }
            if (dTop < min_d) { min_d = dTop; normal = glm::vec2(0.0, 1.0); }

            float overlap = particle.getRadius() + min_d;
            particle.translatePosition(normal * (overlap + 0.001f));

        }
        else { // Outside, but overlapping
            float overlap = particle.getRadius() - dist;
            if (overlap > 0.0001f) {
                normal /= dist;
                particle.translatePosition(normal * (overlap + 0.001f));
            }
            else {
                return; // No collision
            }
        }

        // Velocity response
        glm::vec2 v = particle.getVelocity();
        float dot = glm::dot(v, normal);
        if (dot < 0) {
            float restitution = 0.6f;
            glm::vec2 vPrime = v - normal * ((1.0f + restitution) * dot);
            particle.setVelocity(vPrime);
        }
    }


    // ==========================================================
    // PhysicsPerformanceReport Implementation
    // ==========================================================
    PhysicsPerformanceReport::PhysicsPerformanceReport(size_t pCount, size_t sCount, size_t cCount, size_t colCount, float avgTime, int updates)
        : particleCount(pCount), springCount(sCount), constraintCount(cCount), colliderCount(colCount), averageUpdateTime(avgTime), totalUpdates(updates) {}

    std::string PhysicsPerformanceReport::toString() const {
        return spdlog::fmt_lib::format(
            "Physics Performance: {} particles, {} springs, {} constraints, {} colliders, Avg update: {:.2f}ms, Total updates: {}",
            particleCount, springCount, constraintCount, colliderCount, averageUpdateTime, totalUpdates
        );
    }

} // namespace Vivid2D::util