#include "compiler.h"
#include "types.h"
#include "include/uthash.h"

#include <tram.h>

// Helper macros to reduce repetition
#define CHECK_ARGC(expected) \
    if (instruction->args.length != (expected)) { \
        fprintf(stderr, "Error at instruction %zu (%s): expected %d argument(s), got %zu\n", \
                i, #expected, (expected), instruction->args.length); \
        return false; \
    }

#define CHECK_ARGC_MIN(min) \
    if (instruction->args.length < (min)) { \
        fprintf(stderr, "Error at instruction %zu: expected at least %d argument(s), got %zu\n", \
                i, (min), instruction->args.length); \
        return false; \
    }

#define CHECK_ARG_TYPE(idx, expected_type) \
    if (instruction->args.args[(idx)].type != (expected_type)) { \
        fprintf(stderr, "Error at instruction %zu: argument %d should be " #expected_type ", got %d\n", \
                i, (idx), instruction->args.args[(idx)].type); \
        return false; \
    }

bool checkForErrors(GroundProgram* program) {
    for (size_t i = 0; i < program->size; i++) {
        GroundInstruction* instruction = &program->instructions[i];
        switch (instruction->type) {

            // if $condition %label
            case IF:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, LINEREF);
                break;

            // jump %label
            case JUMP:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, LINEREF);
                break;

            // end $status
            case END:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // input &variable
            case INPUT:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, DIRREF);
                break;

            // print $value
            case PRINT:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // println $value
            case PRINTLN:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // set &variable $value
            case SET:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, VALREF);
                break;

            // gettype $value &variable
            case GETTYPE:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // exists &variable &output
            case EXISTS:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // setlist &varname $val1 $val2 ...  (at least 1 arg, rest are valrefs)
            case SETLIST:
                CHECK_ARGC_MIN(1);
                CHECK_ARG_TYPE(0, DIRREF);
                for (size_t j = 1; j < instruction->args.length; j++) {
                    if (instruction->args.args[j].type != VALREF) {
                        fprintf(stderr, "Error at instruction %zu (SETLIST): argument %zu should be VALREF, got %d\n",
                                i, j, instruction->args.args[j].type);
                        return false;
                    }
                }
                break;

            // setlistat &listname $index $value
            case SETLISTAT:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, VALREF);
                break;

            // getlistat &list $index &variable
            case GETLISTAT:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // getlistsize &list &variable
            case GETLISTSIZE:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // listappend $value &list
            case LISTAPPEND:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // getstrsize $string &variable
            case GETSTRSIZE:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // getstrcharat $string $index &variable
            case GETSTRCHARAT:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // add $value $value &variable
            case ADD:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // subtract $value $value &variable
            case SUBTRACT:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // multiply $value $value &variable
            case MULTIPLY:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // divide $value $value &variable
            case DIVIDE:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // equal $value $value &variable
            case EQUAL:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // inequal $value $value &variable
            case INEQUAL:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // not $value &variable
            case NOT:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // greater $value $value &variable
            case GREATER:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // lesser $value $value &variable
            case LESSER:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, VALREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // stoi $value &variable
            case STOI:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // stod $value &variable
            case STOD:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // itoc, ctoi $value &variable
            case ITOC:
            case CTOI:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // tostring $value &variable
            case TOSTRING:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                break;

            // fun !functionName -returnType [-argType &arg ...]
            // Minimum 2 args: the function ref and return type
            case FUN:
                CHECK_ARGC_MIN(2);
                CHECK_ARG_TYPE(0, FNREF);
                CHECK_ARG_TYPE(1, TYPEREF);
                // Remaining args must be alternating TYPEREF, DIRREF pairs
                if ((instruction->args.length - 2) % 2 != 0) {
                    fprintf(stderr, "Error at instruction %zu (FUN): argument pairs after return type must be -type &name pairs\n", i);
                    return false;
                }
                for (size_t j = 2; j < instruction->args.length; j += 2) {
                    CHECK_ARG_TYPE(j,     TYPEREF);
                    CHECK_ARG_TYPE(j + 1, DIRREF);
                }
                break;

            // return $value
            case RETURN:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // endfun - no args
            case ENDFUN:
                CHECK_ARGC(0);
                break;

            // call !function [$arg1 $arg2 ...] &returnVal
            // Minimum 2 args: function ref and return variable
            case CALL:
                CHECK_ARGC_MIN(2);
                CHECK_ARG_TYPE(0, FNREF);
                // Middle args are valrefs, last arg is dirref
                for (size_t j = 1; j < instruction->args.length - 1; j++) {
                    if (instruction->args.args[j].type != VALREF) {
                        fprintf(stderr, "Error at instruction %zu (CALL): argument %zu should be VALREF, got %d\n",
                                i, j, instruction->args.args[j].type);
                        return false;
                    }
                }
                CHECK_ARG_TYPE(instruction->args.length - 1, DIRREF);
                break;

            // call &object !methodName [$arg1 $arg2 ...] &returnVal
            // Minimum 3 args: object ref, method name, return variable
            case CALLMETHOD:
                CHECK_ARGC_MIN(3);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, FNREF);
                // Middle args are valrefs, last arg is dirref
                for (size_t j = 2; j < instruction->args.length - 1; j++) {
                    if (instruction->args.args[j].type != VALREF) {
                        fprintf(stderr, "Error at instruction %zu (CALL): argument %zu should be VALREF, got %d\n",
                                i, j, instruction->args.args[j].type);
                        return false;
                    }
                }
                CHECK_ARG_TYPE(instruction->args.length - 1, DIRREF);
                break;

            // struct -structName
            case STRUCT:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, TYPEREF);
                break;

            // endstruct - no args
            case ENDSTRUCT:
                CHECK_ARGC(0);
                break;

            // init &var -type
            case INIT:
                CHECK_ARGC(2);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, TYPEREF);
                break;

            // getfield $object &fieldName &outputVar
            case GETFIELD:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, VALREF);
                CHECK_ARG_TYPE(1, DIRREF);
                CHECK_ARG_TYPE(2, DIRREF);
                break;

            // setfield &object &fieldName $value
            case SETFIELD:
                CHECK_ARGC(3);
                CHECK_ARG_TYPE(0, DIRREF);
                CHECK_ARG_TYPE(1, DIRREF);
                CHECK_ARG_TYPE(2, VALREF);
                break;

            // use $libraryName
            case USE:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // extern $libraryName
            case EXTERN:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, VALREF);
                break;

            // createlabel @label
            case CREATELABEL:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, LABEL);
                break;

            // pause, licence - no args
            case PAUSE:
            case LICENSE:
                break;
            // drop &variable
            case DROP:
                CHECK_ARGC(1);
                CHECK_ARG_TYPE(0, DIRREF);
                break;
        }
    }
    return true;
}

