#include "Character.h"
#include "Animator.h"
#include "CharacterMind.h"
#include "Ragdoll.h"
#include "Physics.h"
#include "tFunctions.h"
#include <cfloat>
#include <cmath>
#include <glm/gtc/quaternion.hpp>

// -------------------------------------------------------
// CharacterBody
// -------------------------------------------------------

void CharacterBody::Perform(const CharacterAction& action) {
    switch (action.type) {
    case CharacterActionType::Move: {
        wishDirection = glm::vec3(action.direction.x, 0.0f, action.direction.z);
        if (glm::length(wishDirection) > 0.0001f) {
            wishDirection = glm::normalize(wishDirection);
            turnDirection = wishDirection;
        }
        break;
    }
    }
}

void CharacterBody::ApplyMovement(bool grounded, float speedFactor, float dt, float referenceDt) {
    const bool moving = glm::length(wishDirection) > 0.0f && speedFactor > 0.0f;
    SetWalking(moving);
    wishDirection = glm::vec3(0.0f);   // used up, the next Move action has to ask again

    // With time stopped there's nothing to apply, and the rescaling below would divide zero by zero
    if (!physics || dt <= 0.0f)
        return;

    // Walks the way it faces rather than straight along the input, so it turns into the direction first
    const float movespeed = walkSpeed * speedFactor;
    const glm::vec3 inputVel = moving ? GetForwardDirection() * movespeed : glm::vec3(0.0f);

    const float accel = 40.0f;
    const float damp = 0.75f;
    const float airControl = 0.5f;
    const float groundTopSpeed = movespeed * accel * referenceDt * damp / (1.0f - damp);

    if (grounded) {
        // Damping and acceleration worked out per reference step, then rescaled so frame rate doesn't change them
        const float dampStep = std::pow(damp, dt / referenceDt);
        const float accelScale = referenceDt * damp * (1.0f - dampStep) / (dt * dampStep * (1.0f - damp));
        physics->linearvelocity.x *= dampStep;
        physics->linearvelocity.z *= dampStep;
        physics->linearvelocity.x += inputVel.x * accel * dt * accelScale;
        physics->linearvelocity.z += inputVel.z * accel * dt * accelScale;
    }
    else if (moving) {
        // In the air it can only steer, up to walking speed
        const glm::vec3 wishDir = glm::normalize(inputVel);
        const float current = glm::dot(glm::vec3(physics->linearvelocity.x, 0.0f, physics->linearvelocity.z), wishDir);
        const float add = glm::clamp(groundTopSpeed - current, 0.0f, accel * airControl * movespeed * dt);
        physics->linearvelocity.x += wishDir.x * add;
        physics->linearvelocity.z += wishDir.z * add;
    }
}

void CharacterBody::Step(float dt) {
    if (glm::length(turnDirection) > 0.0001f) {
        const glm::quat target = LookAt(-turnDirection);
        const glm::quat current = t.GetRotationQuaternion();
        const float blend = 1.0f - std::exp(-turnSpeed * dt);
        t.RotateToQuaternion(glm::normalize(glm::slerp(current, target, blend)));
    }

    const float target = walking ? 1.0f : 0.0f;
    walkBlend += (target - walkBlend) * (1.0f - std::exp(-walkFadeSpeed * dt));

    for (Animator* a : animators)
        a->Step(dt);

    // After stepping, since the animator ramps its own weight up to full when it plays
    if (walkAnimator)
        walkAnimator->blendWeight = walkBlend;

    if (rig && animator->m_CurrentAnimation) {
        blendedBones = animator->GetFinalBoneMatrices();
        BlendBones(animator->m_CurrentAnimation->GetRootNode(), glm::mat4(1.0f));
    }
}

