#include "ReplicatorClient.hpp"

#include "PlayerMove.pb.h"
#include "debug/metrics/NetworkMetrics.hpp"
#include "network/EntityIdMap.hpp"
#include "world/CharacterBody.hpp"
#include "world/JoltPhysicsWorld.hpp"
#include "world/Transform.hpp"
#include <chrono>

namespace tw::net {

ReplicatorClient::ReplicatorClient(
        ProtobufMessages* messages,
        EntityIdMap* entity_id_map,
        dbg::NetworkMetrics *network_metrics,
        JoltPhysicsWorld *physics_world,
        World *world,
        net::ServerConnection* server_connection,
        EntityPositionInterpolator* entity_writer
) :
    m_messages(messages),
    m_entity_id_map(entity_id_map),
    m_network_metrics(network_metrics),
    m_rollback(physics_world),
    m_world(world),
    m_server_connection(server_connection),
    m_input_send_times(INPUT_SEND_TIME_COUNT),
    m_entity_interpolator(entity_writer)
{

}

int64_t ReplicatorClient::now_ms() {
    return std::chrono::duration_cast<std::chrono::milliseconds>(
            Clock::now().time_since_epoch()).count();
}

void ReplicatorClient::set_player_entity(entt::entity entity) {
    m_player_entity = entity;
}

void ReplicatorClient::measure_response_time(uint32_t current_frame_idx, uint32_t snapshot_frame_idx) {
    // Snapshots carry frame zero until the server has an input to answer, and
    // repeat the same frame whenever no newer one arrived in between.
    if(snapshot_frame_idx == 0 || snapshot_frame_idx <= m_last_measured_frame) {
        return;
    }

    // Anything the send times no longer cover, including a frame we never sent,
    // which underflows into a large distance.
    if(current_frame_idx - snapshot_frame_idx >= INPUT_SEND_TIME_COUNT) {
        return;
    }

    m_last_measured_frame = snapshot_frame_idx;

    auto sent_at = m_input_send_times[snapshot_frame_idx % INPUT_SEND_TIME_COUNT];
    m_network_metrics->record_response_time(Clock::now() - sent_at);
}

void ReplicatorClient::set_input(uint32_t frame_idx, glm::vec3 input) {
    m_rollback.set_input(frame_idx, input);

    mmo::PlayerMoveMessage player_move_message = {};

    player_move_message.set_frame_idx(frame_idx);
    mmo::PlayerInput* player_input = new mmo::PlayerInput();
    player_input->set_x(input.x);
    player_input->set_y(input.y);
    player_input->set_z(input.z);

    player_move_message.set_allocated_input(player_input);

    auto send_result = m_messages->send(m_server_connection->server(), player_move_message, false);
    if(!send_result) {
        spdlog::error("Failed to send message: {}", send_result.error().message());
    }

    m_input_send_times[frame_idx % INPUT_SEND_TIME_COUNT] = Clock::now();
}

void ReplicatorClient::record_prediction(uint32_t frame_idx) {
    if(!m_player_entity.has_value()) {
        return;
    }

    Transform* transform = m_world->registry().try_get<Transform>(m_player_entity.value());
    if(!transform) {
        return;
    }

    glm::vec3 position = transform->position();

    m_rollback.record_prediction(frame_idx, position);
}

void ReplicatorClient::update() {
    if(!m_player_entity.has_value()) {
        return;
    }

    // m_entity_interpolator.interpolate_smoothed_entities(std::chrono::high_resolution_clock::now() - std::chrono::milliseconds(100));

    // Transform* transform = m_world->registry().try_get<Transform>(m_player_entity.value());
    // if(!transform) {
    //     return;
    // }
    //
    // transform->set_position(m_render_position);
}

void ReplicatorClient::snap_player_to(uint32_t frame_idx, entt::entity entity, glm::vec3 position) {
    CharacterBody* body = m_world->registry().try_get<CharacterBody>(entity);
    if(body) {
        body->m_character->SetPosition(JPH::RVec3(position.x, position.y, position.z));
        body->m_character->SetLinearVelocity(JPH::Vec3::sZero());
        body->m_desired_velocity = JPH::Vec3::sZero();
    }

    Transform* transform = m_world->registry().try_get<Transform>(entity);
    if(transform) {
        transform->set_position(position);
    }

    m_render_position = position;

    // m_rollback.reset_at(frame_idx);
}

void ReplicatorClient::handle_snapshot_entity(
        uint32_t current_frame_idx,
        uint32_t record_frame_idx,
        serial::EntityRecord& record
) {
    std::optional<entt::entity> entity = m_entity_id_map->get_local(record.id);

    if(!entity.has_value()) {
        spdlog::warn("Received position update for unknown entity {}", record.id);
        return;
    }

    glm::vec3 p = {record.position.x, record.position.y, record.position.z};

    if(m_player_entity.has_value() &&
       entity.value() == m_player_entity.value()
    ) {
        if(!m_player_position_initialized) {
            m_player_position_initialized = true;
            // snap_player_to(current_frame_idx, entity.value(), p);
            // return;
        }

        //snap_player_to(current_frame_idx, entity.value(), p);

        // bool reconcile_happened = m_reconciler.reconcile(record_frame_idx, p, entity.value(), &m_world->registry(), current_frame_idx);
        //
        // if(reconcile_happened) {
        //     m_network_metrics->record_rollback();
        //     m_network_metrics->record_correction_distance(m_reconciler.last_correction_distance());
        //     m_network_metrics->record_replayed_frames(m_reconciler.last_replayed_frames());
        // }

        // if(record_frame_idx != 0) {
        //     uint32_t ack_lag = 0;
        //     if(current_frame_idx >= record_frame_idx) {
        //         ack_lag = current_frame_idx - record_frame_idx;
        //     }
        //     m_network_metrics->record_ack_lag(ack_lag);
        // }
    } 

    m_entity_interpolator->set_position(std::chrono::high_resolution_clock::now(), entity.value(), p);
}

void ReplicatorClient::handle_snapshot(uint32_t current_frame_idx, serial::WorldStateReader& reader) {
    auto header = reader.read_header();
    measure_response_time(current_frame_idx, header.frame_idx);

    while(reader.has_entity()) {
        auto entity_r = reader.read_entity();

        handle_snapshot_entity(current_frame_idx, header.frame_idx, entity_r);
    }
}

}