#include "../resources/resources.h"
#include "puzzle.h"
#include "../wren_inter/wren_inter.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <wren.h>
#include <stdbool.h>

static const char *USER_MODULE_NAME = "main";

/// Returns 1 test puzzle for the GetTestPuzzle segment
Test *getTestPuzzleTests() {
    Test *outp = malloc(sizeof(Test));
    outp->RecieverName = strdup("main");
    outp->MainMethodSignature = strdup("call(_)");

    outp->argc = 1;
    outp->argv = malloc(sizeof(PuzzleIO));
    outp->argv->type = WREN_TYPE_STRING;
    outp->argv->value = strdup("Hello");

    outp->ExpectedResult.type = WREN_TYPE_STRING;
    outp->ExpectedResult.value = strdup("Hello");
    return outp;
}

Puzzle *GetTestPuzzle() {
    Puzzle *outp = malloc(sizeof(Puzzle));
    if (outp == NULL) return NULL;

    outp->Name = strdup("Test puzzle");
    outp->Description = strdup("In this puzzle you will have to return the input string");
    outp->InitialCode = strdup( "var main = Fn.new { |name| \n    System.print(\"Hello, %(name)!\")\n}\n\nmain.call(\"world\")");
    outp->TestC = 1;
    outp->Tests = getTestPuzzleTests();

    return outp;
}

void FreePuzzleIO(PuzzleIO *io) {
    if (io->value != NULL)
        free(io->value);
}

void FreeTests(Test *Tests, size_t TestC) {
    if (TestC == 0 || Tests == NULL) return;
    for (size_t i = 0; i < TestC; i++) {
        Test *t = &Tests[i];
        if (t->RecieverName != NULL) free(t->RecieverName);
        if (t->MainMethodSignature != NULL) free(t->MainMethodSignature);
        for (size_t j = 0; j < t->argc; j++) {
            FreePuzzleIO(&t->argv[j]);
        }
        if (t->argv != NULL) free(t->argv);
    }
    free(Tests);
}

void FreePuzzle(Puzzle *self) {
    if (self == NULL) return;
    if (self->Name != NULL) free(self->Name);
    if (self->Description != NULL) free(self->Description);
    if (self->InitialCode != NULL) free(self->InitialCode);
    FreeTests(self->Tests, self->TestC);
    free(self);
}

void PrintArgumentToSB(PuzzleIO *io, StringBuilder *sb) {
    if (sb == NULL || io == NULL) return;
    char buf[128];
    switch (io->type) {
        case (WREN_TYPE_BOOL):
            bool value = *(bool *)io->value;
            AppendStringBuilder(sb, value ? "true" : "false");
            return;
        case (WREN_TYPE_NUM):
            sprintf(buf, "%lf", *(double *)io->value) ;
            AppendStringBuilder(sb, buf);
            break;
        case (WREN_TYPE_LIST):
            AppendStringBuilder(sb, "[TODO]");
            break;
        case (WREN_TYPE_MAP):
            // TODO
            return;
        case (WREN_TYPE_NULL):
            AppendStringBuilder(sb, "null");
            break;
        case (WREN_TYPE_STRING):
            sprintf(buf, "\"%s\"", (char *)io->value) ;
            AppendStringBuilder(sb, buf);
            break;
        default:
            return;
    }
}

bool PrintTestCallToSB(Test *test, StringBuilder *sb) {
    if (test == NULL || test->MainMethodSignature == NULL || test->RecieverName == NULL)
        return false;
    AppendStringBuilder(sb, test->RecieverName);
    AppendStringBuilder(sb, ".");
    size_t sig_len = strlen(test->MainMethodSignature);
    size_t arg_ptr = 0;
    for (size_t i = 0; i < sig_len; i++) {
        char ch = test->MainMethodSignature[i];

        if (ch != '_') {
            char buf[2];
            buf[0] = ch; buf[1] = 0;
            AppendStringBuilder(sb, buf);
            continue;
        }
        if (arg_ptr >= test->argc) {
            AppendStringBuilder(sb, "null");
            continue;
        }

        PuzzleIO arg = test->argv[arg_ptr];
        PrintArgumentToSB(&arg, sb);
        arg_ptr += 1;
    }
    return true;
}

// We are currenttly 
void PutIOIntoSlot(int slot, PuzzleIO *io, WrenVM *vm) {
    switch (io->type) {
        case (WREN_TYPE_BOOL):
            wrenSetSlotBool(vm, slot, *(bool *)io->value);
            break;
        case (WREN_TYPE_NUM):
            wrenSetSlotDouble(vm, slot, *(double *)io->value);
            break;
        case (WREN_TYPE_LIST):
            wrenSetSlotDouble(vm, slot, *(double *)io->value);
            break;
        case (WREN_TYPE_MAP):
            // TODO
            return;
        case (WREN_TYPE_NULL):
            wrenSetSlotNull(vm, slot);
            break;
        case (WREN_TYPE_STRING):
            wrenSetSlotString(vm, slot, io->value);
            break;
        default:
            return;
    }
}

