#pragma once
#include "ast.hpp"
#include <sstream>

class RubyGenerator
{
    std::ostringstream output;
    int indent_level = 0;

    void write_indent()
    {
        for (int i = 0; i < indent_level; i++)
        {
            output << "  ";
        }
    }

    void write(const std::string &s)
    {
        output << s;
    }

public:
    std::string generate(const Module &module)
    {
        output.str("");
        indent_level = 0;

        for (const auto &stmt : module.body)
        {
            generate_stmt(*stmt);
        }

        return output.str();
    }

private:
    void generate_stmt(const Stmt &stmt)
    {
        if (auto *expr_stmt = dynamic_cast<const ExprStmt *>(&stmt))
        {
            write_indent();
            generate_expr(*expr_stmt->expr);
            write("\n");
        }
        else if (auto *assign = dynamic_cast<const Assign *>(&stmt))
        {
            write_indent();
            generate_expr(*assign->target);
            write(" = ");
            generate_expr(*assign->value);
            write("\n");
        }
        else if (auto *if_stmt = dynamic_cast<const If *>(&stmt))
        {
            write_indent();
            write("if ");
            generate_expr(*if_stmt->condition);
            write("\n");

            indent_level++;
            for (const auto &s : if_stmt->then_body)
            {
                generate_stmt(*s);
            }
            indent_level--;

            if (!if_stmt->else_body.empty())
            {
                write("else\n");
                indent_level++;
                for (const auto &s : if_stmt->else_body)
                {
                    generate_stmt(*s);
                }
                indent_level--;
            }
            write_indent();
            write("end\n");
        }
        else if (auto *ret = dynamic_cast<const Return *>(&stmt))
        {
            write_indent();
            write("return ");
            if (ret->value)
            {
                generate_expr(*ret->value);
            }
            write("\n");
        }
        else if (auto *for_stmt = dynamic_cast<const For *>(&stmt))
        {
            write_indent();
            output << "for " << for_stmt->target << " in ";
            generate_expr(*for_stmt->iter);
            output << "\n";
            indent_level++;
            for (const auto &s : for_stmt->body)
            {
                generate_stmt(*s);
            }
            indent_level--;
            write_indent();
            write("end\n");
        }
        else if (auto *while_stmt = dynamic_cast<const While *>(&stmt))
        {
            write_indent();
            output << "while ";
            generate_expr(*while_stmt->condition);
            output << "\n";
            indent_level++;
            for (const auto &s : while_stmt->body)
            {
                generate_stmt(*s);
            }
            indent_level--;
            write_indent();
            write("end\n");
        }
        else if (auto *pass_stmt = dynamic_cast<const Pass *>(&stmt))
        {
            // #todo
        }
        else if (auto *break_stmt = dynamic_cast<const Break *>(&stmt))
        {
            write_indent();
            output << "break\n";
        }
        else if (auto *continue_stmt = dynamic_cast<const Continue *>(&stmt))
        {
            write_indent();
            output << "continue\n";
        }
        else if (auto *function_def = dynamic_cast<const FunctionDef *>(&stmt))
        {
            write_indent();
            output << "def " << function_def->name << "(";
            if (function_def->params.size() > 0)
            {
                output << function_def->params[0];
            }
            for (int i = 1; i < function_def->params.size(); i++)
            {
                output << ", " << function_def->params[i];
            }
            output << ")\n";
            indent_level++;
            for (const auto &s : function_def->body)
            {
                generate_stmt(*s);
            }
            indent_level--;
            write_indent();
            write("end\n");
        }
        else if (auto *class_def = dynamic_cast<const ClassDef *>(&stmt))
        {
            write_indent();
            output << "Class " << class_def->name << "\n";
            indent_level++;
            for (const auto &s : class_def->body)
            {
                generate_stmt(*s);
            }
            indent_level--;
            write_indent();
            write("end\n");
        }
    }
    void generate_expr(const Expr &expr)
    {
        if (auto *num = dynamic_cast<const Number *>(&expr))
        {
            write(num->value);
        }
        else if (auto *id = dynamic_cast<const Identifier *>(&expr))
        {
            write(id->name);
        }
        else if (auto *str = dynamic_cast<const String *>(&expr))
        {
            write("\"" + str->value + "\"");
        }
        else if (auto *binop = dynamic_cast<const BinOp *>(&expr))
        {
            generate_expr(*binop->left);
            write(" " + binop->op.symbol() + " ");
            generate_expr(*binop->right);
        }
        else if (auto *call = dynamic_cast<const Call *>(&expr))
        {
            generate_expr(*call->func);
            write("(");
            for (size_t i = 0; i < call->args.size(); i++)
            {
                if (i > 0)
                    write(", ");
                generate_expr(*call->args[i]);
            }
            write(")");
        }
        else if (auto *unary_op = dynamic_cast<const UnaryOpExpr *>(&expr))
        {
            output << unary_op->op.symbol();
            generate_expr(*unary_op->operand);
        }
        else if (auto *attr = dynamic_cast<const Attribute *>(&expr))
        {
            generate_expr(*attr->value);
            output << "." << attr->attr;
        }
        else if (auto *subscript = dynamic_cast<const Subscript *>(&expr))
        {
            generate_expr(*subscript->value);
            output << "[";
            generate_expr(*subscript->index);
            output << "]";
        }
        else if (auto *slice_sub = dynamic_cast<const SliceSubscript *>(&expr))
        {
            generate_expr(*slice_sub->value);
            write("[");

            auto &slice = *slice_sub->index;

            if (slice.step)
            {
                write("# TODO: slice with step not directly translatable");
            }
            else
            {
                if (slice.lower)
                {
                    generate_expr(*slice.lower);
                }
                else
                {
                    write("0");
                }

                write("...");

                if (slice.upper)
                {
                    generate_expr(*slice.upper);
                }
                else
                {
                    write("-1");
                }
            }

            write("]");
        }
        else if (auto *list = dynamic_cast<const List *>(&expr))
        {
            write("[");
            for (size_t i = 0; i < list->elements.size(); i++)
            {
                if (i > 0)
                    write(", ");
                generate_expr(*list->elements[i]);
            }
            write("]");
        }
        else if (auto *dict = dynamic_cast<const Dict *>(&expr))
        {
            write("{");
            for (size_t i = 0; i < dict->pairs.size(); i++)
            {
                if (i > 0)
                    write(", ");
                generate_expr(*dict->pairs[i].first);
                write(" => ");
                generate_expr(*dict->pairs[i].second);
            }
            write("}");
        }
        else if (auto *tuple = dynamic_cast<const Tuple *>(&expr))
        {
            // #note no direct translation available
            write("[");
            for (size_t i = 0; i < tuple->elements.size(); i++)
            {
                if (i > 0)
                    write(", ");
                generate_expr(*tuple->elements[i]);
            }
            write("]");
        }
    }
};