/*
 * transpiler from catskill to the c programming language.
 * takes a direct catskill syntax tree and produces a c source file,
 * without an intermediate representation.
 *
 * Copyright (c) 2025-2026, Mel G. <mel@rnrd.eu>
 *
 * SPDX-License-Identifier: MPL-2.0
 */

#pragma once

#include "catboot.h"

// the output of the transpilation pass.
// used for both file output and string buffer output.
struct Transpile_Output
{
    FILE* file;
    struct String_Buffer string;
};

struct Transpile_Output
transpile_output_from_file(FILE* file)
{
    check(file != nil, "transpile output file is nil");
    return (struct Transpile_Output){
        .file = file,
        .string = string_buffer_empty(),
    };
}

#define TRANSPILE_OUTPUT_MAX_LENGTH 262144 // 256 KiB

struct Transpile_Output
transpile_output_from_string(void)
{
    return (struct Transpile_Output){
        .file = nil,
        .string = string_buffer_new(TRANSPILE_OUTPUT_MAX_LENGTH),
    };
}

#define TRANSPILE_OUTPUT_MAX_WRITE_LENGTH 4096

void
transpile_output_write(struct Transpile_Output* output, const ascii* format, ...)
{
    va_list args;
    va_start(args, format);

    if (output->file) {
        vfprintf(output->file, format, args);
    } else {
        ascii buffer[TRANSPILE_OUTPUT_MAX_WRITE_LENGTH];
        int written = vsnprintf(buffer, TRANSPILE_OUTPUT_MAX_WRITE_LENGTH, format, args);
        check(written >= 0 && (uint)written < TRANSPILE_OUTPUT_MAX_LENGTH, "transpile output string buffer overflow");

        string_buffer_append_c_str(&output->string, buffer);
    }

    va_end(args);
}

#define TRANSPILE_WRITE(...) \
    transpile_output_write(&transpiler->output, __VA_ARGS__)

struct String
transpile_output_string(struct Transpile_Output* output)
{
    check(output->file == nil, "transpile output is not a string");
    return string_buffer_to_string(&output->string);
}

FILE*
transpile_output_file(struct Transpile_Output* output)
{
    check(output->file != nil, "transpile output is not a file");
    return output->file;
}

struct Transpile_Output
transpile_output_to_string(void)
{
    return (struct Transpile_Output){
        .file = nil,
        .string = string_buffer_empty(),
    };
}

struct Transpile_Context
{
    bool in_function;
    struct String function_name;
    struct String function_return_type;
    bool main_function_found;
    bool main_function_takes_args;
};

#define CONTEXT_START(name)        \
    context = transpiler->context; \
    transpiler->context.name = true;

#define CONTEXT_END(name) transpiler->context.name = context.name;

struct Transpiler
{
    struct Transpile_Output output;
    struct Transpile_Context context;
};

#define TRANSPILER_PREAMBLE                                 \
    DATA_FOR_VISIT(struct Transpiler, transpiler)           \
    struct Transpile_Context context = transpiler->context; \
    (void)context;

void
transpiler_new(struct Transpiler* transpiler, struct Transpile_Output output)
{
    *transpiler = (struct Transpiler){
        .output = output,
        .context = {
            .in_function = false,
            .function_name = string_empty(),
            .function_return_type = string_empty(),
            .main_function_found = false,
            .main_function_takes_args = false,
        },
    };
}

void
transpiler_visit_type_node(struct Visit* visit, struct Type_Node* node)
{
    TRANSPILER_PREAMBLE

    if (!node || node->type == TYPE_NODE_NONE) {
        TRANSPILE_WRITE("void");
        return;
    }

    switch (node->type) {
    case TYPE_NODE_NAME: {
        struct String name = node->value.name.name;
        if (strcmp(name.data, "int") == 0) {
            TRANSPILE_WRITE("integer");
        } else if (strcmp(name.data, "string") == 0) {
            TRANSPILE_WRITE("struct String");
        } else if (strcmp(name.data, "bool") == 0) {
            TRANSPILE_WRITE("bool");
        } else if (strcmp(name.data, "float") == 0) {
            TRANSPILE_WRITE("real");
        } else if (strcmp(name.data, "uint") == 0) {
            TRANSPILE_WRITE("uint");
        } else if (strcmp(name.data, "byte") == 0) {
            TRANSPILE_WRITE("byte");
        } else if (strcmp(name.data, "ascii") == 0) {
            TRANSPILE_WRITE("ascii");
        } else {
            TRANSPILE_WRITE("%.*s", (int)name.length, name.data);
        }
        break;
    }
    default:
        TRANSPILE_WRITE("/* unknown type %d */", node->type);
        break;
    }
}

