#include "../../include/ground.h"
#include <uthash.h>

static inline void doLabels(GroundProgram* program, GroundState* state) {

    for (GroundSize i = 0; i < program->len; i++) {
        GroundInstruction* instruction = &program->at[i];
        if (instruction->type == GroundInstruction_CREATELABEL) {
            if (instruction->args.len > 0) {
                GroundArg* arg = &instruction->args.at[0];
                GroundLabel* label = NULL;
                HASH_FIND_STR(state->labels, arg->as.ref->string, label);

                if (label == NULL) {
                    label = malloc(sizeof(GroundLabel));
                    if (label == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doLabels");
                        Ground.Flags.error = true;
                        return;
                    }
                    strncpy(label->name, arg->as.ref->string, 2047);
                    label->name[2047] = '\0';
                    label->lineNum = i;
                    HASH_ADD_STR(state->labels, name, label);
                } 
            }
        }
    }

    for (GroundSize i = 0; i < program->len; i++) {
        GroundInstruction* instruction = &program->at[i];
        if (instruction->type == GroundInstruction_JUMP) {
            GroundArg* arg = &instruction->args.at[0];
            GroundSize* line = Ground.State.findLabel(state, arg->as.ref->string);
            if (line == NULL) {
                char buf[2048];
                snprintf(buf, 2047, "couldn't find label '%s' (instruction JUMP at %zu) in Ground.New.Bytecode -> doLabels", arg->as.ref->string, i);
                Ground.Log.Error(buf);
                Ground.Flags.error = true;
                return;
            }

            arg->_offset = *line;

        } else if (instruction->type == GroundInstruction_IF) {
            GroundArg* arg = &instruction->args.at[1];
            GroundSize* line = Ground.State.findLabel(state, arg->as.ref->string);
            if (line == NULL) {
                char buf[2048];
                snprintf(buf, 2047, "couldn't find label '%s' (instruction IF at %zu) in Ground.New.Bytecode -> doLabels", arg->as.ref->string, i);
                Ground.Log.Error(buf);
                Ground.Flags.error = true;
                return;
            }

            arg->_offset = *line;
            
        } else if (instruction->type == GroundInstruction_CATCH) {
            GroundArg* arg = &instruction->args.at[1];
            GroundSize* line = Ground.State.findLabel(state, arg->as.ref->string);
            if (line == NULL) {
                char buf[2048];
                snprintf(buf, 2047, "couldn't find label '%s' (instruction CATCH at %zu) in Ground.New.Bytecode -> doLabels", arg->as.ref->string, i);
                Ground.Log.Error(buf);
                Ground.Flags.error = true;
                return;
            }

            arg->_offset = *line;

        }
    }

}

/*
 * Assigns an offset to each variable referenced.
 */