// Blends two parent-relative transforms, position and scale linearly and rotation along the shortest arc
static glm::mat4 BlendLocalTransforms(const glm::mat4& from, const glm::mat4& to, float t) {
    const glm::vec3 fromScale(glm::length(glm::vec3(from[0])), glm::length(glm::vec3(from[1])), glm::length(glm::vec3(from[2])));
    const glm::vec3 toScale(glm::length(glm::vec3(to[0])), glm::length(glm::vec3(to[1])), glm::length(glm::vec3(to[2])));
    const glm::mat3 fromRotation(glm::vec3(from[0]) / fromScale.x, glm::vec3(from[1]) / fromScale.y, glm::vec3(from[2]) / fromScale.z);
    const glm::mat3 toRotation(glm::vec3(to[0]) / toScale.x, glm::vec3(to[1]) / toScale.y, glm::vec3(to[2]) / toScale.z);

    const glm::quat rotation = glm::normalize(glm::slerp(glm::quat_cast(fromRotation), glm::quat_cast(toRotation), t));
    glm::mat4 result = glm::mat4_cast(rotation) * glm::scale(glm::mat4(1.0f), glm::mix(fromScale, toScale, t));
    result[3] = glm::vec4(glm::mix(glm::vec3(from[3]), glm::vec3(to[3]), t), 1.0f);
    return result;
}

// Every node's rest transform relative to the model, and its parent's name, from one clip's hierarchy
static void CollectRestPose(const AssimpNodeData& node, const glm::mat4& parentTransform, const std::string& parentName,
    std::unordered_map<std::string, glm::mat4>& globals, std::unordered_map<std::string, std::string>& parents) {
    const glm::mat4 global = parentTransform * node.transformation;
    globals[node.name] = global;
    parents[node.name] = parentName;
    for (const AssimpNodeData& child : node.children)
        CollectRestPose(child, global, node.name, globals, parents);
}

// A clip exported with only some bones, like an arm on its own, has them hung off whatever ancestor made it into
// the file. Its tracks are then relative to that ancestor instead of the bone's real parent. Going through the rest
// pose of both files turns one into the other, so the arm still rides on the torso.
const std::unordered_map<std::string, glm::mat4>& CharacterBody::GetParentCorrections(const Animation* layer) {
    const Animation* base = animator->m_CurrentAnimation;
    const auto key = std::make_pair(base, layer);
    const auto cached = parentCorrections.find(key);
    if (cached != parentCorrections.end())
        return cached->second;

    std::unordered_map<std::string, glm::mat4>& corrections = parentCorrections[key];
    if (!base || !layer || base == layer)
        return corrections;

    std::unordered_map<std::string, glm::mat4> baseGlobals, layerGlobals;
    std::unordered_map<std::string, std::string> baseParents, layerParents;
    CollectRestPose(base->GetRootNode(), glm::mat4(1.0f), "", baseGlobals, baseParents);
    CollectRestPose(layer->GetRootNode(), glm::mat4(1.0f), "", layerGlobals, layerParents);

    for (const std::string& bone : layer->m_BoneNames) {
        const auto baseParent = baseParents.find(bone);
        const auto layerParent = layerParents.find(bone);
        if (baseParent == baseParents.end() || layerParent == layerParents.end())
            continue;
        if (baseParent->second == layerParent->second)
            continue;

        const auto baseParentGlobal = baseGlobals.find(baseParent->second);
        const auto layerParentGlobal = layerGlobals.find(layerParent->second);
        const glm::mat4 from = layerParentGlobal != layerGlobals.end() ? layerParentGlobal->second : glm::mat4(1.0f);
        const glm::mat4 to = baseParentGlobal != baseGlobals.end() ? baseParentGlobal->second : glm::mat4(1.0f);
        corrections[bone] = glm::inverse(to) * from;
    }
    return corrections;
}

// Walks the base animation's hierarchy. Each bone starts from its rest pose, then every layer with a track for it
// blends in on top in weight order, before it's passed down to its children.
void CharacterBody::BlendBones(const AssimpNodeData& node, const glm::mat4& parentTransform) {
    glm::mat4 local = node.transformation;
    for (Animator* a : animators) {
        if (!a->m_CurrentAnimation || a->blendWeight <= 0.0f)
            continue;
        if (a != animator && !a->isPlaying)
            continue;
        const auto& transforms = a->GetLocalTransforms();
        const auto found = transforms.find(node.name);
        if (found == transforms.end())
            continue;

        glm::mat4 layerLocal = found->second;
        if (a != animator) {
            const auto& corrections = GetParentCorrections(a->m_CurrentAnimation);
            const auto correction = corrections.find(node.name);
            if (correction != corrections.end())
                layerLocal = correction->second * layerLocal;
        }

        local = a == animator || a->blendWeight >= 1.0f
            ? layerLocal
            : BlendLocalTransforms(local, layerLocal, a->blendWeight);
    }

    glm::mat4 global = parentTransform * local;
    if (node.name == "Head")
        global = glm::scale(global, glm::vec3(HeadScale));

    const auto info = rig->m_BoneInfoMap.find(node.name);
    if (info != rig->m_BoneInfoMap.end()) {
        const int index = info->second.id;
        if (index >= 0 && index < (int)blendedBones.size())
            blendedBones[index] = global * info->second.offset;
    }

    for (const AssimpNodeData& child : node.children)
        BlendBones(child, global);
}