void
transpiler_visit_function_header_node(struct Visit* visit, struct Function_Header_Node* header)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("(");
    struct Type_Node* param = header->parameters_type_and_name;
    if (!param) {
        TRANSPILE_WRITE("void");
    } else {
        bool first = true;
        while (param) {
            if (!first) { TRANSPILE_WRITE(", "); }
            VISIT(visit_type_node, param);
            TRANSPILE_WRITE(" %.*s", (int)param->value_name.length, param->value_name.data);
            first = false;
            param = param->next;
        }
    }
    TRANSPILE_WRITE(")");
}

void
transpiler_visit_block_node(struct Visit* visit, struct Block_Node* node)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("{\n");
    FOR_EACH (struct Statement*, statement, node->statements) { VISIT(visit_statement, statement); }
    TRANSPILE_WRITE("}\n");
}

void
transpiler_visit_argument_group_node(struct Visit* visit, struct Argument_Group_Node* node)
{
    TRANSPILER_PREAMBLE

    struct Expression* arg = node->arguments;
    bool first = true;
    while (arg) {
        if (!first) { TRANSPILE_WRITE(", "); }
        VISIT(visit_expression, arg);
        first = false;
        arg = arg->next;
    }
}

void
transpiler_visit_statement_declaration(struct Visit* visit, struct Statement* stmt)
{
    TRANSPILER_PREAMBLE

    struct Statement_Value_Declaration* declaration = &stmt->value.declaration;
    struct Expression* initializer = declaration->inner.initializer;

    if (initializer && initializer->kind == EXPRESSION_FUNCTION) {
        struct Expression_Function* fun = &initializer->value.function;
        struct Function_Header_Node* header = &fun->header;

        VISIT(visit_type_node, header->return_type);
        struct String name = *array_at(struct String, &declaration->inner.names, 0);
        TRANSPILE_WRITE(" %.*s", (int)name.length, name.data);
        VISIT(visit_function_header_node, header);
        TRANSPILE_WRITE(" ");
        VISIT(visit_block_node, &fun->body);
    } else {
        if (declaration->kind == STATEMENT_DECLARATION_CONSTANT) {
            TRANSPILE_WRITE("const ");
        }
        VISIT(visit_type_node, declaration->inner.type);
        struct String name = *array_at(struct String, &declaration->inner.names, 0);
        TRANSPILE_WRITE(" %.*s", (int)name.length, name.data);
        if (initializer) {
            TRANSPILE_WRITE(" = ");
            VISIT(visit_expression, initializer);
        }
    }
}

void
transpiler_visit_statement_return(struct Visit* visit, struct Statement* stmt)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("return");
    if (stmt->value.return_value.value) {
        TRANSPILE_WRITE(" ");
        VISIT(visit_expression, stmt->value.return_value.value);
    }
}

void
transpiler_visit_statement_conditional(struct Visit* visit, struct Statement* stmt)
{
    TRANSPILER_PREAMBLE

    struct Statement_Value_Conditional* conditional = &stmt->value.conditional;
    for (uint i = 0; i < conditional->condition_count; ++i) {
        struct Statement_Conditional_Branch* branch = &conditional->conditions[i];
        if (i > 0) { TRANSPILE_WRITE("else "); }
        if (branch->when) {
            TRANSPILE_WRITE("if (");
            VISIT(visit_expression, branch->when);
            TRANSPILE_WRITE(") ");
        }
        VISIT(visit_block_node, &branch->then);
    }
}

