#ifndef TACTICAL_SHOOTER_BITSTREAM_H
#define TACTICAL_SHOOTER_BITSTREAM_H

#include <cstdint>
#include <cstring>
#include <cassert>
#include <climits>
#include <algorithm>
#include <vector>

namespace tactical_shooter {

/**
 * Bit-level read/write stream for compact network serialization.
 *
 * All multi-byte values are written in little-endian order regardless of
 * host endianness (network byte order). Booleans pack as single bits.
 * Floats can be quantized to arbitrary bit depths for bandwidth savings.
 *
 * Buffers are dynamically resized. Pre-allocate with reserve() to avoid
 * reallocation in hot paths.
 */
class Bitstream {
public:
    static constexpr size_t kMaxBufferSize = 1024 * 1024; // 1MB safety limit

    Bitstream() : buffer_(), bits_written_(0), bits_read_(0) {}

    explicit Bitstream(std::vector<uint8_t> data)
        : buffer_(std::move(data)), bits_written_(buffer_.size() * 8), bits_read_(0) {}

    // ---- Write -----------------------------------------------------------

    void write_bool(bool value) {
        write_bits(value ? 1 : 0, 1);
    }

    void write_uint8(uint8_t value) {
        write_bits(value, 8);
    }

    void write_uint16(uint16_t value) {
        write_bits(value, 16);
    }

    void write_uint32(uint32_t value) {
        write_bits(value, 32);
    }

    void write_int32(int32_t value) {
        // Zigzag encoding for efficient negative-number packing
        uint32_t zigzag = static_cast<uint32_t>((value << 1) ^ (value >> 31));
        write_bits(zigzag, 32);
    }

    /**
     * Write a float quantized to `nbits` within [min, max].
     * Storage: nbits bits. Resolution: (max-min) / (2^nbits - 1).
     * Pass nbits=32 for full-precision float (no quantization loss).
     */
    void write_float_quantized(float value, float min, float max, uint8_t nbits) {
        assert(nbits > 0 && nbits <= 32);
        if (nbits >= 32) {
            // Full precision: store as raw bits
            uint32_t raw;
            memcpy(&raw, &value, sizeof(raw));
            write_bits(raw, 32);
            return;
        }
        float clamped = std::clamp(value, min, max);
        float normalized = (clamped - min) / (max - min);
        uint32_t quantized = static_cast<uint32_t>(normalized * ((1u << nbits) - 1));
        write_bits(quantized, nbits);
    }

    /**
     * Write up to `nbits` bits of `value`. LSB first packing.
     */
    void write_bits(uint32_t value, uint8_t nbits) {
        assert(nbits > 0 && nbits <= 32);
        ensure_capacity(nbits);

        uint8_t *data = buffer_.data();
        size_t byte_pos = bits_written_ / 8;
        uint8_t bit_offset = bits_written_ % 8;

        for (uint8_t i = 0; i < nbits; ++i) {
            if (value & (1u << i)) {
                data[byte_pos] |= (1u << bit_offset);
            }
            ++bit_offset;
            if (bit_offset >= 8) {
                bit_offset = 0;
                ++byte_pos;
            }
        }
        bits_written_ += nbits;
    }

    // ---- Read ------------------------------------------------------------

    bool read_bool() {
        return read_bits(1) != 0;
    }

    uint8_t read_uint8() {
        return static_cast<uint8_t>(read_bits(8));
    }

    uint16_t read_uint16() {
        return static_cast<uint16_t>(read_bits(16));
    }

    uint32_t read_uint32() {
        return read_bits(32);
    }

    int32_t read_int32() {
        uint32_t zigzag = read_bits(32);
        return static_cast<int32_t>((zigzag >> 1) ^ -(static_cast<int32_t>(zigzag & 1)));
    }

