#include "serializer.h"
#include <cstring>
#include <sstream>
#include <iostream>

namespace nix_irc {

struct Serializer::Impl {
    std::vector<uint8_t> buffer;

    void write_u32(uint32_t val) {
        buffer.push_back((val >> 0) & 0xFF);
        buffer.push_back((val >> 8) & 0xFF);
        buffer.push_back((val >> 16) & 0xFF);
        buffer.push_back((val >> 24) & 0xFF);
    }

    void write_u64(uint64_t val) {
        for (int i = 0; i < 8; i++) {
            buffer.push_back((val >> (i * 8)) & 0xFF);
        }
    }

    void write_u8(uint8_t val) {
        buffer.push_back(val);
    }

    void write_string(const std::string& str) {
        write_u32(str.size());
        buffer.insert(buffer.end(), str.begin(), str.end());
    }

    NodeType get_node_type(const Node& node) {
        if (node.holds<ConstIntNode>()) return NodeType::CONST_INT;
        if (node.holds<ConstStringNode>()) return NodeType::CONST_STRING;
        if (node.holds<ConstPathNode>()) return NodeType::CONST_PATH;
        if (node.holds<ConstBoolNode>()) return NodeType::CONST_BOOL;
        if (node.holds<ConstNullNode>()) return NodeType::CONST_NULL;
        if (node.holds<VarNode>()) return NodeType::VAR;
        if (node.holds<LambdaNode>()) return NodeType::LAMBDA;
        if (node.holds<AppNode>()) return NodeType::APP;
        if (node.holds<BinaryOpNode>()) return NodeType::BINARY_OP;
        if (node.holds<UnaryOpNode>()) return NodeType::UNARY_OP;
        if (node.holds<AttrsetNode>()) return NodeType::ATTRSET;
        if (node.holds<SelectNode>()) return NodeType::SELECT;
        if (node.holds<WithNode>()) return NodeType::WITH;
        if (node.holds<IfNode>()) return NodeType::IF;
        if (node.holds<LetNode>()) return NodeType::LET;
        if (node.holds<LetRecNode>()) return NodeType::LETREC;
        if (node.holds<AssertNode>()) return NodeType::ASSERT;
        return NodeType::ERROR;
    }

    uint32_t get_node_line(const Node& node) {
        return std::visit([](const auto& n) { return n.line; }, node.data);
    }

    void write_node(const Node& node) {
        write_u8(static_cast<uint8_t>(get_node_type(node)));
        write_u32(get_node_line(node));

        if (auto* n = node.get_if<ConstIntNode>()) {
            write_u64(static_cast<uint64_t>(n->value));
        } else if (auto* n = node.get_if<ConstStringNode>()) {
            write_string(n->value);
        } else if (auto* n = node.get_if<ConstPathNode>()) {
            write_string(n->value);
        } else if (auto* n = node.get_if<ConstBoolNode>()) {
            write_u8(n->value ? 1 : 0);
        } else if (auto* n = node.get_if<ConstNullNode>()) {
            // No data for null
        } else if (auto* n = node.get_if<VarNode>()) {
            write_u32(n->index);
        } else if (auto* n = node.get_if<LambdaNode>()) {
            write_u32(n->arity);
            if (n->body) write_node(*n->body);
        } else if (auto* n = node.get_if<AppNode>()) {
            if (n->func) write_node(*n->func);
            if (n->arg) write_node(*n->arg);
        } else if (auto* n = node.get_if<BinaryOpNode>()) {
            write_u8(static_cast<uint8_t>(n->op));
            if (n->left) write_node(*n->left);
            if (n->right) write_node(*n->right);
        } else if (auto* n = node.get_if<UnaryOpNode>()) {
            write_u8(static_cast<uint8_t>(n->op));
            if (n->operand) write_node(*n->operand);
        } else if (auto* n = node.get_if<AttrsetNode>()) {
            write_u8(n->recursive ? 1 : 0);
            write_u32(n->attrs.size());
            for (const auto& [key, val] : n->attrs) {
                write_string(key);
                if (val) write_node(*val);
            }
        } else if (auto* n = node.get_if<SelectNode>()) {
            if (n->expr) write_node(*n->expr);
            if (n->attr) write_node(*n->attr);
            if (n->default_expr && *n->default_expr) {
                write_u8(1);
                write_node(**n->default_expr);
            } else {
                write_u8(0);
            }
        } else if (auto* n = node.get_if<WithNode>()) {
            if (n->attrs) write_node(*n->attrs);
            if (n->body) write_node(*n->body);
        } else if (auto* n = node.get_if<IfNode>()) {
            if (n->cond) write_node(*n->cond);
            if (n->then_branch) write_node(*n->then_branch);
            if (n->else_branch) write_node(*n->else_branch);
        } else if (auto* n = node.get_if<LetNode>()) {
            write_u32(n->bindings.size());
            for (const auto& [key, val] : n->bindings) {
                write_string(key);
                if (val) write_node(*val);
            }
            if (n->body) write_node(*n->body);
        } else if (auto* n = node.get_if<LetRecNode>()) {
            write_u32(n->bindings.size());
            for (const auto& [key, val] : n->bindings) {
                write_string(key);
                if (val) write_node(*val);
            }
            if (n->body) write_node(*n->body);
        } else if (auto* n = node.get_if<AssertNode>()) {
            if (n->cond) write_node(*n->cond);
            if (n->body) write_node(*n->body);
        }
    }
};

Serializer::Serializer() : pImpl(std::make_unique<Impl>()) {}
Serializer::~Serializer() = default;

void Serializer::serialize(const IRModule& module, const std::string& path) {
    auto bytes = serialize_to_bytes(module);
    std::ofstream out(path, std::ios::binary);
    out.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
}

std::vector<uint8_t> Serializer::serialize_to_bytes(const IRModule& module) {
    pImpl->buffer.clear();

    pImpl->write_u32(IR_MAGIC);
    pImpl->write_u32(IR_VERSION);

    pImpl->write_u32(module.sources.size());
    for (const auto& src : module.sources) {
        pImpl->write_string(src.path);
        pImpl->write_string(src.content);
    }

    pImpl->write_u32(module.imports.size());
    for (const auto& [from, to] : module.imports) {
        pImpl->write_string(from);
        pImpl->write_string(to);
    }

    pImpl->write_u32(module.string_table.size());
    for (const auto& [str, id] : module.string_table) {
        pImpl->write_string(str);
        pImpl->write_u32(id);
    }

    if (module.entry && module.entry != nullptr) {
        pImpl->write_u8(1);
        pImpl->write_node(*module.entry);
    } else {
        pImpl->write_u8(0);
    }

    return pImpl->buffer;
}

struct Deserializer::Impl {
    std::vector<uint8_t> buffer;
    size_t pos = 0;

