#ifndef NIX_IRC_TYPES_H
#define NIX_IRC_TYPES_H

#include <cstdint>
#include <string>
#include <vector>
#include <unordered_map>
#include <optional>
#include <memory>
#include <variant>
#include <fstream>
#include <sstream>

namespace nix_irc {

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

enum class NodeType : uint8_t {
    CONST_INT = 0x01,
    CONST_STRING = 0x02,
    CONST_PATH = 0x03,
    CONST_BOOL = 0x04,
    CONST_NULL = 0x05,
    VAR = 0x10,
    LAMBDA = 0x20,
    APP = 0x21,
    BINARY_OP = 0x22,
    UNARY_OP = 0x23,
    ATTRSET = 0x30,
    SELECT = 0x31,
    HAS_ATTR = 0x34,
    WITH = 0x32,
    IF = 0x40,
    LET = 0x50,
    LETREC = 0x51,
    ASSERT = 0x52,
    THUNK = 0x60,
    FORCE = 0x61,
    ERROR = 0xFF
};

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

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 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 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 AttrsetNode {
    std::vector<std::pair<std::string, std::shared_ptr<Node>>> 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 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);
};

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

    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); }
};

// Constructor implementations
inline LambdaNode::LambdaNode(uint32_t a, std::shared_ptr<Node> b, uint32_t l)
    : arity(a), body(b), line(l) {}

inline AppNode::AppNode(std::shared_ptr<Node> f, std::shared_ptr<Node> a, uint32_t l)
    : func(f), arg(a), line(l) {}

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

inline UnaryOpNode::UnaryOpNode(UnaryOp o, std::shared_ptr<Node> operand, uint32_t l)
    : op(o), operand(operand), line(l) {}

inline SelectNode::SelectNode(std::shared_ptr<Node> e, std::shared_ptr<Node> a, uint32_t l)
    : expr(e), attr(a), line(l) {}

inline HasAttrNode::HasAttrNode(std::shared_ptr<Node> e, std::shared_ptr<Node> a, uint32_t l)
    : expr(e), attr(a), line(l) {}

inline WithNode::WithNode(std::shared_ptr<Node> a, std::shared_ptr<Node> b, uint32_t l)
    : attrs(a), body(b), line(l) {}

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

inline LetNode::LetNode(std::shared_ptr<Node> b, uint32_t l)
    : body(b), line(l) {}

inline LetRecNode::LetRecNode(std::shared_ptr<Node> b, uint32_t l)
    : body(b), line(l) {}

inline AssertNode::AssertNode(std::shared_ptr<Node> c, std::shared_ptr<Node> b, uint32_t l)
    : cond(c), body(b), line(l) {}

inline ThunkNode::ThunkNode(std::shared_ptr<Node> e, uint32_t l)
    : expr(e), line(l) {}

inline ForceNode::ForceNode(std::shared_ptr<Node> e, uint32_t l)
    : expr(e), line(l) {}

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;
};

}
#endif