#ifndef NIX_IRC_TYPES_H
#define NIX_IRC_TYPES_H

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

namespace nix_irc {

constexpr uint32_t IR_MAGIC = 0x4E495258;
constexpr uint32_t IR_VERSION = 3;

enum class NodeType : uint8_t {
  CONST_INT = 0x01,
  CONST_FLOAT = 0x06,
  CONST_STRING = 0x02,
  CONST_PATH = 0x03,
  CONST_BOOL = 0x04,
  CONST_NULL = 0x05,
  CONST_URI = 0x07,
  CONST_LOOKUP_PATH = 0x08,
  VAR = 0x10,
  LAMBDA = 0x20,
  APP = 0x21,
  BINARY_OP = 0x22,
  UNARY_OP = 0x23,
  IMPORT = 0x24,
  ATTRSET = 0x30,
  SELECT = 0x31,
  HAS_ATTR = 0x34,
  WITH = 0x32,
  LIST = 0x33,
  IF = 0x40,
  LET = 0x50,
  LETREC = 0x51,
  ASSERT = 0x52,
  THUNK = 0x60,
  FORCE = 0x61,
  LAMBDA_PATTERN = 0x70,
  INHERIT = 0x71,
  INHERIT_FROM = 0x72,
  STRING_INTERPOLATION = 0x73,
  BUILTIN_CALL = 0x74,
  ERROR = 0xFF
};

enum class BinaryOp : uint8_t {
  ADD,
  SUB,
  MUL,
  DIV,
  CONCAT,
  EQ,
  NE,
  LT,
  GT,
  LE,
  GE,
  AND,
  OR,
  IMPL,
  MERGE
};

enum class UnaryOp : uint8_t { NEG, NOT };

// Forward declare Node for use in shared_ptr
class Node;

struct ConstIntNode {
  int64_t value;
  uint32_t line = 0;
  ConstIntNode(int64_t v = 0, uint32_t l = 0) : value(v), line(l) {}
};

struct ConstStringNode {
  std::string value;
  uint32_t line = 0;
  ConstStringNode(std::string v = "", uint32_t l = 0) : value(std::move(v)), line(l) {}
};

struct ConstPathNode {
  std::string value;
  uint32_t line = 0;
  ConstPathNode(std::string v = "", uint32_t l = 0) : value(std::move(v)), line(l) {}
};

struct ConstBoolNode {
  bool value;
  uint32_t line = 0;
  ConstBoolNode(bool v = false, uint32_t l = 0) : value(v), line(l) {}
};

struct ConstNullNode {
  uint32_t line = 0;
  ConstNullNode(uint32_t l = 0) : line(l) {}
};

struct ConstFloatNode {
  double value;
  uint32_t line = 0;
  ConstFloatNode(double v = 0.0, uint32_t l = 0) : value(v), line(l) {}
};

struct ConstURINode {
  std::string value;
  uint32_t line = 0;
  ConstURINode(std::string v = "", uint32_t l = 0) : value(std::move(v)), line(l) {}
};

struct ConstLookupPathNode {
  std::string value; // e.g., "nixpkgs" or "nixpkgs/lib"
  uint32_t line = 0;
  ConstLookupPathNode(std::string v = "", uint32_t l = 0) : value(std::move(v)), line(l) {}
};

struct VarNode {
  uint32_t index = 0;
  std::optional<std::string> name;
  uint32_t line = 0;
  VarNode(uint32_t idx = 0, std::string n = "", uint32_t l = 0)
      : index(idx), name(n.empty() ? std::nullopt : std::optional<std::string>(n)), line(l) {}
};

struct LambdaNode {
  uint32_t arity = 1;
  std::shared_ptr<Node> body;
  std::optional<std::string> param_name;
  bool strict_pattern = true;
  uint32_t line = 0;
  LambdaNode(uint32_t a, std::shared_ptr<Node> b, uint32_t l = 0);
};

struct PatternField {
  std::string name;
  std::optional<std::shared_ptr<Node>> default_value;