    uint32_t read_u32() {
        uint32_t val = 0;
        val |= buffer[pos + 0];
        val |= (uint32_t)buffer[pos + 1] << 8;
        val |= (uint32_t)buffer[pos + 2] << 16;
        val |= (uint32_t)buffer[pos + 3] << 24;
        pos += 4;
        return val;
    }

    uint64_t read_u64() {
        uint64_t val = 0;
        for (int i = 0; i < 8; i++) {
            val |= (uint64_t)buffer[pos + i] << (i * 8);
        }
        pos += 8;
        return val;
    }

    uint8_t read_u8() {
        return buffer[pos++];
    }

    std::string read_string() {
        uint32_t len = read_u32();
        std::string str(reinterpret_cast<const char*>(&buffer[pos]), len);
        pos += len;
        return str;
    }
};

Deserializer::Deserializer() : pImpl(std::make_unique<Impl>()) {}
Deserializer::~Deserializer() = default;

IRModule Deserializer::deserialize(const std::string& path) {
    std::ifstream in(path, std::ios::binary | std::ios::ate);
    size_t size = in.tellg();
    in.seekg(0);
    pImpl->buffer.resize(size);
    in.read(reinterpret_cast<char*>(pImpl->buffer.data()), size);
    pImpl->pos = 0;
    return deserialize(pImpl->buffer);
}

IRModule Deserializer::deserialize(const std::vector<uint8_t>& data) {
    pImpl->buffer = data;
    pImpl->pos = 0;

    IRModule module;

    uint32_t magic = pImpl->read_u32();
    if (magic != IR_MAGIC) {
        throw std::runtime_error("Invalid IR file");
    }

    uint32_t version = pImpl->read_u32();
    if (version != IR_VERSION) {
        throw std::runtime_error("Unsupported IR version");
    }

    uint32_t num_sources = pImpl->read_u32();
    for (uint32_t i = 0; i < num_sources; i++) {
        SourceFile src;
        src.path = pImpl->read_string();
        src.content = pImpl->read_string();
        module.sources.push_back(src);
    }

    uint32_t num_imports = pImpl->read_u32();
    for (uint32_t i = 0; i < num_imports; i++) {
        module.imports.push_back({pImpl->read_string(), pImpl->read_string()});
    }

    uint32_t num_strings = pImpl->read_u32();
    for (uint32_t i = 0; i < num_strings; i++) {
        std::string str = pImpl->read_string();
        uint32_t id = pImpl->read_u32();
        module.string_table[str] = id;
    }

    if (pImpl->read_u8()) {
        // TODO: deserialize AST
    }

    return module;
}

}