#include "PlayerRollback.hpp"

#include "world/JoltPhysicsWorld.hpp"
#include "world/CharacterBody.hpp"
#include "world/Transform.hpp"
#include <spdlog/spdlog.h>
#include <glm/glm.hpp>
#include <glm/gtx/norm.hpp>

namespace tw::net {

PlayerRollback::PlayerRollback(JoltPhysicsWorld* physics)
    : m_physics(physics), m_last_reconciled_ack(0), m_rollback_count(0),
      m_last_correction_distance(0.0f), m_last_replayed_frames(0), m_last_ack_frame(0)
{
    for (auto& record : m_records) {
        record.frame = 0;
        record.valid = false;
        record.input = glm::vec3(0.0f);
        record.predicted_position = glm::vec3(0.0f);
    }
}

void PlayerRollback::set_input(uint32_t frame, glm::vec3 input) {
    size_t idx = frame % RING_SIZE;
    m_records[idx].frame = frame;
    m_records[idx].valid = true;
    m_records[idx].input = input;
}

void PlayerRollback::record_prediction(uint32_t frame, glm::vec3 position) {
    size_t idx = frame % RING_SIZE;
    if (m_records[idx].frame == frame && m_records[idx].valid) {
        m_records[idx].predicted_position = position;
    }
}

bool PlayerRollback::apply_correction(uint32_t ack_frame, glm::vec3 authoritative_position,
                                  entt::entity player, entt::registry* registry,
                                  uint32_t current_frame)
{
    if (ack_frame == 0 || ack_frame <= m_last_reconciled_ack || ack_frame >= current_frame) {
        return false;
    }

    m_last_reconciled_ack = ack_frame;
    m_last_ack_frame = ack_frame;

    Record& record = m_records[ack_frame % RING_SIZE];
    const bool has_prediction = record.valid && record.frame == ack_frame;

    // With a prediction to compare against, an answer that already matches costs
    // nothing further. This is the case almost every frame.
    if (has_prediction) {
        float distance = glm::distance(record.predicted_position, authoritative_position);
        m_last_correction_distance = distance;

        if (distance < kPositionEpsilon) {
            return false;
        }
    }

    // Restoring the frame the answer describes keeps everything the simulation
    // derived from it, so only the character has to be moved. Without a stored
    // frame there is nothing to restore and the answer is taken as it stands.
    const bool restored = has_prediction && m_physics->rollback(ack_frame);

    place_character(player, registry, authoritative_position, !restored);
    replay_from(ack_frame, player, registry, current_frame);

    m_last_replayed_frames = current_frame - 1 - ack_frame;
    m_rollback_count++;

    spdlog::debug("Corrected at frame {}: distance {}, replayed {}, restored {}",
                  ack_frame, m_last_correction_distance, m_last_replayed_frames, restored);

    return true;
}

void PlayerRollback::place_character(entt::entity player, entt::registry* registry,
                                       glm::vec3 position, bool clear_velocity) {
    CharacterBody* body = registry->try_get<CharacterBody>(player);
    Transform* transform = registry->try_get<Transform>(player);
    if (!body) {
        return;
    }

    body->m_character->SetPosition(JPH::RVec3(position.x, position.y, position.z));
    transform->set_position(position);

    if (clear_velocity) {
        body->m_character->SetLinearVelocity(JPH::Vec3::sZero());
        body->m_desired_velocity = JPH::Vec3::sZero();
    }
}

void PlayerRollback::replay_from(uint32_t from_frame, entt::entity player,
                                   entt::registry* registry, uint32_t current_frame) {
    for (uint32_t f = from_frame + 1; f < current_frame; ++f) {
        m_physics->step(f, tw::JoltPhysicsWorld::FIXED_DELTA_TIME, true);

        Transform* transform = registry->try_get<Transform>(player);
        if (!transform) {
            continue;
        }

        // Frames the ring never saw still need an entry, or the answer to them
        // arrives with nothing to compare against and forces another correction.
        Record& replayed = m_records[f % RING_SIZE];
        replayed.frame = f;
        replayed.valid = true;
        replayed.predicted_position = transform->position();
    }
}

void PlayerRollback::reset_at(uint32_t frame) {
    m_last_reconciled_ack = frame;

    for (auto& record : m_records) {
        record.valid = false;
    }
}

}