#include "../../include/ground.h"

static void doFunction(GroundProgram* program, GroundProgram* newProgram, GroundState* state, GroundSize* i) {

    GroundFunction function = Ground.New.Function(state);

    // Parse signature
    GroundInstruction* sig = &program->at[*i];
    char* id = sig->args.at[0].as.ref->string;

    for (GroundSize j = 1; j < sig->args.len; j++) {
        GroundArg* arg = &sig->args.at[j];

        // Ignore type annotations, this can be taken care of in another function
        if (arg->type == GroundArg_TypeRef) {
            continue;
        }
        char* argId = arg->as.ref->string;
        Ground.Function.appendArg(&function, argId);

        GroundVariable* argVar = malloc(sizeof(GroundVariable));
        if (argVar == NULL) {
            Ground.Flags.error = true;
            Ground.Log.Error("malloc failed in Ground.Program.Preprocess");
            return;
        }
        snprintf(argVar->name, 2047, "%s", argId);
        argVar->value = Ground.New.Value.Int(0);
        HASH_ADD_STR(function.closure->variables, name, argVar);
    }

    // Add instructions to function
    GroundSize funCount = 1;
    GroundSize structCount = 0;

    for (;;) {
        (*i)++;
        if (*i >= program->len) {
            Ground.Flags.error = true;
            Ground.Log.Error("expecting ENDFUN instruction, reached end of program in Ground.Program.Preprocess");
            return;
        }
        GroundInstruction* inst = &program->at[*i];
        switch (inst->type) {
            case GroundInstruction_FUN: {
                funCount++;
                break;
            }
            case GroundInstruction_ENDFUN: {
                if (funCount < 1) {
                    Ground.Flags.error = true;
                    Ground.Log.Error("extra ENDFUN instruction (expecting ENDSTRUCT before) in Ground.Program.Preprocess");
                    return;
                }
                funCount--;
                break;
            }
            case GroundInstruction_STRUCT: {
                structCount++;
                break;
            }
            case GroundInstruction_ENDSTRUCT: {
                if (structCount < 1) {
                    Ground.Flags.error = true;
                    Ground.Log.Error("extra ENDSTRUCT instruction (expecting ENDFUN before) in Ground.Program.Preprocess");
                    return;
                }
                structCount--;
                break;
            }
            default: break;
        }

        if (funCount == 0 && structCount == 0) {
            // We done
            break;
        }

        Ground.Function.appendInstruction(&function, *inst);
    }

    // Preprocess function's body
    GroundProgram functionBody = Ground.Program.preprocess(function.program.ground, function.closure);
    if (Ground.Flags.error) {
        return;
    }

    Ground.Free.Program(function.program.ground);
    *function.program.ground = functionBody;

    GroundVariable* var = malloc(sizeof(GroundVariable));
    if (var == NULL) {
        Ground.Flags.error = true;
        Ground.Log.Error("malloc failed in Ground.Program.Preprocess");
        return;
    }

    snprintf(var->name, 2047, "%s", id);
    var->value = Ground.New.Value.Function(function);

    HASH_ADD_STR(state->variables, name, var);

    // Add instruction to attach closure to function
    GroundInstruction inst = Ground.New.Instruction(GroundInstruction_FUN);
    Ground.Instruction.append(&inst, Ground.New.Arg.FunctionRef(sig->args.at[0].as.ref));
    Ground.Program.append(newProgram, inst);
}

