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

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;
      new_bindings.reserve(n->bindings.size());
      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;
      new_bindings.reserve(n->bindings.size());
      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);
}

} // namespace nix_irc