#ifndef AST_H
#define AST_H

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

#include "codegen.h"
#include "types.h"
#include "tokens.h"
#include "typechecker.h"

namespace AST {
    class Node {
    public:
        token::Metadata m_meta;

        Node(token::Metadata meta) : m_meta{ meta } {}
        virtual std::string formatted() = 0;
        virtual ~Node() = default;
    };

    class Expression : public Node {
    public:
        Expression(token::Metadata meta) : Node{ meta } {}
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) = 0;
    };

    class Statement : public Node {
    public:
        Statement(token::Metadata meta) : Node{ meta } {}
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) = 0;
    };

    class IntLiteralExpression : public Expression {
    private:
        int m_value;
    public:
        IntLiteralExpression(token::Metadata meta, int value) : Expression{ meta }, m_value{ value } {}
        virtual ~IntLiteralExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    class StringLiteralExpression : public Expression {
    private:
        std::string m_value;
    public:
        StringLiteralExpression(token::Metadata meta, std::string value) : Expression{ meta }, m_value{ value } {}
        virtual ~StringLiteralExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    class ValueReferenceExpression : public Expression {
    private:
        std::string m_name;
    public:
        ValueReferenceExpression(token::Metadata meta, std::string name) : Expression{ meta }, m_name{ name } {}
        virtual ~ValueReferenceExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    class BinaryOperationExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_lhs;
        types::BinOp m_binop;
        std::unique_ptr<Expression> m_rhs;
    public:
        BinaryOperationExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> lhs,
            types::BinOp op,
            std::unique_ptr<Expression> rhs)
            : Expression{ meta }
            , m_lhs{ std::move(lhs) }
            , m_binop{ op }
            , m_rhs{ std::move(rhs) } {
        }
        virtual ~BinaryOperationExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    class FunctionCallExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_fn_expr;
        std::vector<std::unique_ptr<Expression>> m_args;
    public:
        FunctionCallExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> fn_expr,
            std::vector<std::unique_ptr<Expression>> args)
            : Expression{ meta }
            , m_fn_expr{ std::move(fn_expr) }
            , m_args{ std::move(args) } {
        }
        virtual ~FunctionCallExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual typecheck::TypecheckResult typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };


    class ReturnStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        ReturnStatement(token::Metadata meta, std::unique_ptr<Expression> expr)
            : Statement{ meta }, m_expr{ std::move(expr) } {
        }
        virtual ~ReturnStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class InitializationStatement : public Statement {
    private:
        std::shared_ptr<types::Type> m_type;
        std::string m_name;
        std::optional<std::unique_ptr<Expression>> m_expr;
    public:
        InitializationStatement(
            token::Metadata meta,
            std::shared_ptr<types::Type> ty,
            std::string name,
            std::optional<std::unique_ptr<Expression>> expr)
            : Statement{ meta }
            , m_type{ std::move(ty) }
            , m_name{ name }
            , m_expr{ std::move(expr) } {
        }
        virtual ~InitializationStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class ExpressionStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        ExpressionStatement(token::Metadata meta, std::unique_ptr<Expression> expr)
            : Statement{ meta }, m_expr{ std::move(expr) } {
        }
        virtual ~ExpressionStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class IfStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_condition;
        std::unique_ptr<Statement> m_then;
        std::optional<std::unique_ptr<Statement>> m_else;
    public:
        IfStatement(token::Metadata meta,
            std::unique_ptr<Expression> condition,
            std::unique_ptr<Statement> then_statement,
            std::optional<std::unique_ptr<Statement>> else_statement)
            : Statement{ meta }
            , m_condition{ std::move(condition) }
            , m_then{ std::move(then_statement) }
            , m_else{ std::move(else_statement) } {
        }
        virtual ~IfStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class TopLevelStatement : public Node {
    public:
        TopLevelStatement(token::Metadata meta) : Node{ meta } {}
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) = 0;
    };

    class Function : public TopLevelStatement {
    private:
        std::shared_ptr<types::Type> m_return_ty;
        std::vector<std::pair<std::optional<std::string>, std::shared_ptr<types::Type>>> m_params;
        bool m_is_vararg;
        std::string m_name;
        std::optional<std::vector<std::unique_ptr<Statement>>> m_statements;
    public:
        Function(
            token::Metadata meta,
            std::shared_ptr<types::Type> return_ty,
            std::vector<std::pair<std::optional<std::string>, std::shared_ptr<types::Type>>> params,
            bool is_vararg,
            std::string name,
            std::optional<std::vector<std::unique_ptr<Statement>>> statements)
            : TopLevelStatement{ meta }
            , m_return_ty{ return_ty }
            , m_params{ params }
            , m_is_vararg{ is_vararg }
            , m_name{ name }
            , m_statements{ std::move(statements) } {
        }
        virtual ~Function() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };
}

#endif