void
transpiler_visit_statement(struct Visit* visit, struct Statement* statement)
{
    TRANSPILER_PREAMBLE

    switch (statement->kind) {
    case STATEMENT_EXPRESSION:
        VISIT(visit_expression, statement->value.expression.inner);
        if (transpiler->context.in_function) TRANSPILE_WRITE(";\n");
        break;
    case STATEMENT_DECLARATION:
        VISIT(visit_statement_declaration, statement);
        if (transpiler->context.in_function) TRANSPILE_WRITE(";\n");
        break;
    case STATEMENT_RETURN:
        VISIT(visit_statement_return, statement);
        if (transpiler->context.in_function) TRANSPILE_WRITE(";\n");
        break;
    case STATEMENT_BLOCK:
        VISIT(visit_block_node, &statement->value.block.inner);
        break;
    default:
        walk_statement(visit, statement);
        break;
    }
}

void
transpiler_visit_expression_integer_literal(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("%ld", expr->value.integer_literal.value);
}

void
transpiler_visit_expression_float_literal(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("%f", expr->value.float_literal.value);
}

void
transpiler_visit_expression_string_literal(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("\"%.*s\"", (int)expr->value.string_literal.value.length,
            expr->value.string_literal.value.data);
}

void
transpiler_visit_expression_boolean_literal(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("%s", expr->value.bool_literal.value ? "true" : "false");
}

void
transpiler_visit_expression_name(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("%.*s", (int)expr->value.name.name.length, expr->value.name.name.data);
}

void
transpiler_visit_expression_unary_operation(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    TRANSPILE_WRITE("(%s", unary_operation_to_string(expr->value.unary_operator.operation));
    VISIT(visit_expression, expr->value.unary_operator.operand);
    TRANSPILE_WRITE(")");
}

void
transpiler_visit_expression_binary_operation(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE
    struct Expression_Binary_Operator* bin_op = &expr->value.binary_operator;

    if (bin_op->operation == BINARY_ASSIGN && bin_op->right_operand->kind == EXPRESSION_FUNCTION) {
        struct Expression* fun_expr = bin_op->right_operand;
        struct Expression_Function* fun = &fun_expr->value.function;
        struct Function_Header_Node* header = &fun->header;
        struct Expression* name_expr = bin_op->left_operand;

        VISIT(visit_type_node, header->return_type);
        TRANSPILE_WRITE(" ");

        // check if this is a main function assignment
        if (name_expr && name_expr->kind == EXPRESSION_NAME) {
            struct Expression_Name* name = &name_expr->value.name;
            if (string_equals_c_str(name->name, "main")) {
                transpiler->context.main_function_found = true;
                transpiler->context.main_function_takes_args =
                    header->parameters_type_and_name != nil;
                TRANSPILE_WRITE("catskill_main");
            } else {
                VISIT(visit_expression, name_expr);
            }
        } else {
            VISIT(visit_expression, name_expr);
        }
        VISIT(visit_function_header_node, header);
        TRANSPILE_WRITE(" ");

        CONTEXT_START(in_function);
        VISIT(visit_block_node, &fun->body);
        CONTEXT_END(in_function);
    } else {
        enum Binary_Operation op = expr->value.binary_operator.operation;
        if (op == BINARY_ASSIGN_AND || op == BINARY_ASSIGN_OR) {
            TRANSPILE_WRITE("(");
            VISIT(visit_expression, expr->value.binary_operator.left_operand);
            TRANSPILE_WRITE(" = ");
            VISIT(visit_expression, expr->value.binary_operator.left_operand);
            if (op == BINARY_ASSIGN_AND)
                TRANSPILE_WRITE(" && ");
            else
                TRANSPILE_WRITE(" || ");
            VISIT(visit_expression, expr->value.binary_operator.right_operand);
            TRANSPILE_WRITE(")");
        } else {
            TRANSPILE_WRITE("(");
            VISIT(visit_expression, expr->value.binary_operator.left_operand);
            TRANSPILE_WRITE(" %s ", binary_operation_to_string(op));
            VISIT(visit_expression, expr->value.binary_operator.right_operand);
            TRANSPILE_WRITE(")");
        }
    }
}