  PatternField(std::string n, std::optional<std::shared_ptr<Node>> def = std::nullopt)
      : name(std::move(n)), default_value(std::move(def)) {}
};

struct LambdaPatternNode {
  std::vector<PatternField> required_fields;
  std::vector<PatternField> optional_fields;
  std::optional<std::string> at_binding;
  bool allow_extra;
  std::shared_ptr<Node> body;
  uint32_t line = 0;

  LambdaPatternNode(std::shared_ptr<Node> b, uint32_t l = 0);
};

struct AppNode {
  std::shared_ptr<Node> func;
  std::shared_ptr<Node> arg;
  uint32_t line = 0;
  AppNode(std::shared_ptr<Node> f, std::shared_ptr<Node> a, uint32_t l = 0);
};

struct BinaryOpNode {
  BinaryOp op;
  std::shared_ptr<Node> left;
  std::shared_ptr<Node> right;
  uint32_t line = 0;
  BinaryOpNode(BinaryOp o, std::shared_ptr<Node> l, std::shared_ptr<Node> r, uint32_t ln = 0);
};

struct UnaryOpNode {
  UnaryOp op;
  std::shared_ptr<Node> operand;
  uint32_t line = 0;
  UnaryOpNode(UnaryOp o, std::shared_ptr<Node> operand, uint32_t l = 0);
};

struct AttrBinding {
  std::optional<std::string> static_name; // Static key like "foo"
  std::shared_ptr<Node> dynamic_name;     // Dynamic key like ${expr}
  std::shared_ptr<Node> value;

  // Static attribute
  AttrBinding(std::string name, std::shared_ptr<Node> val)
      : static_name(std::move(name)), value(std::move(val)) {}

  // Dynamic attribute
  AttrBinding(std::shared_ptr<Node> name_expr, std::shared_ptr<Node> val)
      : dynamic_name(std::move(name_expr)), value(std::move(val)) {}

  bool is_dynamic() const { return !static_name.has_value(); }
};

struct AttrsetNode {
  std::vector<AttrBinding> attrs;
  bool recursive = false;
  uint32_t line = 0;
  AttrsetNode(bool rec = false, uint32_t l = 0) : recursive(rec), line(l) {}
};

struct SelectNode {
  std::shared_ptr<Node> expr;
  std::shared_ptr<Node> attr;
  std::optional<std::shared_ptr<Node>> default_expr;
  uint32_t line = 0;
  SelectNode(std::shared_ptr<Node> e, std::shared_ptr<Node> a, uint32_t l = 0);
};

struct HasAttrNode {
  std::shared_ptr<Node> expr;
  std::shared_ptr<Node> attr;
  uint32_t line = 0;
  HasAttrNode(std::shared_ptr<Node> e, std::shared_ptr<Node> a, uint32_t l = 0);
};

struct WithNode {
  std::shared_ptr<Node> attrs;
  std::shared_ptr<Node> body;
  uint32_t line = 0;
  WithNode(std::shared_ptr<Node> a, std::shared_ptr<Node> b, uint32_t l = 0);
};

struct IfNode {
  std::shared_ptr<Node> cond;
  std::shared_ptr<Node> then_branch;
  std::shared_ptr<Node> else_branch;
  uint32_t line = 0;
  IfNode(std::shared_ptr<Node> c, std::shared_ptr<Node> t, std::shared_ptr<Node> e, uint32_t l = 0);
};

struct LetNode {
  std::vector<std::pair<std::string, std::shared_ptr<Node>>> bindings;
  std::shared_ptr<Node> body;
  uint32_t line = 0;
  LetNode(std::shared_ptr<Node> b, uint32_t l = 0);
};

struct LetRecNode {
  std::vector<std::pair<std::string, std::shared_ptr<Node>>> bindings;
  std::shared_ptr<Node> body;
  uint32_t line = 0;
  LetRecNode(std::shared_ptr<Node> b, uint32_t l = 0);
};

struct AssertNode {
  std::shared_ptr<Node> cond;
  std::shared_ptr<Node> body;
  uint32_t line = 0;
  AssertNode(std::shared_ptr<Node> c, std::shared_ptr<Node> b, uint32_t l = 0);
};

struct ImportNode {
  std::shared_ptr<Node> path; // Path expression to import
  uint32_t line = 0;
  ImportNode(std::shared_ptr<Node> p, uint32_t l = 0);
};

struct ThunkNode {
  std::shared_ptr<Node> expr;
  uint32_t line = 0;
  ThunkNode(std::shared_ptr<Node> e, uint32_t l = 0);
};

struct ForceNode {
  std::shared_ptr<Node> expr;
  uint32_t line = 0;
  ForceNode(std::shared_ptr<Node> e, uint32_t l = 0);
};

struct ListNode {
  std::vector<std::shared_ptr<Node>> elements;
  uint32_t line = 0;
  ListNode(std::vector<std::shared_ptr<Node>> elems = {}, uint32_t l = 0)
      : elements(std::move(elems)), line(l) {}
};

struct InheritNode {
  std::vector<std::string> names;
  uint32_t line = 0;