const char *GetTypeName(WrenType type) {
    switch (type) {
        case WREN_TYPE_BOOL: return "bool";
        case WREN_TYPE_NUM: return "num";
        case WREN_TYPE_FOREIGN: return "foreign";
        case WREN_TYPE_LIST: return "list";
        case WREN_TYPE_MAP: return "map";
        case WREN_TYPE_NULL: return "null";
        case WREN_TYPE_STRING: return "string";
        case WREN_TYPE_UNKNOWN: return "unknown";
    }
}

bool CheckResult(PuzzleIO *io, WrenVM *vm, char *outp) {
    WrenType actual_type = wrenGetSlotType(vm, 0);
    if (io->type != actual_type) {
        sprintf(outp, "Expected type %s, got %s", 
                GetTypeName(io->type),
                GetTypeName(actual_type)); 
        return false;
    }

    bool output = false;
    switch (io->type) {
        case (WREN_TYPE_BOOL): 
            {
                if (io->value == NULL) return false; // Invalid puzzle
                bool result = wrenGetSlotBool(vm, 0);
                output = result == *((bool *)io->value);
                if (!output)
                    sprintf(outp, "Expected '%s', got '%s'", 
                            result ? "true" : "false",
                            output ? "true" : "false");
                return output;
            }

        case (WREN_TYPE_NUM): 
            {
                if (io->value == NULL) return false; // Invalid puzzle
                double result = wrenGetSlotDouble(vm, 0);
                output = result == *((double *)io->value);
                if (!output)
                    sprintf(outp, "Expected '%lf', got '%lf'", 
                            *((double *)io->value),
                            result);
                return output;
            }

        case (WREN_TYPE_NULL):
            return actual_type == WREN_TYPE_NULL;

        case (WREN_TYPE_STRING):
            {
                if (io->value == NULL) return false; // Invalid puzzle
 
                const char *result = wrenGetSlotString(vm, 0);
                output = strcmp(result, io->value) == 0;
                if (!output)
                    sprintf(outp, "Expected '%s', got '%s'", 
                            (char *)io->value,
                            result);
                return output;
            }

        case (WREN_TYPE_LIST):
        case (WREN_TYPE_MAP):
        default:
            return false;
    }
}



TestResult RunTest(Test *test, WrenVM *vm, size_t max_op) {
    wrenEnsureSlots(vm, test->argc + 1); // finction handle + arguments
    wrenGetVariable(vm, USER_MODULE_NAME, test->RecieverName, 0);
    char outp_buf[512];

    WrenHandle *obj_handle = wrenGetSlotHandle(vm, 0);
    if (obj_handle == NULL) {
        sprintf(outp_buf, "No such class as %s", test->RecieverName);
        return (TestResult) {
            .res = RESULT_COMPTIME_NO_ENTRY,
            .message = strdup(outp_buf),
        };
    }

    WrenHandle *fn_handle = wrenMakeCallHandle(vm, test->MainMethodSignature);
    if (fn_handle == NULL) {
        wrenReleaseHandle(vm, obj_handle);
        sprintf(outp_buf, "No such method as %s in a class %s", 
                test->MainMethodSignature,
                test->RecieverName);

        return (TestResult) {
            .res = RESULT_COMPTIME_NO_ENTRY,
            .message = strdup(outp_buf),
        };
    }

    wrenSetSlotHandle(vm, 0, obj_handle);
    for (int i = 0; i < test->argc; ++i) {
        PutIOIntoSlot(i+1, &test->argv[i], vm);
    }

    WrenInterpretResult call_result = wrenCallLim(vm, fn_handle, max_op);

    switch (call_result) {
        case WREN_RESULT_COMPILE_ERROR:
            return (TestResult) { .res = RESULT_COMPTIME_ERR, .message = NULL };
        case WREN_RESULT_RUNTIME_ERROR: 
            return (TestResult) { .res = RESULT_RUNTIME_ERR, .message = NULL };
        case WREN_RESULT_RUNTIME_LIMIT:
            return (TestResult) { 
                .res = RESULT_RUNTIME_ERR, 
                .message = strdup("The execution halted")
            };
        default: break;
    }

    bool result = CheckResult(&test->ExpectedResult, vm, outp_buf);

    wrenReleaseHandle(vm, fn_handle);
    wrenReleaseHandle(vm, obj_handle);

    if (result) {
        return (TestResult) { .res = RESULT_OK, .message = NULL };
    }

    return (TestResult) { 
        .res = RESULT_WRONG_ANSWER, 
        .message = strdup(outp_buf) 
    };
}