void
transpiler_visit_expression_call(struct Visit* visit, struct Expression* expr)
{
    TRANSPILER_PREAMBLE

    VISIT(visit_expression, expr->value.call.subject);
    TRANSPILE_WRITE("(");
    VISIT(visit_argument_group_node, &expr->value.call.argument_group);
    TRANSPILE_WRITE(")");
}

void
transpiler_visit_expression(struct Visit* visit, struct Expression* expression)
{
    TRANSPILER_PREAMBLE

    switch (expression->kind) {
    case EXPRESSION_INTEGER_LITERAL:
        VISIT(visit_expression_integer_literal, expression);
        break;
    case EXPRESSION_FLOAT_LITERAL:
        VISIT(visit_expression_float_literal, expression);
        break;
    case EXPRESSION_STRING_LITERAL:
        VISIT(visit_expression_string_literal, expression);
        break;
    case EXPRESSION_BOOLEAN_LITERAL:
        VISIT(visit_expression_boolean_literal, expression);
        break;
    case EXPRESSION_NAME:
        VISIT(visit_expression_name, expression);
        break;
    case EXPRESSION_UNARY_OPERATION:
        VISIT(visit_expression_unary_operation, expression);
        break;
    case EXPRESSION_BINARY_OPERATION:
        VISIT(visit_expression_binary_operation, expression);
        break;
    case EXPRESSION_CALL:
        VISIT(visit_expression_call, expression);
        break;
    default:
        walk_expression(visit, expression);
        break;
    }
}

void
transpiler_visit_tree(struct Visit* visit, struct Tree* tree)
{
    TRANSPILER_PREAMBLE

    // include the catskill runtime core library
    // which provides all the necessary types and functions
    // for transpiled catskill programs.
    // other headers can be included by the user
    // with the pragma `| c-header "header.h"`.
    // TODO: for now we just reference the path to it in this repo
    // exactly and tell the backend in `./build.c` to look for includes there,
    // but in the real implementation we should embed all runtime files into
    // this executable and then write them out into the temporary build directory.
    TRANSPILE_WRITE("#include \"core.c\"\n");

    FOR_EACH (struct Statement*, statement, tree->top_level_statements) {
        VISIT(visit_statement, statement);
    }

    // check if we found a main function and define the appropriate macro for the runtime
    // TODO: create a nice lookup table of all the functions and types we found through
    // the catskill source.
    if (transpiler->context.main_function_found) {
        if (transpiler->context.main_function_takes_args) {
            TRANSPILE_WRITE("\n#define CATSKILL_MAIN_TAKES_ARGS\n");
        }
        TRANSPILE_WRITE("#include \"runtime.c\"\n");
    }
}

struct Visit_Table transpiler_visit_functions = {
    .visit_tree = transpiler_visit_tree,
    .visit_statement = transpiler_visit_statement,
    .visit_statement_declaration = transpiler_visit_statement_declaration,
    .visit_statement_conditional = transpiler_visit_statement_conditional,
    .visit_statement_return = transpiler_visit_statement_return,
    .visit_expression = transpiler_visit_expression,
    .visit_expression_integer_literal = transpiler_visit_expression_integer_literal,
    .visit_expression_float_literal = transpiler_visit_expression_float_literal,
    .visit_expression_string_literal = transpiler_visit_expression_string_literal,
    .visit_expression_boolean_literal = transpiler_visit_expression_boolean_literal,
    .visit_expression_name = transpiler_visit_expression_name,
    .visit_expression_unary_operation = transpiler_visit_expression_unary_operation,
    .visit_expression_binary_operation = transpiler_visit_expression_binary_operation,
    .visit_expression_call = transpiler_visit_expression_call,
    .visit_type_node = transpiler_visit_type_node,
    .visit_function_header_node = transpiler_visit_function_header_node,
    .visit_block_node = transpiler_visit_block_node,
    .visit_argument_group_node = transpiler_visit_argument_group_node,
};

int
transpiler_catskill_to_c(struct Transpiler* transpiler, struct Tree* tree)
{
    struct Visit visit = { .table = &transpiler_visit_functions, .user_data = transpiler };
    visit_table_fill_defaults(visit.table);

    walk(&visit, tree);

    return 0;
}