static inline GroundSize doOffsets(GroundProgram* gp, GroundState* state, GroundBytecodeProgram* gbp) {
    GroundSize size = 0;

    // Add size of state and assign offsets to pre-existing variables
    GroundVariable *s, *tmp;
    HASH_ITER(hh, state->variables, s, tmp) {
        s->_offset = size++;
    }

    for (GroundSize i = 0; i < gp->len; i++) {
        // Handle instructions which deal with struct fields
        switch (gp->at[i].type) {
            case GroundInstruction_STRUCT: {
                GroundInstruction* inst = &gp->at[i];
                GroundVariable* gsv = NULL;
                HASH_FIND_STR(state->variables, inst->args.at[0].as.ref->string, gsv);
                if (gsv == NULL) {
                    Ground.Log.Error("unexpected null variable in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                if (gsv->value.type.type != GroundType_Struct) {
                    Ground.Log.Error("unexpected null variable in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                GroundStruct* gs = &gsv->value.as.Struct;
                
                // Set offset for each struct field 
                GroundObjectField *s, *tmp;
                GroundSize fi = 0;
                HASH_ITER(hh, gs->fields, s, tmp) {
                    s->offset = fi++;
                }

                // Create bytecode instruction
                GroundBytecodeInstruction newInst = {
                    .type = GroundInstruction_STRUCT,
                    .args = {
                        .at = malloc(sizeof(GroundSize) * inst->args.len),
                        .capacity = inst->args.len,
                        .len = inst->args.len
                    }
                };

                if (newInst.args.at == NULL) {
                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }

                // Resolve the struct type variable to a heap offset
                GroundArg* arg = &inst->args.at[0];
                GroundVariable* structItem = NULL;
                HASH_FIND_STR(state->variables, arg->as.ref->string, structItem);
                if (structItem == NULL) {
                    structItem = malloc(sizeof(GroundVariable));
                    if (structItem == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    structItem->value = Ground.New.Value.Int(0);
                    strncpy(structItem->name, arg->as.ref->string, 2047);
                    structItem->_offset = size++;
                    HASH_ADD_STR(state->variables, name, structItem);
                }
                newInst.args.at[0] = structItem->_offset;

                // Now go through and add offsets to everything (skip struct name at index 0)
                for (GroundSize j = 1; j < inst->args.len; j++) {
                    // struct field name
                    arg = &inst->args.at[j];
                    GroundObjectField* field = NULL;
                    HASH_FIND_STR(gs->fields, arg->as.ref->string, field);
                    if (field == NULL) {
                        Ground.Log.Error("unexpected null struct field in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    newInst.args.at[j] = field->offset;
                    // 2nd arg is the op code, what we do next depends on it
                    j++;
                    // 0 -> none
                    // 1 -> init 
                    // 2 -> set
                    // 3 -> attach closure to function
                    GroundSize op = inst->args.at[j]._offset;
                    newInst.args.at[j] = op;
                    switch (op) {
                        case 0: case 3: break; // these do not require extra values
                        case 1:                // but these do, so we process them
                        case 2: {
                            j++;
                            arg = &inst->args.at[j];
                            GroundVariable* item = NULL;

                            if (arg->type == GroundArg_Value) {
                                // Literal value: create a constant variable for it
                                item = malloc(sizeof(GroundVariable));
                                if (item == NULL) {
                                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                                    Ground.Flags.error = true;
                                    return 0;
                                }
                                item->value = Ground.Copy.Value(&arg->as.value);
                                item->_offset = size++;
                                snprintf(item->name, sizeof(item->name) - 1, "_._.ground_internal_constant_%zu", item->_offset);
                                HASH_ADD_STR(state->variables, name, item);
                            } else {
                                // Variable reference: look it up or create it
                                HASH_FIND_STR(state->variables, arg->as.ref->string, item);

                                if (item == NULL) {
                                    item = malloc(sizeof(GroundVariable));
                                    if (item == NULL) {
                                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                                        Ground.Flags.error = true;
                                        return 0;
                                    }
                                    item->value = Ground.New.Value.Int(0);
                                    strncpy(item->name, arg->as.ref->string, 2047);
                                    item->_offset = size++;
                                    HASH_ADD_STR(state->variables, name, item);
                                }
                            }

                            newInst.args.at[j] = item->_offset;
                        }
                    }
                }

                gbp->at[i] = newInst;
                gbp->len++;

                continue;
            }
            case GroundInstruction_GETFIELD: {
                // getfield $parent -type (&field -type)... &field &output
                // The last field does not need a type annotation.
                // Bytecode: parent_offset, field0_offset, ..., fieldF_offset, output_offset
                GroundInstruction* inst = &gp->at[i];
                GroundSize gbiSize = 2 + ((inst->args.len - 2) / 2);

                GroundBytecodeInstruction newInst = {
                    .type = GroundInstruction_GETFIELD,
                    .args = {
                        .at = malloc(sizeof(GroundSize) * gbiSize),
                        .capacity = gbiSize,
                        .len = gbiSize
                    }
                };

                if (newInst.args.at == NULL) {
                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }

                // Resolve the parent variable to a heap offset
                GroundArg* arg = &inst->args.at[0];

                GroundVariable* item = NULL;
                HASH_FIND_STR(state->variables, arg->as.ref->string, item);

                if (item == NULL) {
                    item = malloc(sizeof(GroundVariable));
                    if (item == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    item->value = Ground.New.Value.Int(0);
                    strncpy(item->name, arg->as.ref->string, 2047);
                    item->_offset = size++;
                    HASH_ADD_STR(state->variables, name, item);
                }

                arg->_offset = item->_offset;
                newInst.args.at[0] = arg->_offset;

                // Get the type of the parent (type annotation is at args[1])
                GroundVariable* var = NULL;
                HASH_FIND_STR(state->variables, inst->args.at[1].as.ref->string, var);
                if (var == NULL) {
                    Ground.Log.Error("unknown struct name in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                if (var->value.type.type != GroundType_Struct) {
                    Ground.Log.Error("known name is not a struct in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }

                GroundStruct* currentType = &var->value.as.Struct;

                // Resolve each field's offset within its containing struct
                for (GroundSize j = 0; j < gbiSize - 2; j++) {
                    GroundArg* arg = &inst->args.at[2 + (j * 2)];

                    GroundObjectField* field;
                    HASH_FIND_STR(currentType->fields, arg->as.ref->string, field);
                    if (field == NULL) {
                        Ground.Log.Error("unknown field in struct in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    newInst.args.at[1 + j] = field->offset;

                    // Resolve the type for the next field (not needed for the last field)
                    if (j < gbiSize - 3) {
                        GroundVariable* typeVar = NULL;
                        HASH_FIND_STR(state->variables, inst->args.at[3 + j * 2].as.ref->string, typeVar);
                        if (typeVar == NULL) {
                            Ground.Log.Error("unknown struct type for field in Ground.New.Bytecode -> doOffsets");
                            Ground.Flags.error = true;
                            return 0;
                        }
                        if (typeVar->value.type.type != GroundType_Struct) {
                            Ground.Log.Error("field type is not a struct in Ground.New.Bytecode -> doOffsets");
                            Ground.Flags.error = true;
                            return 0;
                        }
                        currentType = &typeVar->value.as.Struct;
                    }
                }

                // Resolve the output variable
                GroundArg* outputArg = &inst->args.at[inst->args.len - 1];
                GroundVariable* outputItem = NULL;
                HASH_FIND_STR(state->variables, outputArg->as.ref->string, outputItem);
                if (outputItem == NULL) {
                    outputItem = malloc(sizeof(GroundVariable));
                    if (outputItem == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    outputItem->value = Ground.New.Value.Int(0);
                    strncpy(outputItem->name, outputArg->as.ref->string, 2047);
                    outputItem->_offset = size++;
                    HASH_ADD_STR(state->variables, name, outputItem);
                }
                outputArg->_offset = outputItem->_offset;
                newInst.args.at[gbiSize - 1] = outputArg->_offset;

                gbp->at[i] = newInst;
                gbp->len++;

                continue;
            }
            case GroundInstruction_SETFIELD: {
                // setfield $parent -type (&field -type)... &field $value
                // Bytecode: parent_offset, field0_offset, ..., fieldF_offset, value_offset
                GroundInstruction* inst = &gp->at[i];
                GroundSize gbiSize = 2 + ((inst->args.len - 2) / 2);

                GroundBytecodeInstruction newInst = {
                    .type = GroundInstruction_SETFIELD,
                    .args = {
                        .at = malloc(sizeof(GroundSize) * gbiSize),
                        .capacity = gbiSize,
                        .len = gbiSize
                    }
                };

                if (newInst.args.at == NULL) {
                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }

                // Resolve the parent variable to a heap offset
                GroundArg* arg = &inst->args.at[0];

                GroundVariable* item = NULL;
                HASH_FIND_STR(state->variables, arg->as.ref->string, item);

                if (item == NULL) {
                    item = malloc(sizeof(GroundVariable));
                    if (item == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    item->value = Ground.New.Value.Int(0);
                    strncpy(item->name, arg->as.ref->string, 2047);
                    item->_offset = size++;
                    HASH_ADD_STR(state->variables, name, item);
                }

                newInst.args.at[0] = item->_offset;

                // Get the type of the parent (type annotation is at args[1])
                GroundVariable* var = NULL;
                HASH_FIND_STR(state->variables, inst->args.at[1].as.ref->string, var);
                if (var == NULL) {
                    Ground.Log.Error("unknown struct name in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                if (var->value.type.type != GroundType_Struct) {
                    Ground.Log.Error("known name is not a struct in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }

                GroundStruct* currentType = &var->value.as.Struct;

                // Resolve each field's offset within its containing struct
                for (GroundSize j = 0; j < gbiSize - 2; j++) {
                    GroundArg* arg = &inst->args.at[2 + (j * 2)];

                    GroundObjectField* field;
                    HASH_FIND_STR(currentType->fields, arg->as.ref->string, field);
                    if (field == NULL) {
                        Ground.Log.Error("unknown field in struct in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    newInst.args.at[1 + j] = field->offset;

                    // Resolve the type for the next field (not needed for the last field)
                    if (j < gbiSize - 3) {
                        GroundVariable* typeVar = NULL;
                        HASH_FIND_STR(state->variables, inst->args.at[3 + j * 2].as.ref->string, typeVar);
                        if (typeVar == NULL) {
                            Ground.Log.Error("unknown struct type for field in Ground.New.Bytecode -> doOffsets");
                            Ground.Flags.error = true;
                            return 0;
                        }
                        if (typeVar->value.type.type != GroundType_Struct) {
                            Ground.Log.Error("field type is not a struct in Ground.New.Bytecode -> doOffsets");
                            Ground.Flags.error = true;
                            return 0;
                        }
                        currentType = &typeVar->value.as.Struct;
                    }
                }

                // Resolve the value variable
                GroundArg* valueArg = &inst->args.at[inst->args.len - 1];
                GroundVariable* valueItem = NULL;
                HASH_FIND_STR(state->variables, valueArg->as.ref->string, valueItem);
                if (valueItem == NULL) {
                    valueItem = malloc(sizeof(GroundVariable));
                    if (valueItem == NULL) {
                        Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                        Ground.Flags.error = true;
                        return 0;
                    }
                    valueItem->value = Ground.New.Value.Int(0);
                    strncpy(valueItem->name, valueArg->as.ref->string, 2047);
                    valueItem->_offset = size++;
                    HASH_ADD_STR(state->variables, name, valueItem);
                }
                newInst.args.at[gbiSize - 1] = valueItem->_offset;

                gbp->at[i] = newInst;
                gbp->len++;

                continue;
            }
            case GroundInstruction_CALLMETHOD: {

                continue;
            }
            default: break;
        }
        for (GroundSize j = 0; j < gp->at[i].args.len; j++) {
            GroundArg* arg = &gp->at[i].args.at[j];

            if (
                arg->type == GroundArg_Label ||
                arg->type == GroundArg_LineRef
               ) continue;

            if (arg->type == GroundArg_Value) {
                GroundVariable* item = malloc(sizeof(GroundVariable));
                if (item == NULL) {
                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                item->value = Ground.Copy.Value(&arg->as.value);
                if (Ground.Flags.error) return 0;
                item->_offset = size++;
                snprintf(item->name, sizeof(item->name) - 1, "_._.ground_internal_constant_%zu", item->_offset);
                HASH_ADD_STR(state->variables, name, item);
                arg->_offset = item->_offset;
                continue;
            }

            GroundVariable* item = NULL;
            HASH_FIND_STR(state->variables, arg->as.ref->string, item);

            if (item == NULL) {
                item = malloc(sizeof(GroundVariable));
                if (item == NULL) {
                    Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
                    Ground.Flags.error = true;
                    return 0;
                }
                item->value = Ground.New.Value.Int(0);
                strncpy(item->name, arg->as.ref->string, 2047);
                item->_offset = size++;
                HASH_ADD_STR(state->variables, name, item);
            }

            arg->_offset = item->_offset;
        }

        // Convert to GroundBytecodeInstruction
        GroundBytecodeInstruction inst = {
            .type = gp->at[i].type,
            .args.at = malloc(sizeof(GroundSize) * gp->at[i].args.len),
            .args.len = gp->at[i].args.len,
            .args.capacity = gp->at[i].args.len
        };

        if (inst.args.at == NULL) {
            Ground.Log.Error("malloc failed in Ground.New.Bytecode -> doOffsets");
            Ground.Flags.error = true;
            return 0;
        }

        for (GroundSize j = 0; j < gp->at[i].args.len; j++) {
            inst.args.at[j] = gp->at[i].args.at[j]._offset;
        }

        gbp->at[i] = inst;
        gbp->len++;

    }
    return size;
}


GroundBytecode _GroundNewBytecode(GroundProgram* program, GroundState* state) {

    GroundBytecode bytecode = {
        .program = {
            .at = malloc(sizeof(GroundBytecodeInstruction) * program->len),
            .capacity = program->len,
            .len = 0,
        },
        .heap = {
            .heap = NULL,
            .capacity = 0,
            .len = 0
        }
    };

    if (bytecode.program.at == NULL) {
        Ground.Log.Error("malloc failed in Ground.New.Bytecode");
        Ground.Flags.error = true;
        return bytecode;
    }

    // Add core types to state 
    for (unsigned int i = 0; i < (unsigned int)GroundType_Object; i++) {
        GroundVariable* var = malloc(sizeof(GroundVariable));
        if (var == NULL) {
            Ground.Log.Error("malloc failed in Ground.New.Bytecode");
            Ground.Flags.error = true;
            return bytecode;
        }
        var->value = (GroundValue) {
            .type = {GroundType_CoreType},
            .as.CoreType = (enum GroundTypeType) i
        };

        switch ((enum GroundTypeType) i) {
            case GroundType_Int: {
                sprintf(var->name, "int");
                break;
            }
            case GroundType_Double: {
                sprintf(var->name, "double");
                break;
            }
            case GroundType_Char: {
                sprintf(var->name, "char");
                break;
            }
            case GroundType_Bool: {
                sprintf(var->name, "bool");
                break;
            }
            case GroundType_String: {
                sprintf(var->name, "string");
                break;
            }
            case GroundType_List: {
                sprintf(var->name, "list");
                break;
            }
            case GroundType_Function: {
                sprintf(var->name, "function");
                break;
            }
            case GroundType_Struct: {
                sprintf(var->name, "struct");
                break;
            }
            default: break; // unreachable (in theory)
        }

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

    }

    state->_size += (unsigned int)GroundType_Object + 1;

    // Precompute some really cool stuff
    doLabels(program, state);
    if (Ground.Flags.error) return bytecode;

    GroundSize size = doOffsets(program, state, &bytecode.program);
    if (Ground.Flags.error) return bytecode;

    // Allocate heap
    bytecode.heap.heap = malloc(sizeof(GroundBytecodeValue) * size);
    if (bytecode.heap.heap == NULL) {
        Ground.Log.Error("malloc failed in Ground.New.Bytecode");
        Ground.Flags.error = true;
        return bytecode;
    }
    bytecode.heap.capacity = size;
    bytecode.heap.len = size;

    GroundVariable *s, *tmp;

    // Copy constants into heap at their assigned offsets
    HASH_ITER(hh, state->variables, s, tmp) {
        bytecode.heap.heap[s->_offset] = Ground.New.BytecodeValue(Ground.Copy.Value(&s->value));
        if (Ground.Flags.error) return bytecode;
    }


    return bytecode;
    
}