#pragma once
#include <iostream>
#include <sstream>
#include "string_utils.hpp"
#include "base.hpp"

template <typename T>
using ptr = std::unique_ptr<T>;

struct Assign : Stmt
{
    ptr<Expr> target;
    ptr<Expr> value;

    Assign(ptr<Expr> t, ptr<Expr> v)
        : target(std::move(t)), value(std::move(v)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "Assign\n";
        oss << target->dump(indent + 1) << "\n";
        oss << value->dump(indent + 1);
        return oss.str();
    }
};

struct ExprStmt : Stmt
{
    ptr<Expr> expr;

    ExprStmt(ptr<Expr> e) : expr(std::move(e)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "ExprStmt\n";
        oss << expr->dump(indent + 1);
        return oss.str();
    }
};

struct Return : Stmt
{
    ptr<Expr> value;

    Return(ptr<Expr> v) : value(std::move(v)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "Return\n";
        if (value)
        {
            oss << value->dump(indent + 1);
        }
        return oss.str();
    }
};

struct Pass : Stmt
{
    std::string dump(int indent = 0) const override
    {
        return indent_str(indent) + "Pass";
    }
};

// #todo elif
struct If : Stmt
{
    ptr<Expr> condition;
    std::vector<ptr<Stmt>> then_body;
    std::vector<ptr<Stmt>> else_body;

    If(ptr<Expr> cond, std::vector<ptr<Stmt>> then_b, std::vector<ptr<Stmt>> else_b)
        : condition(std::move(cond)),
          then_body(std::move(then_b)),
          else_body(std::move(else_b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "If\n";
        oss << condition->dump(indent + 1);
        for (const auto &stmt : then_body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        if (!else_body.empty())
        {
            oss << "\n"
                << indent_str(indent + 1) << "Else";
            for (const auto &stmt : else_body)
            {
                oss << "\n"
                    << stmt->dump(indent + 1);
            }
        }
        return oss.str();
    }
};

struct For : Stmt
{
    std::string target;
    ptr<Expr> iter;
    std::vector<ptr<Stmt>> body;

    For(std::string t, ptr<Expr> i, std::vector<ptr<Stmt>> b)
        : target(t), iter(std::move(i)), body(std::move(b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "For(" << target << ")\n";
        oss << iter->dump(indent + 1);
        for (const auto &stmt : body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        return oss.str();
    }
};

struct FunctionDef : Stmt
{
    std::string name;
    std::vector<std::string> params;
    std::vector<ptr<Stmt>> body;

    FunctionDef(std::string n, std::vector<std::string> p, std::vector<ptr<Stmt>> b)
        : name(n), params(std::move(p)), body(std::move(b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "FunctionDef(" << name << ")";
        for (const auto &param : params)
        {
            oss << "\n"
                << indent_str(indent + 1) << "Param(" << param << ")";
        }
        for (const auto &stmt : body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        return oss.str();
    }
};

struct ClassDef : Stmt
{
    std::string name;
    std::vector<ptr<Stmt>> body;

    ClassDef(std::string n, std::vector<ptr<Stmt>> b)
        : name(n), body(std::move(b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "ClassDef(" << name << ")";
        for (const auto &stmt : body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        return oss.str();
    }
};

struct Module : Node
{
    std::vector<ptr<Stmt>> body;

    Module(std::vector<ptr<Stmt>> b) : body(std::move(b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << "Module";
        for (const auto &stmt : body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        return oss.str();
    }
};

struct While : Stmt
{
    ptr<Expr> condition;
    std::vector<ptr<Stmt>> body;

    While(ptr<Expr> cond, std::vector<ptr<Stmt>> b)
        : condition(std::move(cond)), body(std::move(b)) {}

    std::string dump(int indent = 0) const override
    {
        std::ostringstream oss;
        oss << indent_str(indent) << "While\n";
        oss << condition->dump(indent + 1);
        for (const auto &stmt : body)
        {
            oss << "\n"
                << stmt->dump(indent + 1);
        }
        return oss.str();
    }
};