#ifndef CAMELLYA_AST_H
#define CAMELLYA_AST_H

#include <memory>
#include <vector>
#include <string>
#include <variant>

namespace camellya {

// Forward declarations
struct Expr;
struct Stmt;

using ExprPtr = std::unique_ptr<Expr>;
using StmtPtr = std::unique_ptr<Stmt>;

// Value types
enum class ValueType {
    NUMBER, STRING, BOOL, LIST, MAP, CLASS, NIL
};

// Expressions
struct Expr {
    virtual ~Expr() = default;
};

struct BinaryExpr : public Expr {
    ExprPtr left;
    std::string op;
    ExprPtr right;
    
    BinaryExpr(ExprPtr left, std::string op, ExprPtr right)
        : left(std::move(left)), op(std::move(op)), right(std::move(right)) {}
};

struct UnaryExpr : public Expr {
    std::string op;
    ExprPtr operand;
    
    UnaryExpr(std::string op, ExprPtr operand)
        : op(std::move(op)), operand(std::move(operand)) {}
};

struct LiteralExpr : public Expr {
    std::variant<double, std::string, bool, std::monostate> value;
    
    explicit LiteralExpr(double value) : value(value) {}
    explicit LiteralExpr(std::string value) : value(std::move(value)) {}
    explicit LiteralExpr(bool value) : value(value) {}
    LiteralExpr() : value(std::monostate{}) {} // nil
};

struct VariableExpr : public Expr {
    std::string name;
    
    explicit VariableExpr(std::string name) : name(std::move(name)) {}
};

struct AssignExpr : public Expr {
    std::string name;
    ExprPtr value;
    
    AssignExpr(std::string name, ExprPtr value)
        : name(std::move(name)), value(std::move(value)) {}
};

struct CallExpr : public Expr {
    ExprPtr callee;
    std::vector<ExprPtr> arguments;
    
    CallExpr(ExprPtr callee, std::vector<ExprPtr> arguments)
        : callee(std::move(callee)), arguments(std::move(arguments)) {}
};

struct GetExpr : public Expr {
    ExprPtr object;
    std::string name;
    
    GetExpr(ExprPtr object, std::string name)
        : object(std::move(object)), name(std::move(name)) {}
};

struct SetExpr : public Expr {
    ExprPtr object;
    std::string name;
    ExprPtr value;
    
    SetExpr(ExprPtr object, std::string name, ExprPtr value)
        : object(std::move(object)), name(std::move(name)), value(std::move(value)) {}
};

struct IndexExpr : public Expr {
    ExprPtr object;
    ExprPtr index;
    
    IndexExpr(ExprPtr object, ExprPtr index)
        : object(std::move(object)), index(std::move(index)) {}
};

struct IndexSetExpr : public Expr {
    ExprPtr object;
    ExprPtr index;
    ExprPtr value;
    
    IndexSetExpr(ExprPtr object, ExprPtr index, ExprPtr value)
        : object(std::move(object)), index(std::move(index)), value(std::move(value)) {}
};

struct ListExpr : public Expr {
    std::vector<ExprPtr> elements;
    
    explicit ListExpr(std::vector<ExprPtr> elements)
        : elements(std::move(elements)) {}
};

struct MapExpr : public Expr {
    std::vector<std::pair<ExprPtr, ExprPtr>> pairs;
    
    explicit MapExpr(std::vector<std::pair<ExprPtr, ExprPtr>> pairs)
        : pairs(std::move(pairs)) {}
};

// Statements
struct Stmt {
    virtual ~Stmt() = default;
};

struct ExprStmt : public Stmt {
    ExprPtr expression;
    
    explicit ExprStmt(ExprPtr expression) : expression(std::move(expression)) {}
};

struct VarDecl : public Stmt {
    std::string type_name;
    std::string name;
    ExprPtr initializer;
    
    VarDecl(std::string type_name, std::string name, ExprPtr initializer = nullptr)
        : type_name(std::move(type_name)), name(std::move(name)), initializer(std::move(initializer)) {}
};

struct BlockStmt : public Stmt {
    std::vector<StmtPtr> statements;
    
    explicit BlockStmt(std::vector<StmtPtr> statements)
        : statements(std::move(statements)) {}
};

struct IfStmt : public Stmt {
    ExprPtr condition;
    StmtPtr then_branch;
    StmtPtr else_branch;
    
    IfStmt(ExprPtr condition, StmtPtr then_branch, StmtPtr else_branch = nullptr)
        : condition(std::move(condition)), then_branch(std::move(then_branch)), 
          else_branch(std::move(else_branch)) {}
};

struct WhileStmt : public Stmt {
    ExprPtr condition;
    StmtPtr body;
    
    WhileStmt(ExprPtr condition, StmtPtr body)
        : condition(std::move(condition)), body(std::move(body)) {}
};

struct ForStmt : public Stmt {
    StmtPtr initializer;
    ExprPtr condition;
    ExprPtr increment;
    StmtPtr body;
    
    ForStmt(StmtPtr initializer, ExprPtr condition, ExprPtr increment, StmtPtr body)
        : initializer(std::move(initializer)), condition(std::move(condition)),
          increment(std::move(increment)), body(std::move(body)) {}
};

struct ReturnStmt : public Stmt {
    ExprPtr value;
    
    explicit ReturnStmt(ExprPtr value = nullptr) : value(std::move(value)) {}
};

struct FunctionDecl : public Stmt {
    std::string name;
    std::vector<std::pair<std::string, std::string>> parameters; // (type, name)
    std::string return_type;
    std::shared_ptr<BlockStmt> body;
    
    FunctionDecl(std::string name, 
                 std::vector<std::pair<std::string, std::string>> parameters,
                 std::string return_type,
                 std::shared_ptr<BlockStmt> body)
        : name(std::move(name)), parameters(std::move(parameters)),
          return_type(std::move(return_type)), body(std::move(body)) {}
};

struct ClassDecl : public Stmt {
    std::string name;
    std::vector<StmtPtr> members; // VarDecl and FunctionDecl
    
    ClassDecl(std::string name, std::vector<StmtPtr> members)
        : name(std::move(name)), members(std::move(members)) {}
};

struct Program {
    std::vector<StmtPtr> statements;
    
    explicit Program(std::vector<StmtPtr> statements)
        : statements(std::move(statements)) {}
};

} // namespace camellya

#endif // CAMELLYA_AST_H