#include "ir_gen.h"
#include <stack>
#include <unordered_map>
#include <algorithm>

namespace nix_irc {

struct NameResolver::Impl {
    std::vector<std::unordered_map<std::string, uint32_t>> scopes;
    std::vector<std::vector<std::string>> scope_names;

    Impl() {
        scopes.push_back({});
        scope_names.push_back({});
    }
};

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

void NameResolver::enter_scope() {
    pImpl->scopes.push_back({});
    pImpl->scope_names.push_back({});
}

void NameResolver::exit_scope() {
    if (!pImpl->scopes.empty()) {
        pImpl->scopes.pop_back();
        pImpl->scope_names.pop_back();
    }
}

void NameResolver::bind(const std::string& name) {
    if (pImpl->scopes.empty()) return;
    uint32_t idx = pImpl->scope_names.back().size();
    pImpl->scopes.back()[name] = idx;
    pImpl->scope_names.back().push_back(name);
}

uint32_t NameResolver::resolve(const std::string& name) {
    for (int i = (int)pImpl->scopes.size() - 1; i >= 0; --i) {
        auto it = pImpl->scopes[i].find(name);
        if (it != pImpl->scopes[i].end()) {
            uint32_t depth = pImpl->scopes.size() - 1 - i;
            uint32_t offset = it->second;
            return depth << 16 | offset;
        }
    }
    return 0xFFFFFFFF;
}

bool NameResolver::is_bound(const std::string& name) const {
    for (auto it = pImpl->scopes.rbegin(); it != pImpl->scopes.rend(); ++it) {
        if (it->count(name)) return true;
    }
    return false;
}

struct IRGenerator::Impl {
    std::unordered_map<std::string, uint32_t> string_table;
    uint32_t next_string_id = 0;
    NameResolver name_resolver;

    Impl() {}

    uint32_t add_string(const std::string& str) {
        auto it = string_table.find(str);
        if (it != string_table.end()) {
            return it->second;
        }
        uint32_t id = next_string_id++;
        string_table[str] = id;
        return id;
    }

    std::shared_ptr<Node> convert(const std::shared_ptr<Node>& node_ptr) {
        if (!node_ptr) return std::make_shared<Node>(ConstNullNode{});

        const Node& node = *node_ptr;

        if (auto* n = node.get_if<ConstIntNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstStringNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstPathNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstBoolNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<ConstNullNode>()) {
            return std::make_shared<Node>(*n);
        }
        if (auto* n = node.get_if<VarNode>()) {
            uint32_t idx = name_resolver.resolve(n->name.value_or(""));
            VarNode converted(idx);
            converted.name = n->name;
            converted.line = n->line;
            return std::make_shared<Node>(converted);
        }
        if (auto* n = node.get_if<LambdaNode>()) {
            name_resolver.enter_scope();
            if (n->param_name) {
                name_resolver.bind(*n->param_name);
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LambdaNode lambda(n->arity, body, n->line);
            lambda.param_name = n->param_name;
            return std::make_shared<Node>(lambda);
        }
        if (auto* n = node.get_if<AppNode>()) {
            auto func = convert(n->func);
            auto arg = convert(n->arg);
            return std::make_shared<Node>(AppNode(func, arg, n->line));
        }
        if (auto* n = node.get_if<AttrsetNode>()) {
            AttrsetNode attrs(n->recursive, n->line);
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->attrs) {
                name_resolver.bind(key);
            }
            for (const auto& [key, val] : n->attrs) {
                attrs.attrs.push_back({key, convert(val)});
            }
            name_resolver.exit_scope();
            return std::make_shared<Node>(attrs);
        }
        if (auto* n = node.get_if<SelectNode>()) {
            auto expr = convert(n->expr);
            auto attr = convert(n->attr);
            SelectNode select(expr, attr, n->line);
            if (n->default_expr) {
                select.default_expr = convert(*n->default_expr);
            }
            return std::make_shared<Node>(select);
        }
        if (auto* n = node.get_if<HasAttrNode>()) {
            auto expr = convert(n->expr);
            auto attr = convert(n->attr);
            return std::make_shared<Node>(HasAttrNode(expr, attr, n->line));
        }
        if (auto* n = node.get_if<WithNode>()) {
            auto attrs = convert(n->attrs);
            auto body = convert(n->body);
            return std::make_shared<Node>(WithNode(attrs, body, n->line));
        }
        if (auto* n = node.get_if<IfNode>()) {
            auto cond = convert(n->cond);
            auto then_b = convert(n->then_branch);
            auto else_b = convert(n->else_branch);
            return std::make_shared<Node>(IfNode(cond, then_b, else_b, n->line));
        }
        if (auto* n = node.get_if<LetNode>()) {
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->bindings) {
                name_resolver.bind(key);
            }
            std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
            for (const auto& [key, val] : n->bindings) {
                new_bindings.push_back({key, convert(val)});
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LetNode let(body, n->line);
            let.bindings = std::move(new_bindings);
            return std::make_shared<Node>(let);
        }
        if (auto* n = node.get_if<LetRecNode>()) {
            name_resolver.enter_scope();
            for (const auto& [key, val] : n->bindings) {
                name_resolver.bind(key);
            }
            std::vector<std::pair<std::string, std::shared_ptr<Node>>> new_bindings;
            for (const auto& [key, val] : n->bindings) {
                new_bindings.push_back({key, convert(val)});
            }
            auto body = convert(n->body);
            name_resolver.exit_scope();
            LetRecNode letrec(body, n->line);
            letrec.bindings = std::move(new_bindings);
            return std::make_shared<Node>(letrec);
        }
        if (auto* n = node.get_if<AssertNode>()) {
            auto cond = convert(n->cond);
            auto body = convert(n->body);
            return std::make_shared<Node>(AssertNode(cond, body, n->line));
        }
        if (auto* n = node.get_if<BinaryOpNode>()) {
            auto left = convert(n->left);
            auto right = convert(n->right);
            return std::make_shared<Node>(BinaryOpNode(n->op, left, right, n->line));
        }
        if (auto* n = node.get_if<UnaryOpNode>()) {
            auto operand = convert(n->operand);
            return std::make_shared<Node>(UnaryOpNode(n->op, operand, n->line));
        }
        return std::make_shared<Node>(ConstNullNode{});
    }
};

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

void IRGenerator::set_string_table(const std::unordered_map<std::string, uint32_t>& table) {
    pImpl->string_table = table;
}

uint32_t IRGenerator::add_string(const std::string& str) {
    return pImpl->add_string(str);
}

std::shared_ptr<Node> IRGenerator::generate(const std::shared_ptr<Node>& ast) {
    return pImpl->convert(ast);
}

}