void CharacterBody::SetWalking(bool isWalking) {
    walking = isWalking;
}

glm::vec3 CharacterBody::GetForwardDirection() const {
    glm::vec3 forward = -t.GetFrontVector();
    forward.y = 0.0f;
    const float length = glm::length(forward);
    if (length <= 0.0001f)
        return glm::vec3(0.0f, 0.0f, -1.0f);
    return forward / length;
}

// -------------------------------------------------------
// Character
// -------------------------------------------------------

void Character::Think() {
    if (!mind)
        return;
    pendingActions.clear();
    mind->ProduceActions(*this, pendingActions);
    for (const CharacterAction& action : pendingActions)
        body->Perform(action);
}

int Character::GetCollisionGroup() {
    // Kept clear of the ragdoll's groups, which count up from 1
    static int nextGroup = 1000;
    if (collisionGroup == 0)
        collisionGroup = nextGroup++;
    return collisionGroup;
}

void Character::CreatePhysicsBody(Physics& physics, float massKg) {
    if (body->physics)
        return;
    p* physicsBody = new p(body);
    physicsBody->mass = massKg;
    physicsBody->lockRotation = true;   // it turns by facing where it walks, never by tipping over
    physicsBody->friction = 0.0f;       // walking does its own stopping
    physicsBody->collisionGroup = GetCollisionGroup();
    physics.AddObject(physicsBody);
    body->physics = physicsBody;
}

void Character::EnableRagdoll(Physics& physics) {
    if (!ragdoll)
        ragdoll = new Ragdoll(*body, GetRigMatrix(), physics);
}

glm::mat4 Character::GetRigMatrix() const {
    // body->t's scale is the hitbox only; the rig is drawn at visualScale so it fills that hitbox
    return body->t.GetTranslationMatrix() * body->t.GetRotationMatrix()
        * glm::scale(glm::mat4(1.0f), glm::vec3(visualScale));
}

bool Character::GetWorldBounds(BoundingBox& bounds) const {
    if (!body)
        return false;

    if (ragdoll && !ragdoll->GetParts().empty()) {
        bounds.min = glm::vec3(FLT_MAX);
        bounds.max = glm::vec3(-FLT_MAX);
        for (const Ragdoll::Part* part : ragdoll->GetParts()) {
            const BoundingBox partBounds = part->transform.t.GetAABB();
            bounds.min = glm::min(bounds.min, partBounds.min);
            bounds.max = glm::max(bounds.max, partBounds.max);
        }
        return true;
    }

    // The rig is about 2 units across natively in every direction it reaches, arms out included.
    // A cube that size around the body covers any pose and any turn.
    const glm::vec3 reach(visualScale);
    bounds.min = body->t.GetTranslation() - reach;
    bounds.max = body->t.GetTranslation() + reach;
    return true;
}

Character::Character() {
    body = new CharacterBody();
    mind = nullptr;
    shadertype = ShaderType::RigShader;
}

void Character::Render(Shader& shader) {
    if (!body || !body->rig || !body->animator) return;
    const std::vector<glm::mat4> transforms = ragdoll ? ragdoll->GetBoneMatrices()
        : body->animator->m_CurrentAnimation ? body->GetBlendedBoneMatrices()
        : std::vector<glm::mat4>(100, glm::mat4(1.0f));
    shader.SetMat4Array("finalBonesMatrices", transforms.data(), glm::min((int)transforms.size(), 100));
    shader.Set4F("colorOverride", colorOverride);
    surface.Apply(shader, 1, glm::vec3(visualScale));
    shader.Set1F("metallic", metallic);
    shader.Set1F("roughness", roughness);
    shader.Set1F("ao", ao);
    body->rig->Render(shader, ragdoll ? glm::mat4(1.0f) : GetRigMatrix());  // ragdoll bones are already in world space
    shader.Set4F("colorOverride", glm::vec4(0.0f));  // don't leak into other rig draws
}