#include "ClientWorldController.hpp"

#include <glm/glm.hpp>
#include <chrono>
#include <spdlog/spdlog.h>

#include "Address.hpp"

#include "Entity.pb.h"
#include "Login.pb.h"
#include "WorldState.pb.h"
#include "PlayerMove.pb.h"

#include "entt/entity/entity.hpp"
#include "entt/entity/fwd.hpp"
#include "messages/PlayerMoveMessage.hpp"
#include "metrics/HistoryBuffer.hpp"
#include "world/CharacterBody.hpp"
#include "world/CharacterController.hpp"
#include "world/JoltPhysicsWorld.hpp"
#include "world/WorldEntity.hpp"
#include "tw/serial/WorldStateWriter.hpp"
#include "network/EntityInterpolation.hpp"

namespace tw {

typedef  HistoryBuffer<long, glm::vec3> EntityPositionHistory;

entt::entity create_player_entity(World* world, JoltPhysicsWorld* physics_world, drw::WorldRenderer* renderer) {
    // entt::entity entity = world->registry().create();

    // drw::Mesh mesh = renderer->add_mesh(drw::MeshData::cube(glm::vec3(1.0f)));

    // world->registry()
    //     .emplace<Transform>(entity,
    //             Transform(glm::vec3(0.0f, 10.0f, 0.0f)));

    // world->registry()
    //     .emplace<EntityPositionHistory>(entity, 0, glm::vec3(0.0f, 10.0f, 0.0f), 1000);

    // world->registry()
    //     .emplace<tw::WorldEntity>(entity,
    //             tw::WorldEntity(std::string("player"), (uint32_t)entity));

    // world->registry()
    //     .emplace<tw::drw::Mesh>(entity, mesh);

    // world->registry()
    //     .emplace<tw::CharacterController>(entity, 20.0f);
    // auto* body = &world->registry()
    //     .emplace<CharacterBody>(entity, physics_world->create_character(
    //             new JPH::BoxShape(JPH::Vec3Arg(0.5f, 0.5f, 0.5f)),
    //             glm::vec3(0.0f, 10.0f, 0.0f)));


    // return entity;
    return entt::entity(0);
}


entt::entity
ClientWorldController::create_entity(const std::string& name, glm::vec3 position) {
    const auto entity = m_world->registry().create();

    if(!m_mesh.has_value()) {
        m_mesh = m_world_renderer->add_mesh(drw::MeshData::cube(glm::vec3(1.0f)));
    }
    spdlog::info("Creating entity {}", name);

    m_world->registry()
        .emplace<Transform>(entity,
                Transform(position));

    m_world->registry()
        .emplace<tw::WorldEntity>(entity,
                tw::WorldEntity(name, (uint32_t)entity));

    m_world->registry()
        .emplace<tw::drw::Mesh>(entity,
                m_mesh.value());

    return entity;
}

static std::unique_ptr<msg::MessageEndpoint> create_endpoint() {
    auto endpoint_r = msg::MessageEndpoint::create();
    if(!endpoint_r) {
        spdlog::error("Failed to create the endpoint: {}", endpoint_r.error().message());
        throw std::runtime_error("Failed to create the endpoint");
    }

    return std::move(endpoint_r.value());
}

static msg::MessageConnection* connect_to_server(msg::MessageEndpoint* endpoint, net::Address address) {
    auto server_r = endpoint->connect(address.ip_string(), address.port());
    if(!server_r) {
        spdlog::error("Failed to connect to server: {}", server_r.error().message());
        throw std::runtime_error("Failed to connect to server");
    }

    spdlog::info("Connected to server at {}", address.to_string());

    return server_r.value();
}

std::optional<entt::entity> ClientWorldController::map_from_server_entity(int id) {
    if(m_entity_mapping.contains(id)) {
        return m_entity_mapping[id];
    }

    return {};
}

void ClientWorldController::map_server_entity(int server_id, entt::entity local_id) {
    m_entity_mapping[server_id] = local_id;
}

void ClientWorldController::apply_entity_positions(const mmo::EntityPosition* const* positions, size_t count) {
    for(int i = 0; i < count; i++) {
        const mmo::EntityPosition* const position = positions[i];
        std::optional<entt::entity> entity = map_from_server_entity(position->id());

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

        glm::vec3 p = {position->x(), position->y(), position->z()};
        m_entity_interpolator.add_position_for_entity(entity.value(), p);

        EntityPositionHistory* history = m_world->registry().try_get<EntityPositionHistory>(entity.value());

        if(history != nullptr) {
            auto millis = std::chrono::duration_cast<std::chrono::milliseconds>(Clock::now().time_since_epoch()).count();
            history->set(millis, p);
        }
    }
}

ClientWorldController::ClientWorldController(
        const io::InputManager* inputs,
        World* world,
        JoltPhysicsWorld* physics_world,
        drw::WorldRenderer* world_renderer,
        tw::net::Address address,
        dbg::NetworkMetrics* network_metrics
) :
    m_input_manager(inputs),
    m_world(world),
    m_physics_world(physics_world),
    m_world_renderer(world_renderer),
    // m_player_entity(/* create_player_entity(world, physics_world, world_renderer) */),
    m_player_controller(&world_renderer->camera(), glm::vec3()),
    m_endpoint(create_endpoint()),
    m_server(connect_to_server(m_endpoint.get(), address)),
    m_messages(m_endpoint.get()),
    m_network_metrics(network_metrics),
    m_tick_step(20),
    m_is_connected(false),
    m_input_send_times(INPUT_SEND_TIME_COUNT),
    m_position_history_exporter("/home/martin/output.csv"),
    m_entity_interpolator(&m_world->registry(), (entt::entity)0, 300),
    m_reconciler(physics_world)
{
    m_messages.set_handler<mmo::LoginResponse>(
        [this](msg::PeerId, const mmo::LoginResponse& mesg) {
            if(!m_is_connected) {
                spdlog::info("Logged in!");
            }
        });

    m_messages.set_handler<mmo::SetControlledEntity>(
        [this](msg::PeerId, const mmo::SetControlledEntity& mesg) {
            spdlog::info("Setting controlled entity from server id {}", mesg.entity_id());
            m_controlled_server_id = mesg.entity_id();
            try_bind_player_entity();
        });

    m_endpoint->set_handler(Message<mmo::WorldStateMessage>::value,
        [this](msg::PeerId, std::span<const std::byte> data) {

            serial::WorldStateReader reader(data);

            auto header = reader.read_header();

            measure_response_time(header.frame_idx);

            while(reader.has_spawn()) {
                auto spawn = reader.read_spawn();
                auto entity = create_entity("test", glm::vec3());
                spdlog::info("Spawning entity {}", spawn);

                map_server_entity(spawn, entity);

                if(m_controlled_server_id.has_value() && m_controlled_server_id.value() == spawn) {
                    try_bind_player_entity();
                } else {
                    m_entity_interpolator.register_entity(entity);
                }
            }

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

                std::optional<entt::entity> entity = map_from_server_entity(entity_r.id);

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

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

                if(m_player_entity.has_value() && entity.value() == m_player_entity.value()) {
                    // The entity was created before its position was known, so the
                    // body sits at the origin until the server places it. There is
                    // no predicted history to reconcile against yet.
                    if(!m_player_position_initialized) {
                        m_player_position_initialized = true;
                        snap_player_to(entity.value(), p);
                        continue;
                    }

                    glm::vec3 position_before = glm::vec3(0.0f);
                    Transform* player_transform = m_world->registry().try_get<Transform>(entity.value());
                    if(player_transform) {
                        position_before = player_transform->position();
                    }

                    bool reconcile_happened = m_reconciler.reconcile(header.frame_idx, p, entity.value(), &m_world->registry(), m_frame_idx);

                    if(reconcile_happened) {
                        // Where the replay actually ended up, which is ahead of the
                        // acked position by the frames that were re-simulated.
                        glm::vec3 position_after = player_transform
                            ? player_transform->position()
                            : p;
                        glm::vec3 correction_delta = position_before - position_after;
                        float correction_magnitude = glm::length(correction_delta);
                        if(correction_magnitude > 5.0f) {
                            correction_delta = glm::normalize(correction_delta) * 5.0f;
                        }
                        m_visual_error += correction_delta;

                        m_render_curr_position = position_after;
                        m_render_prev_position = position_after;
                        m_tick_accumulator = 0.0;

                        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(header.frame_idx != 0) {
                        uint32_t ack_lag = 0;
                        if(m_frame_idx >= header.frame_idx) {
                            ack_lag = m_frame_idx - header.frame_idx;
                        }
                        m_network_metrics->record_ack_lag(ack_lag);
                    }
                } else {
                    m_entity_interpolator.add_position_for_entity(entity.value(), p);
                }

                EntityPositionHistory* history = m_world->registry().try_get<EntityPositionHistory>(entity.value());

                if(history != nullptr) {
                    auto millis = std::chrono::duration_cast<std::chrono::milliseconds>(Clock::now().time_since_epoch()).count();
                    history->set(millis, p);
                }
            }

            // apply_entity_positions();
        });

    m_messages.set_handler<mmo::EntitySpawnMessage>(
        [this](msg::PeerId, const mmo::EntitySpawnMessage& mesg) {
            auto entity = create_entity(mesg.name(), glm::vec3());

            map_server_entity(mesg.entity_id(), entity);

            if(m_controlled_server_id.has_value() && m_controlled_server_id.value() == mesg.entity_id()) {
                try_bind_player_entity();
            } else {
                m_entity_interpolator.register_entity(entity);
            }
        });
}

ClientWorldController::~ClientWorldController() {
}

void ClientWorldController::try_bind_player_entity() {
    if(!m_controlled_server_id.has_value()) {
        return;
    }

    auto local_entity = map_from_server_entity(m_controlled_server_id.value());
    if(!local_entity.has_value()) {
        return;
    }

    if(m_player_entity.has_value() && m_player_entity.value() == local_entity.value()) {
        return;
    }

    entt::entity entity = local_entity.value();
    spdlog::info("Binding player entity");

    m_player_entity = entity;

    Transform* transform = m_world->registry().try_get<Transform>(entity);
    glm::vec3 position = transform ? transform->position() : glm::vec3(0.0f);

    m_world->registry().emplace<CharacterController>(entity, 20.0f);
    m_world->registry().emplace<CharacterBody>(entity, m_physics_world->create_character(
        new JPH::BoxShape(JPH::Vec3Arg(0.5f, 0.5f, 0.5f)),
        position
    ));

    if(m_world->registry().all_of<net::EntityPositionInterpolation>(entity)) {
        m_world->registry().remove<net::EntityPositionInterpolation>(entity);
    }
}

void ClientWorldController::snap_player_to(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_prev_position = position;
    m_render_curr_position = position;
    m_tick_accumulator     = 0.0;
    m_visual_error         = glm::vec3(0.0f);

    // Frames simulated before the player was placed describe a position it never
    // actually had, so answers to them must not be reconciled against.
    m_reconciler.reset_at(m_frame_idx);
}

void ClientWorldController::export_entity_history() {
}

void ClientWorldController::measure_response_time(uint32_t 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(frame_idx == 0 || 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(m_frame_idx - frame_idx >= INPUT_SEND_TIME_COUNT) {
        return;
    }

    m_last_measured_frame = frame_idx;

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

void ClientWorldController::update(double delta_time) {
    m_player_controller.update(m_input_manager, delta_time);

    ImGui::Begin("Player Controller");
    if(m_player_entity.has_value()) {
        ImGui::Text("Player entity ID: %d", (uint32_t)m_player_entity.value());
    }

    ImGui::Text("Player count: %ld", m_entity_mapping.size());
    for(auto mapping : m_entity_mapping) {
        ImGui::Text("%d -> %d", (uint32_t)mapping.first, mapping.second);
    }

    ImGui::End();

    if(m_tick_step.update()) {
        auto network_start = Clock::now();
        m_endpoint->update();
        m_network_metrics->record_update_time(Clock::now() - network_start);

        m_network_metrics->sample({
            .bytes_sent        = m_endpoint->bytes_sent(),
            .bytes_received    = m_endpoint->bytes_received(),
            .messages_sent     = m_endpoint->messages_sent(),
            .messages_received = m_endpoint->messages_received()
        });

        if(!m_is_connected && false) {
            return;
        } else {
            glm::vec3 input = m_player_controller.input();

            if(m_player_entity.has_value()) {
                CharacterController* controller = m_world->registry().try_get<CharacterController>(m_player_entity.value());
                if(controller) {
                    controller->set_input(m_frame_idx, input);
                    m_reconciler.record_input(m_frame_idx, input);
                }
            }

            m_physics_world->step(m_frame_idx, JoltPhysicsWorld::FIXED_DELTA_TIME, true);

            if(m_player_entity.has_value()) {
                Transform* player_transform = m_world->registry().try_get<Transform>(m_player_entity.value());
                if(player_transform) {
                    glm::vec3 true_position = player_transform->position();
                    m_reconciler.record_prediction(m_frame_idx, true_position);

                    m_render_prev_position = m_render_curr_position;
                    m_render_curr_position = true_position;
                    m_tick_accumulator = 0.0;
                }
            }

            mmo::PlayerMoveMessage player_move_message = {};

            player_move_message.set_frame_idx(m_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 r = m_messages.send(m_server, player_move_message, false);

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

            export_entity_history();

            m_frame_idx++;

            m_world->registry().view<Transform>()
                .each([&](const auto e, Transform& t) {
                    m_position_history_exporter.write((uint32_t)e,
                        std::chrono::duration_cast<std::chrono::milliseconds>(Clock::now().time_since_epoch()).count(), t.position());
                });
        }
    }

    m_entity_interpolator.update();

    if(m_player_entity.has_value()) {
        Transform* player_transform = m_world->registry().try_get<Transform>(m_player_entity.value());
        if(player_transform) {
            m_visual_error *= std::exp(-delta_time * kVisualErrorDecayRate);
            if(glm::length(m_visual_error) < 0.001f) {
                m_visual_error = glm::vec3(0.0f);
            }

            m_tick_accumulator += delta_time;
            float alpha = glm::clamp(
                static_cast<float>(m_tick_accumulator / JoltPhysicsWorld::FIXED_DELTA_TIME),
                0.0f, 1.0f
            );
            glm::vec3 smoothed_position = glm::mix(m_render_prev_position, m_render_curr_position, alpha) + m_visual_error;
            player_transform->set_position(smoothed_position);

            m_player_controller.set_target(smoothed_position);
        }
    }
}

}