// Returns the res ptr. Does not modify the address
SubmitionResult *TryEvaluateUserCode(WrenVM *vm, char *script, SubmitionResult *res, size_t max_op) {
    WrenInterpretResult result = wrenInterpretLim(vm, "main", script, max_op);
    res->result = RESULT_OK;
    switch (result) {
        case WREN_RESULT_COMPILE_ERROR: 
            res->result = RESULT_COMPTIME_ERR;
            return res;
        case WREN_RESULT_RUNTIME_ERROR:
            res->result = RESULT_RUNTIME_ERR;
            return res;
        case WREN_RESULT_RUNTIME_LIMIT:
            res->result = RESULT_TIMEOUT;
            printf("Timeout with max_op %zu\n", max_op);
            return res;
        default: break;
    }
    return res;
}

// Frees the vm. Does not modify the self address.
SubmitionResult *PrepareResult(SubmitionResult *self, 
        WrenVM *vm, 
        enum PUZZLE_TEST_RESULT result) {
    if (vm == NULL || self == NULL) return self;
    self->result = result;
    self->stdout_sb = GetMachineOutput(vm);
    InterpreterFreeMachine(vm);
    return self;
}

SubmitionResult *SubmitSolution(Puzzle *self, char *script) {
    Config *cfg = GetConfig();
    if (self == NULL || script == NULL) return NULL;
    SubmitionResult *res = malloc(sizeof(SubmitionResult));
    
    WrenVM *vm = InterpreterPrepareMachine();
    if (vm == NULL) {
        res->result = RESULT_INIT_ERROR;
        return res;
    }

    StringBuilder *sb = GetOutputSb(vm);
    AppendStringBuilder(sb, "Evaluating user code...\n");
    res = TryEvaluateUserCode(vm, script, res, cfg->MaxOperations);
    switch (res->result) {
        case RESULT_OK: break;
        case RESULT_COMPTIME_ERR: 
                        return PrepareResult(res, vm, res->result);
                        break;
        case RESULT_RUNTIME_ERR:
                        return PrepareResult(res, vm, res->result);
                        break;
        case RESULT_TIMEOUT:
                        AppendStringBuilder(sb, "The execution halted\n");
                        return PrepareResult(res, vm, res->result);
                        break;
        default: break;
    }
    AppendStringBuilder(sb, "Success!\n");

    for (size_t i = 0; i < self->TestC; i++) {
        Test *test = &self->Tests[i];
        AppendStringBuilder(sb, "\nRunning test ");
        PrintTestCallToSB(test, sb);
        AppendStringBuilder(sb, "... \n");
        TestResult test_res = RunTest(test, vm, cfg->MaxOperations);
        if (test_res.message != NULL) {
            AppendStringBuilder(sb, test_res.message);
            AppendStringBuilder(sb, "\n");
            free(test_res.message);
        }
        if (test_res.res != RESULT_OK) {
            AppendStringBuilder(sb, "Test failed.\n");
            return PrepareResult(res, vm, test_res.res);
        }
        AppendStringBuilder(sb, "Success!\n");
    }

    return PrepareResult(res, vm, RESULT_OK);
}

void FreeSunmitionResult(SubmitionResult *res) {
    if (res->stdout_sb != NULL) free(res->stdout_sb);
    free(res);
}

void PrintPuzzle(Puzzle *puzzle) {
    printf("Puzzle name: %s\n", puzzle->Name);
    printf("Puzzle description: %s\n", puzzle->Description);
    printf("Puzzle code: %s\n", puzzle->InitialCode);
    printf("Puzzle TestC: %zu\n", puzzle->TestC);

    for (int i = 0; i < puzzle->TestC; i++) {
        Test t = puzzle->Tests[i];
        printf("\nTest %d:\n", i);
        printf("\tTest rec: %s\n", t.RecieverName);
        printf("\tTest sig: %s\n", t.MainMethodSignature);
        printf("\tTest argc: %zu\n", t.argc);
        printf("\tArguments:\n");

        for (int j = 0; j < t.argc; j++) {
            PuzzleIO io = t.argv[j];
            switch (io.type) {
                case WREN_TYPE_BOOL:
                    printf("\t\t(bool)%s\n", *(bool *)io.value ? "true" : "false");
                    break;
                case WREN_TYPE_NUM:
                    printf("\t\t(num)%f\n", *(float *)io.value);
                    break;
                case WREN_TYPE_LIST:
                    printf("\t\t(list)NOT SUPPORTED\n");
                    break;
                case WREN_TYPE_MAP:
                    printf("\t\t(list)NOT SUPPORTED\n");
                    break;
                case WREN_TYPE_NULL:
                    printf("\t\t(null)\n");
                    break;
                case WREN_TYPE_STRING:
                    printf("\t\t(string)%s\n", (char *)io.value);
                    break;
                default: break;
            }
        }
    }
}