  InheritNode(std::vector<std::string> n = {}, uint32_t l = 0) : names(std::move(n)), line(l) {}
};

struct InheritFromNode {
  std::shared_ptr<Node> source;
  std::vector<std::string> names;
  uint32_t line = 0;

  InheritFromNode(std::shared_ptr<Node> src, std::vector<std::string> n, uint32_t l = 0)
      : source(std::move(src)), names(std::move(n)), line(l) {}
};

struct StringPart {
  enum class Type { LITERAL, EXPR };
  Type type;
  std::string literal;
  std::shared_ptr<Node> expr;

  static StringPart make_literal(std::string lit) {
    StringPart part;
    part.type = Type::LITERAL;
    part.literal = std::move(lit);
    return part;
  }

  static StringPart make_expr(std::shared_ptr<Node> e) {
    StringPart part;
    part.type = Type::EXPR;
    part.expr = std::move(e);
    return part;
  }
};

struct StringInterpolationNode {
  std::vector<StringPart> parts;
  uint32_t line = 0;

  StringInterpolationNode(std::vector<StringPart> p = {}, uint32_t l = 0)
      : parts(std::move(p)), line(l) {}
};

struct BuiltinCallNode {
  std::string builtin_name;
  std::vector<std::shared_ptr<Node>> args;
  uint32_t line = 0;

  BuiltinCallNode(std::string name, std::vector<std::shared_ptr<Node>> a = {}, uint32_t l = 0)
      : builtin_name(std::move(name)), args(std::move(a)), line(l) {}
};

// Node wraps a variant for type-safe AST
class Node {
public:
  using Variant =
      std::variant<ConstIntNode, ConstFloatNode, ConstStringNode, ConstPathNode, ConstBoolNode,
                   ConstNullNode, ConstURINode, ConstLookupPathNode, VarNode, LambdaNode, AppNode,
                   BinaryOpNode, UnaryOpNode, ImportNode, AttrsetNode, SelectNode, HasAttrNode,
                   WithNode, IfNode, LetNode, LetRecNode, AssertNode, ThunkNode, ForceNode,
                   ListNode, LambdaPatternNode, InheritNode, InheritFromNode,
                   StringInterpolationNode, BuiltinCallNode>;

  Variant data;

  template <typename T> Node(T&& value) : data(std::forward<T>(value)) {}

  template <typename T> T* get_if() { return std::get_if<T>(&data); }

  template <typename T> const T* get_if() const { return std::get_if<T>(&data); }

  template <typename T> bool holds() const { return std::holds_alternative<T>(data); }
};



struct SourceFile {
  std::string path;
  std::string content;
  std::shared_ptr<Node> ast;
};

struct IRModule {
  uint32_t version = IR_VERSION;
  std::vector<SourceFile> sources;
  std::vector<std::pair<std::string, std::string>> imports;
  std::shared_ptr<Node> entry;
  std::unordered_map<std::string, uint32_t> string_table;
};

} // namespace nix_irc
#endif