    /**
     * Read a quantized float matching write_float_quantized().
     */
    float read_float_quantized(float min, float max, uint8_t nbits) {
        assert(nbits > 0 && nbits <= 32);
        if (nbits >= 32) {
            uint32_t raw = read_bits(32);
            float value;
            memcpy(&value, &raw, sizeof(value));
            return value;
        }
        uint32_t quantized = read_bits(nbits);
        float normalized = static_cast<float>(quantized) / static_cast<float>((1u << nbits) - 1);
        return min + normalized * (max - min);
    }

    /**
     * Read up to `nbits` bits, returned as LSB-packed uint32.
     */
    uint32_t read_bits(uint8_t nbits) {
        assert(nbits > 0 && nbits <= 32);
        assert((bits_read_ + nbits) <= bits_written_);

        const uint8_t *data = buffer_.data();
        size_t byte_pos = bits_read_ / 8;
        uint8_t bit_offset = bits_read_ % 8;
        uint32_t result = 0;

        for (uint8_t i = 0; i < nbits; ++i) {
            if (data[byte_pos] & (1u << bit_offset)) {
                result |= (1u << i);
            }
            ++bit_offset;
            if (bit_offset >= 8) {
                bit_offset = 0;
                ++byte_pos;
            }
        }
        bits_read_ += nbits;
        return result;
    }

    // ---- Array helpers ---------------------------------------------------

    /**
     * Write a dense array of booleans packed bit-by-bit.
     */
    void write_bool_array(const bool *values, size_t count) {
        for (size_t i = 0; i < count; ++i) {
            write_bool(values[i]);
        }
    }

    void read_bool_array(bool *values, size_t count) {
        for (size_t i = 0; i < count; ++i) {
            values[i] = read_bool();
        }
    }

    /**
     * Write a variable-length array of uint8 values with a uint16 count prefix.
     */
    void write_uint8_array(const uint8_t *values, uint16_t count) {
        write_uint16(count);
        for (uint16_t i = 0; i < count; ++i) {
            write_uint8(values[i]);
        }
    }

    std::vector<uint8_t> read_uint8_array() {
        uint16_t count = read_uint16();
        std::vector<uint8_t> result(count);
        for (uint16_t i = 0; i < count; ++i) {
            result[i] = read_uint8();
        }
        return result;
    }

    // ---- State -----------------------------------------------------------

    /// Total bytes consumed by written data
    size_t byte_size() const {
        return (bits_written_ + 7) / 8;
    }

    /// Number of bits written so far
    size_t bits_written() const { return bits_written_; }

    /// Number of bits read so far
    size_t bits_read() const { return bits_read_; }

    /// Remaining readable bits
    size_t bits_remaining() const {
        return bits_written_ - bits_read_;
    }

    /// Raw buffer (const access)
    const uint8_t *data() const { return buffer_.data(); }

    /// Clear everything, rewind
    void reset() {
        buffer_.clear();
        bits_written_ = 0;
        bits_read_ = 0;
    }

    /// Pre-allocate capacity in bytes
    void reserve(size_t bytes) {
        buffer_.reserve(bytes);
    }

    /// Steal the internal buffer
    std::vector<uint8_t> take_buffer() {
        std::vector<uint8_t> result = std::move(buffer_);
        reset();
        return result;
    }

private:
    void ensure_capacity(uint8_t extra_bits) {
        size_t needed_bytes = (bits_written_ + extra_bits + 7) / 8;
        if (needed_bytes > buffer_.size()) {
            if (needed_bytes > kMaxBufferSize) {
                // TODO: log error instead of assert in production
                assert(!"Bitstream overflow — reduce snapshot size or increase kMaxBufferSize");
            }
            buffer_.resize(std::max(buffer_.size() * 2, needed_bytes));
        }
    }

    std::vector<uint8_t> buffer_;
    size_t bits_written_ = 0;
    size_t bits_read_ = 0;
};

} // namespace tactical_shooter

#endif // TACTICAL_SHOOTER_BITSTREAM_H