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

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: {

                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 newProgram;
                    }
                    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 newProgram;
                    }
                    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 newProgram;
                            }
                            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 newProgram;
                            }
                            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 newProgram;
                }

                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 newProgram;
                }

                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);

                break;
            }
            case GroundInstruction_STRUCT: {
                // TODO: Preprocess structs
                break;
            }
            default: {
                Ground.Program.append(&newProgram, program->at[i]);
            }
        }
    }

    return newProgram;

}