static void doStruct(GroundProgram* program, GroundProgram* newProgram, GroundState* state, GroundSize* i);
static void doStruct(GroundProgram* program, GroundProgram* newProgram, GroundState* state, GroundSize* i) {
    GroundIdentifier* structName = Ground.Copy.Identifier(program->at[*i].args.at[0].as.ref);
    GroundStruct gs = Ground.New.Struct();

    // The aim is to compile an instruction which looks like this:
    // struct &field OP $value &field OP $value &field OP $value...
    // where:
    //      &field is the field to initialise
    //      OP is a numeric identifier:
    //          0 -> none
    //          1 -> init 
    //          2 -> set
    //          3 -> attach closure to function
    //      $value is either a type reference (OP=0), a value reference (OP=1), or none (OP=2)
    // Functions will be moved to the end of the struct, and a closure will be attached to all of them
    //
    // Oh yeah, we also need to make a GroundStruct value

    GroundSize structCount = 0;

    GroundInstruction output = Ground.New.Instruction(GroundInstruction_STRUCT);
    Ground.Instruction.append(&output, Ground.New.Arg.TypeRef(structName));

    // Parse struct body 
    bool parsing = true;
    while (parsing) {
        (*i)++;
        if (*i >= program->len) {
            Ground.Flags.error = true;
            Ground.Log.Error("expecting ENDFUN instruction, reached end of program in Ground.Program.Preprocess");
            return;
        }

        GroundInstruction* instruction = &program->at[*i];
        switch (instruction->type) {
            case GroundInstruction_SET: {

                GroundObjectField* field = malloc(sizeof(GroundObjectField));
                if (field == NULL) {
                    Ground.Log.Error("malloc failed in Ground.Program.preprocess -> doStruct");
                    Ground.Flags.error = true;
                    return;
                } 

                snprintf(field->name, 2047, "%s", instruction->args.at[0].as.ref->string);

                // If we know the value type now, add it to the instruction 
                if (instruction->args.at[1].type == GroundArg_Value) {
                    field->value = Ground.Copy.Value(&instruction->args.at[1].as.value);
                    break;
                }

                HASH_ADD_STR(gs.fields, name, field);

                // set &field $value -> &field 2 $value
                Ground.Instruction.append(&output, Ground.Copy.Arg(&instruction->args.at[0]));
                Ground.Instruction.append(&output, (GroundArg) {.type = GroundArg_DirectRef, ._offset = 2});
                Ground.Instruction.append(&output, Ground.Copy.Arg(&instruction->args.at[1]));
                break;
            }
            case GroundInstruction_INIT: {
                GroundObjectField* field = malloc(sizeof(GroundObjectField));
                if (field == NULL) {
                    Ground.Log.Error("malloc failed in Ground.Program.preprocess -> doStruct");
                    Ground.Flags.error = true;
                    return;
                } 

                snprintf(field->name, 2047, "%s", instruction->args.at[0].as.ref->string);

                // If we know the type now, add it to the instruction 
                if (strcmp(field->name, "int") == 0) {
                    field->value = Ground.New.Value.Int(0);
                    break;
                }
                if (strcmp(field->name, "double") == 0) {
                    field->value = Ground.New.Value.Double(0.0);
                    break;
                }
                if (strcmp(field->name, "string") == 0) {
                    field->value = Ground.New.Value.String(Ground.New.String(""));
                    break;
                }
                if (strcmp(field->name, "char") == 0) {
                    field->value = Ground.New.Value.Char(0);
                    break;
                }
                if (strcmp(field->name, "bool") == 0) {
                    field->value = Ground.New.Value.Bool(false);
                    break;
                }
                if (strcmp(field->name, "function") == 0) {
                    field->value = Ground.New.Value.Function(Ground.New.Function(state));
                    break;
                }

                HASH_ADD_STR(gs.fields, name, field);

                // init &field -type -> &field 1 -type
                Ground.Instruction.append(&output, Ground.Copy.Arg(&instruction->args.at[0]));
                Ground.Instruction.append(&output, (GroundArg) {.type = GroundArg_DirectRef, ._offset = 1});
                Ground.Instruction.append(&output, Ground.Copy.Arg(&instruction->args.at[1]));
                break;
            }
            case GroundInstruction_STRUCT: {
                // TODO structs in structs
                break;
            }
            case GroundInstruction_EXTERN: {
                // TODO external functions as struct members
                break;
            }
            case GroundInstruction_FUN: {
                // Parse the function itself
                doFunction(program, newProgram, state, i);
                // Read and remove the emitted instruction 
                GroundInstruction* inst = &newProgram->at[newProgram->len - 1];
                Ground.Instruction.append(&output, Ground.New.Arg.FunctionRef(Ground.Copy.Identifier(inst->args.at[0].as.ref)));
                Ground.Free.Instruction(inst);
                newProgram->len--;

                // Add 3 to signal addition of closure 
                Ground.Instruction.append(&output, (GroundArg) {.type = GroundArg_DirectRef, ._offset = 1});
                break;
            }
            case GroundInstruction_ENDSTRUCT: {
                parsing = false;
            }
            default: {
                Ground.Log.Error("invalid instruction inside struct in ");
            }
        }

    }

    GroundVariable* variable = malloc(sizeof(GroundVariable));
    if (variable == NULL) {
        Ground.Log.Error("malloc failed in Ground.Program.Preprocess -> doStruct");
        Ground.Flags.error = true;
        return;
    }

    snprintf(variable->name, 2047, "%s", structName->string);
    variable->value = Ground.New.Value.Struct(gs);

    HASH_ADD_STR(state->variables, name, variable);

}

GroundProgram _GroundProgramPreprocess(GroundProgram* program, GroundState* state) {

    GroundProgram newProgram = Ground.New.Program();

    for (GroundSize i = 0; i < program->len; i++) {
        switch (program->at[i].type) {
            case GroundInstruction_FUN: {
                doFunction(program, &newProgram, state, &i);
                if (Ground.Flags.error) return newProgram;
                break;
            }
            case GroundInstruction_STRUCT: {
                doStruct(program, &newProgram, state, &i);
                if (Ground.Flags.error) return newProgram;
                break;
            }
            default: {
                Ground.Program.append(&newProgram, program->at[i]);
            }
        }
    }

    return newProgram;

}