void compileGroundProgram(GroundProgram* program, char* name) {

    if (!checkForErrors(program)) {
        printf("Error: Invalid program\n");
        return;
    }

    Tram_Program tramProgram = Tram_Program_Create();

    Tram_ParameterList parameters = Tram_ParameterList_Create(1, (Tram_Parameter[]){Tram_Parameter_Variable("main")});
    Tram_Program_AddInstruction(&tramProgram, Tram_Instruction_Create(Tram_InstructionType_CreateLabel, parameters));

    for (size_t i = 0; i < program->size; i++) {
        compileGroundInstruction(&program->instructions[i], &tramProgram);
    }

    Tram_Compiler* compiler = Tram_Compiler_Create(tramProgram);
    Tram_Compiler_SetLogLevel(compiler, Tram_LogLevel_Debug);
    Tram_Compiler_SetTarget(compiler, Tram_Target_x86_64_Linux_libc);
    Tram_Compiler_Compile(compiler);

    if (Tram_Compiler_HasCompiled(compiler)) {
        Tram_Compiler_AssembleAndLink(compiler, name, (Tram_LinkerArgs){.size = 0, .data = NULL});
    } else {
        printf("Compilation failed\n");
        exit(EXIT_FAILURE);
    }

}

typedef struct GroundCompilerVariable {
    GroundValueType type;
    char id[MAX_ID_LEN];
    UT_hash_handle hh;
} GroundCompilerVariable;

typedef struct GroundState {
    GroundCompilerVariable* variables;
    GroundLabel* labels;
} GroundState;

GroundState createGroundState() {
    GroundState state = {
        .variables = NULL,
        .labels = NULL
    };
    return state;
}

GroundCompilerVariable* cFindVariable(GroundState* state, const char* id) {
    GroundCompilerVariable* var;
    HASH_FIND_STR(state->variables, id, var);
    return var;
}

GroundCompilerVariable* cDeleteVariable(GroundState* state, const char* id) {
    GroundCompilerVariable* var;
    HASH_FIND_STR(state->variables, id, var);
    HASH_DEL(state->variables, var);
}

GroundCompilerVariable* cAddVariable(GroundState* state, const char* id, GroundValueType type) {
    GroundCompilerVariable* var = malloc(sizeof(GroundCompilerVariable));
    snprintf(var->id, MAX_ID_LEN, "%s", id);
    var->type = type;
    HASH_ADD_STR(state->variables, id, var);
    return var;
}

void compileGroundInstruction(GroundInstruction* instruction, Tram_Program* program) {
    switch (instruction->type) {
        case IF:
            break;
        case JUMP:
            break;
        case END:
            break;
        case INPUT:
            break;
        case PRINT:
            break;
        case PRINTLN:
            break;
        case SET:
            break;
        case GETTYPE:
            break;
        case EXISTS:
            break;
        case SETLIST:
            break;
        case SETLISTAT:
            break;
        case GETLISTAT:
            break;
        case GETLISTSIZE:
            break;
        case LISTAPPEND:
            break;
        case GETSTRSIZE:
            break;
        case GETSTRCHARAT:
            break;
        case ADD:
            break;
        case SUBTRACT:
            break;
        case MULTIPLY:
            break;
        case DIVIDE:
            break;
        case EQUAL:
            break;
        case INEQUAL:
            break;
        case NOT:
            break;
        case GREATER:
            break;
        case LESSER:
            break;
        case AND:
            break;
        case OR:
            break;
        case XOR:
            break;
        case NEG:
            break;
        case SHIFT:
            break;
        case STOI:
            break;
        case STOD:
            break;
        case ITOC:
            break;
        case CTOI:
            break;
        case TOSTRING:
            break;
        case FUN:
            break;
        case RETURN:
            break;
        case ENDFUN:
            break;
        case PUSHARG:
            break;
        case CALL:
            break;
        case STRUCT:
            break;
        case ENDSTRUCT:
            break;
        case INIT:
            break;
        case GETFIELD:
            break;
        case SETFIELD:
            break;
        case USE:
            break;
        case EXTERN:
            break;
        case CREATELABEL: {
            Tram_ParameterList parameters = Tram_ParameterList_Create(1, (Tram_Parameter[]){Tram_Parameter_Variable("main")});
            Tram_Program_AddInstruction(program, Tram_Instruction_Create(Tram_InstructionType_CreateLabel, parameters));
            break;
        }
        case PAUSE:
            break;
        case DROP:
            break;
        case LICENSE:
            break;
        case ERRORCMD:
            break;
    }
}