#include "state_serializer.h"

#include <cassert>

namespace tactical_shooter {

StateSerializer::StateSerializer(const SerializationConfig &config)
    : config_(config) {}

// ---- Write Full --------------------------------------------------------------

void StateSerializer::write_full_snapshot(
    Bitstream &stream,
    uint32_t tick,
    const Entity *const *entities,
    uint16_t count
) {
    // Header: tick, entity count, base_tick=0 (full)
    stream.write_uint32(tick);
    stream.write_uint16(count);
    stream.write_uint16(0); // base_tick=0 signals "full snapshot"

    for (uint16_t i = 0; i < count; ++i) {
        if (entities[i] && entities[i]->is_alive()) {
            write_entity_full(stream, *entities[i]);
        }
    }
}

// ---- Write Delta -------------------------------------------------------------

void StateSerializer::write_delta_snapshot(
    Bitstream &stream,
    uint32_t tick,
    const Entity *const *entities,
    uint16_t count,
    const std::unordered_map<uint16_t, EntitySnapshot> &base
) {
    // Count changed entities first
    uint16_t changed_count = 0;
    for (uint16_t i = 0; i < count; ++i) {
        if (!entities[i] || !entities[i]->is_alive()) continue;
        uint16_t id = entities[i]->get_entity_id();
        auto it = base.find(id);
        if (it == base.end()) {
            ++changed_count; // new entity = full
        } else {
            EntitySnapshot::ChangeMask mask =
                entities[i]->compute_change_mask(it->second);
            if (mask != EntitySnapshot::CHANGED_NONE) {
                ++changed_count;
            }
        }
    }

    // Header: tick, count, delta_tick
    stream.write_uint32(tick);
    stream.write_uint16(changed_count);
    // We need the tick of the base snapshot — for now write 0 (delta from
    // client's last acked tick which the client tracks itself).
    stream.write_uint16(0);

    // Write delta entities
    for (uint16_t i = 0; i < count; ++i) {
        if (!entities[i] || !entities[i]->is_alive()) continue;
        uint16_t id = entities[i]->get_entity_id();
        auto it = base.find(id);
        if (it == base.end()) {
            // New entity: write full state with CHANGED_ALL
            stream.write_uint16(id);
            stream.write_uint32(EntitySnapshot::CHANGED_ALL);
            write_entity_full(stream, *entities[i]);
        } else {
            EntitySnapshot::ChangeMask mask =
                entities[i]->compute_change_mask(it->second);
            if (mask != EntitySnapshot::CHANGED_NONE) {
                stream.write_uint16(id);
                stream.write_uint32(static_cast<uint32_t>(mask));
                write_entity_delta(stream, *entities[i], it->second, mask);
            }
        }
    }
}

// ---- Read Full ---------------------------------------------------------------

std::vector<EntitySnapshot> StateSerializer::read_full_snapshot(
    Bitstream &stream,
    uint32_t *out_tick
) {
    std::vector<EntitySnapshot> result;

    *out_tick = stream.read_uint32();
    uint16_t count = stream.read_uint16();
    uint16_t base_tick = stream.read_uint16(); // should be 0 for full

    (void)base_tick; // unused in full read

    result.reserve(count);
    for (uint16_t i = 0; i < count; ++i) {
        uint16_t id = stream.read_uint16();
        EntitySnapshot::ChangeMask mask = read_change_mask(stream);
        EntitySnapshot snap = read_entity(stream, mask);
        snap.entity_id = id;
        result.push_back(snap);
    }

    return result;
}

// ---- Read Delta --------------------------------------------------------------

void StateSerializer::read_delta_snapshot(
    Bitstream &stream,
    std::unordered_map<uint16_t, EntitySnapshot> &state
) {
    uint32_t tick = stream.read_uint32();
    uint16_t count = stream.read_uint16();
    uint16_t base_tick = stream.read_uint16();

    (void)tick;
    (void)base_tick;

    for (uint16_t i = 0; i < count; ++i) {
        uint16_t id = stream.read_uint16();
        EntitySnapshot::ChangeMask mask = read_change_mask(stream);
        EntitySnapshot snap = read_entity(stream, mask);
        snap.entity_id = id;

        if (mask == EntitySnapshot::CHANGED_ALL) {
            state[id] = snap; // replace entirely
        } else {
            // Merge delta into existing state
            auto it = state.find(id);
            if (it != state.end()) {
                if (mask & EntitySnapshot::CHANGED_POSITION) it->second.position = snap.position;
                if (mask & EntitySnapshot::CHANGED_VELOCITY) it->second.velocity = snap.velocity;
                if (mask & EntitySnapshot::CHANGED_ROTATION) { it->second.yaw = snap.yaw; it->second.pitch = snap.pitch; }
                if (mask & EntitySnapshot::CHANGED_HEALTH)   it->second.health = snap.health;
                if (mask & EntitySnapshot::CHANGED_ARMOR)    it->second.armor = snap.armor;
                if (mask & EntitySnapshot::CHANGED_WEAPON)   it->second.weapon_id = snap.weapon_id;
                if (mask & EntitySnapshot::CHANGED_AMMO)     it->second.ammo = snap.ammo;
                if (mask & EntitySnapshot::CHANGED_FLAGS)    it->second.flags = snap.flags;
                if (mask & EntitySnapshot::CHANGED_INPUT)    it->second.last_processed_input = snap.last_processed_input;
            } else {
                state[id] = snap; // missing base, fall back to full replace
            }
        }
    }
}

// ---- Private: write ----------------------------------------------------------

void StateSerializer::write_entity_full(Bitstream &stream, const Entity &entity) {
    EntitySnapshot snap = entity.capture_snapshot();
    write_entity_delta(stream, entity, EntitySnapshot{}, EntitySnapshot::CHANGED_ALL);
}

void StateSerializer::write_entity_delta(
    Bitstream &stream,
    const Entity &entity,
    const EntitySnapshot &base,
    EntitySnapshot::ChangeMask mask
) {
    EntitySnapshot snap = entity.capture_snapshot();

    // Position
    if (mask & EntitySnapshot::CHANGED_POSITION) {
        stream.write_float_quantized(snap.position.x, config_.pos_min, config_.pos_max, config_.pos_bits);
        stream.write_float_quantized(snap.position.y, config_.pos_min, config_.pos_max, config_.pos_bits);
        stream.write_float_quantized(snap.position.z, config_.pos_min, config_.pos_max, config_.pos_bits);
    }

    // Velocity
    if (mask & EntitySnapshot::CHANGED_VELOCITY) {
        stream.write_float_quantized(snap.velocity.x, config_.vel_min, config_.vel_max, config_.vel_bits);
        stream.write_float_quantized(snap.velocity.y, config_.vel_min, config_.vel_max, config_.vel_bits);
        stream.write_float_quantized(snap.velocity.z, config_.vel_min, config_.vel_max, config_.vel_bits);
    }

    // Rotation
    if (mask & EntitySnapshot::CHANGED_ROTATION) {
        stream.write_float_quantized(snap.yaw, config_.yaw_min, config_.yaw_max, config_.yaw_bits);
        stream.write_float_quantized(snap.pitch, config_.pitch_min, config_.pitch_max, config_.pitch_bits);
    }

    // Health/Armor
    if (mask & EntitySnapshot::CHANGED_HEALTH) {
        stream.write_float_quantized(snap.health, config_.health_min, config_.health_max, config_.health_bits);
    }
    if (mask & EntitySnapshot::CHANGED_ARMOR) {
        stream.write_float_quantized(snap.armor, config_.health_min, config_.health_max, config_.health_bits);
    }

    // Weapon/Ammo
    if (mask & EntitySnapshot::CHANGED_WEAPON) {
        stream.write_uint8(snap.weapon_id);
    }
    if (mask & EntitySnapshot::CHANGED_AMMO) {
        stream.write_uint16(snap.ammo);
    }

    // Flags
    if (mask & EntitySnapshot::CHANGED_FLAGS) {
        stream.write_uint16(snap.flags);
    }

    // Input sequence
    if (mask & EntitySnapshot::CHANGED_INPUT) {
        stream.write_uint32(snap.last_processed_input);
    }
}

// ---- Private: read -----------------------------------------------------------

EntitySnapshot::ChangeMask StateSerializer::read_change_mask(Bitstream &stream) {
    return static_cast<EntitySnapshot::ChangeMask>(stream.read_uint32());
}

EntitySnapshot StateSerializer::read_entity(
    Bitstream &stream,
    EntitySnapshot::ChangeMask mask
) {
    EntitySnapshot snap;

    if (mask & EntitySnapshot::CHANGED_POSITION) {
        snap.position.x = stream.read_float_quantized(config_.pos_min, config_.pos_max, config_.pos_bits);
        snap.position.y = stream.read_float_quantized(config_.pos_min, config_.pos_max, config_.pos_bits);
        snap.position.z = stream.read_float_quantized(config_.pos_min, config_.pos_max, config_.pos_bits);
    }
    if (mask & EntitySnapshot::CHANGED_VELOCITY) {
        snap.velocity.x = stream.read_float_quantized(config_.vel_min, config_.vel_max, config_.vel_bits);
        snap.velocity.y = stream.read_float_quantized(config_.vel_min, config_.vel_max, config_.vel_bits);
        snap.velocity.z = stream.read_float_quantized(config_.vel_min, config_.vel_max, config_.vel_bits);
    }
    if (mask & EntitySnapshot::CHANGED_ROTATION) {
        snap.yaw   = stream.read_float_quantized(config_.yaw_min,   config_.yaw_max,   config_.yaw_bits);
        snap.pitch = stream.read_float_quantized(config_.pitch_min, config_.pitch_max, config_.pitch_bits);
    }
    if (mask & EntitySnapshot::CHANGED_HEALTH) {
        snap.health = stream.read_float_quantized(config_.health_min, config_.health_max, config_.health_bits);
    }
    if (mask & EntitySnapshot::CHANGED_ARMOR) {
        snap.armor = stream.read_float_quantized(config_.health_min, config_.health_max, config_.health_bits);
    }
    if (mask & EntitySnapshot::CHANGED_WEAPON) {
        snap.weapon_id = stream.read_uint8();
    }
    if (mask & EntitySnapshot::CHANGED_AMMO) {
        snap.ammo = stream.read_uint16();
    }
    if (mask & EntitySnapshot::CHANGED_FLAGS) {
        snap.flags = stream.read_uint16();
    }
    if (mask & EntitySnapshot::CHANGED_INPUT) {
        snap.last_processed_input = stream.read_uint32();
    }

    return snap;
}

} // namespace tactical_shooter