#include "../includes/parser.h"
#include <string.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>

Symbol* CreateSymbol(Token* token, int length, int address, int resolved);
const List* TokensList;
int CurrentToken = 0;
//int HeapStart = 4; // zero indexed, used for declared variables.
//int BinaryEnd = 4; // Zero indexed (binary starts with a 4 byte header)

IROpcode* CreateIROpcode(TokenType mnemonic);
Symbol* AddSymbol(Token* token, int length, int address, int resolved);
void PrintSymbols(void);
void HandleOperation(void);
void HandleAssemblerDirective(void);
void AdvanceParser(void);
Token* PeekToken(void);
int ParserAtEnd(void);
int Expect(int, ...);
Symbol* GetSymbol(char*);
//List* SymbolsTable;
unsigned int HeapTop = 4;
unsigned int ProgramCounter = 0;
// unsigned char Heap[HIGHMEMORY];
// unsigned char Memory[HIGHMEMORY];
IRState MachineState;

IRState* ParseTokens(List* tokens) {
    if (!tokens) return NULL;

    // memset(Memory, 0, sizeof(Memory));
    // memset(Heap, 0, sizeof(Heap));

    TokensList = tokens;

    MachineState.SymbolsTable = CreateList();
    MachineState.Opcodes = CreateList();

    while(!ParserAtEnd()) {
        Token* t = PeekToken();

        switch(t->token_class) {
            case Directive:
                HandleAssemblerDirective();
                break;
            case Mnemonic:
                HandleOperation();
                break;
            case Address:
                if (t->type == LABEL) {
                    printf("%s: \n", t->lexeme);
                    AddSymbol(t, 2, ProgramCounter, 1);
                }
                break;
            default:
                break;
        }

        AdvanceParser();
    }

    if (MachineState.SymbolsTable->size > 0) PrintSymbols();

    //DestroyList(SymbolsTable);

    //memcpy(&Heap[HeapTop], Memory, ProgramCounter);
    //int totalSize = HeapTop + ProgramCounter;
    // Heap[0] = 2 >> HeapTop & 0xFF;
    // Heap[1] = HeapTop & 0xFF;
    // Heap[2] = 2 >> totalSize & 0xFF;
    // Heap[3] = totalSize & 0xFF;

    // for(int i = 0; i < 33; i++) { 
    //     if (i != 0 && i % 3 == 0) printf("\n");
    //     printf("%02X ", Heap[i] & 0xFF);
    // }
    // printf("\n");
    // printf("Program Counter: %d\n", ProgramCounter);
    // printf("Heap Top: %#06X\n", HeapTop);

    return &MachineState;
}

void HandleAssemblerDirective(void) {

    if (PeekToken()->type == DB) {
        AdvanceParser();

        Token* identifier = PeekToken();

        if (!Expect(1, IDENTIFIER)) {
            fprintf(stderr, "Expected identifier on line %d\n", identifier->line);
            exit(1);
        }

        Symbol* symbol = AddSymbol(identifier, 0, HeapTop, 1);

        identifier = PeekToken();

        switch(identifier->type) {
            case STRING:
                //printf("[INFO] %s = '%s' (%#04X)\n", identifier->lexeme, identifier->lexeme, HeapTop);

                // for(int i = 0; i < strlen(identifier->lexeme); i++) {
                //     //Memory[ProgramCounter + i] = PeekToken()->lexeme[i];
                //     Heap[HeapTop + i] = identifier->lexeme[i];
                // }

                //ProgramCounter += strlen(PeekToken()->lexeme);
                HeapTop += strlen(identifier->lexeme);
                symbol->length = strlen(identifier->lexeme);

                AdvanceParser();

                if (PeekToken()->type == COMMA) {
                    AdvanceParser();

                    if (PeekToken()-> type != NUMBER) {
                        fprintf(stderr, "Expected string termination byte on line %d\n", PeekToken()->line);
                        exit(1);
                    }
                    else {
                        //Hideous...
                        // int oldLength = strlen(identifier->lexeme);
                        // identifier->lexeme = realloc(identifier->lexeme, oldLength + 2);
                        // identifier->value = identifier->lexeme;
                        // identifier->lexeme[oldLength + 1] = '\0';//*PeekToken()->lexeme & 0x0F;
                        symbol->length++; //Basically, read in the NULL byte at the end of the string.
                        HeapTop++;
                        return;
                    }
                }
            case NUMBER:
                break;
            default:
                fprintf(stderr, "[Line: %d] Expected either a number or string after identifier '%s'.\n", PeekToken()->line, identifier->lexeme);
                exit(1);
        }
    }
}

void HandleRegisterBasedOpcode(unsigned char mask) {
    //Memory[ProgramCounter] = mask;

    ProgramCounter++;

    Token* token = PeekToken();

    IROpcode* opcode = CreateIROpcode(token->type);

    AdvanceParser();

    token = PeekToken();

    if (token->token_class == Reg) {
        //Memory[ProgramCounter - 1] |= (PeekToken()->type - R1) & 0x07;
        opcode->ParameterOneType = Reg;
        AdvanceParser();
    }
    else if (token->token_class == Address && token->type == COPY) {
        opcode->ParameterOneType = Address;

        Symbol* symbol = GetSymbol(token->lexeme);
        
        if (!symbol) opcode->ParameterOne.Address = AddSymbol(token, 2, 0, 0);
        else opcode->ParameterOne.Address = symbol;//CreateSymbol(token, symbol->address, symbol->resolved);
    }
    else {
        fprintf(stderr, "[Error] Line %d: Expected register operand.\n", token->line);
        exit(1);
    }

    token = PeekToken();

    if (!Expect(1, COMMA)) {
        fprintf(stderr, "[Error] Expected comma on line %d.\n", token->line);
        exit(1);
    }

    token = PeekToken();

    if (token->type == IDENTIFIER || token->type == LABEL) {
        opcode->ParameterTwoType = Address;
        Symbol* symbol = GetSymbol(token->lexeme);

        // if (!symbol) {
        //     fprintf(stderr, "[Error] Line %d: %s is undefined.\n", token->line, token->lexeme);
        //     exit(1);
        // }
        if (!symbol) opcode->ParameterOne.Address = AddSymbol(token, 2, 0, 0);//CreateSymbol(token, 0, 0);
        else opcode->ParameterOne.Address = symbol;//CreateSymbol(token, symbol->address, symbol->resolved);

        // Memory[ProgramCounter - 1] |= 0x10;//0b00010000;

        // Memory[ProgramCounter] = 2 >> symbol->address & 0xFF;
        // Memory[ProgramCounter + 1] = symbol->address & 0xFF;

        AdvanceParser();
    }
    else if (token->type == NUMBER) {
        opcode->ParameterTwoType = Constant;

        opcode->ParameterTwo.Constant = *(int*)PeekToken()->value;
        // Memory[ProgramCounter - 1] |= 0x08; //0b00001000;

        // int* value = PeekToken()->value;

        // Memory[ProgramCounter] = 2 >> *value & 0xFF;
        // Memory[ProgramCounter + 1] = *value & 0xFF;

        AdvanceParser();
    }
    else if (token->token_class == Reg) {
        opcode->ParameterTwoType = Reg;
        opcode->ParameterTwo.Constant = (PeekToken()->type - R1) &0x07;
        //Memory[ProgramCounter] = (PeekToken()->type - R1) & 0x07;

        AdvanceParser();

        ProgramCounter++;
        return;
    } 
    else {
        fprintf(stderr, "[Error] Expected operand, got %s\n", token->lexeme);
        exit(1);
    }

    ProgramCounter += 2;

    AddListItem(opcode, sizeof(IROpcode), MachineState.Opcodes);
}

void HandleOperation(void) {
    Token* token = PeekToken();
    //const Instruction* inst = GetOpcodeDetails(opcode->type);
    const unsigned char registerBasedOpcodeMask = 0xE0;

    //if (!inst) return;

    // AdvanceParser();

    unsigned char mask = GetInstructionMask(token->type);

    if (mask & registerBasedOpcodeMask) {
        HandleRegisterBasedOpcode(mask);
        return;
    }

    IROpcode* opcode = CreateIROpcode(token->type);
    opcode->ParameterTwoType = None;

    AdvanceParser();

    switch (mask) {
        case 0: //NOP
            //Memory[ProgramCounter] = 0x00;
            opcode->ParameterOneType = None;

            ProgramCounter++;
            return;
        case 1: //JZ
            //Memory[ProgramCounter] = 0x01;

            ProgramCounter++;

            if (PeekToken()->type == IDENTIFIER || PeekToken()->type == LABEL) {
                Symbol* symbol = GetSymbol(PeekToken()->lexeme);
                opcode->ParameterOneType = Address;
                
                if (!symbol) opcode->ParameterOne.Address = AddSymbol(PeekToken(), 2, 0, 0);
                else opcode->ParameterOne.Address = symbol;
                // if (!symbol) {
                //     fprintf(stderr, "[Error] %s is undefined.\n", PeekToken()->lexeme);
                //     exit(1);
                // }

                // Memory[ProgramCounter] = 2 >> symbol->address & 0xFF;
                // Memory[ProgramCounter + 1] = symbol->address & 0xFF;
            }
            else if (PeekToken()->type == NUMBER) {
                opcode->ParameterOneType = Constant;
                opcode->ParameterOne.Constant = *(int*)PeekToken()->value;
                // int* value = PeekToken()->value;

                // Memory[ProgramCounter] = 2 >> *value & 0xFF;
                // Memory[ProgramCounter + 1] = *value & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line: %d: Expected address after jump if zero (JZ) instruction.\n", ((Token*) TokensList->content[CurrentToken - 1])->line);
                exit(1);
            }

            ProgramCounter += 2;

            return;
        case 2: //INT
            //Memory[ProgramCounter] = 0x02;

            ProgramCounter++;

            if (PeekToken()->type == NUMBER) {
                // int* value = PeekToken()->value;
                opcode->ParameterOneType = Constant;
                opcode->ParameterOne.Constant = *(int*)PeekToken()->value;

                // Memory[ProgramCounter] = 2 >> *value & 0xFF;
                // Memory[ProgramCounter + 1] = *value & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line %d: Expected Interrupt vector.\n", PeekToken()->line);
                exit(1);
            }

            ProgramCounter += 2;

            return;
        case 3: //YLD
            //Memory[ProgramCounter] = 0x03;
            opcode->ParameterOneType = None;
            
            ProgramCounter++;
            break;
        case 4: //RET
            //Memory[ProgramCounter] = 0x04;
            opcode->ParameterOneType = None;
            
            ProgramCounter++;
            break;
        case 5: //CALL
            //Memory[ProgramCounter] = 0x05;

            ProgramCounter++;

            if (PeekToken()->type == IDENTIFIER || PeekToken()->type == LABEL) {
                Symbol* symbol = GetSymbol(PeekToken()->lexeme);
                opcode->ParameterOneType = Address;

                // if (!symbol) {
                //     fprintf(stderr, "[Error] Line %d: %s is undefined.\n", PeekToken()->line, PeekToken()->lexeme);
                //     exit(1);
                // }
                if (!symbol) opcode->ParameterOne.Address = AddSymbol(PeekToken(), 2, 0, 0);
                else opcode->ParameterOne.Address = symbol;
                //Memory[ProgramCounter] = 2 >> symbol->address & 0xFF;
                //Memory[ProgramCounter + 1] = symbol->address & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line %d: Expected address or label.\n", PeekToken()->line);
                exit(1);
            }
            
            ProgramCounter += 2;
            break;
        case 6: //JMP
            //Memory[ProgramCounter] = 0x06;

            ProgramCounter++;

            if (PeekToken()->type == IDENTIFIER || PeekToken()->type == LABEL) {
                Symbol* symbol = GetSymbol(PeekToken()->lexeme);
                opcode->ParameterOneType = Address;

                // if (!symbol) {
                //     fprintf(stderr, "[Error] Line %d: %s is undefined.\n", PeekToken()->line, PeekToken()->lexeme);
                //     exit(1);
                // }
                if (!symbol) opcode->ParameterOne.Address = AddSymbol(PeekToken(), 2, 0, 0);
                else opcode->ParameterOne.Address = symbol;

                // Memory[ProgramCounter] = 2 >> symbol->address & 0xFF;
                // Memory[ProgramCounter + 1] = symbol->address & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line %d: Expected address or label.\n", PeekToken()->line);
                exit(1);
            }
            
            ProgramCounter += 2;
            break;
        case 7: //IN
            //Memory[ProgramCounter] = 0x07;

            ProgramCounter++;

            if (PeekToken()->type == NUMBER) {
                opcode->ParameterOneType = Constant;
                opcode->ParameterOne.Constant = *(int*)PeekToken()->value;
                // int* value = PeekToken()->value;

                // Memory[ProgramCounter] = 2 >> *value & 0xFF;
                // Memory[ProgramCounter + 1] = *value & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line %d: Expected port number.\n", PeekToken()->line);
                exit(1);
            }

            ProgramCounter += 2;
            break;
        case 8: //OUT
            //Memory[ProgramCounter] = 0x08;

            ProgramCounter++;

            if (PeekToken()->type == NUMBER) {
                opcode->ParameterOneType = Constant;
                opcode->ParameterOne.Constant = *(int*)PeekToken()->value;
                // int* value = PeekToken()->value;

                // Memory[ProgramCounter] = 2 >> *value & 0xFF;
                // Memory[ProgramCounter + 1] = *value & 0xFF;
            }
            else {
                fprintf(stderr, "[Error] Line %d: Expected port number.\n", PeekToken()->line);
                exit(1);
            }

            ProgramCounter += 2;
            break;
    }

    AddListItem(opcode, sizeof(IROpcode), MachineState.Opcodes);

    return;
}

int Expect(int count, ...) {
    va_list list;
    Token* token = PeekToken();

    va_start(list, count);

    for(int i = 0; i < count; i++) {
        if (va_arg(list, TokenType) == token->type) {
            va_end(list);
            AdvanceParser();
            return 1;
        }
    }

    va_end(list);

    return 0;
}

void PrintSymbols(void) {
    printf("-----SYMBOLS-----\n");
    for(int i = 0; i < MachineState.SymbolsTable->size; i++) {
        Symbol* symbol = MachineState.SymbolsTable->content[i];
        printf("[%#06X] %s\n", symbol->address, symbol->token->lexeme);
    }
    printf("-----SYMBOLS-----\n");
}

Symbol* AddSymbol(Token* token, int length, int address, int resolved) {
    if (!token) return NULL;
    if (token->type != IDENTIFIER && token->type != LABEL) return NULL;

    for(int i = 0; i < MachineState.SymbolsTable->size; i++) {
        Symbol* s = MachineState.SymbolsTable->content[i];

        if (strcmp(s->token->lexeme, token->lexeme) == 0) {
            if (s->resolved) {
                fprintf(stderr, "[Error] Line %d: Attempted to redeclare %s.\n", token->line, token->lexeme);
                exit(1);
            }
            s->resolved = resolved;
            s->address = address;
            s->length = length;
            return s;
        }
    }

    Symbol* symbol = CreateSymbol(token, length, address, resolved);

    AddListItem(symbol, sizeof(Symbol), MachineState.SymbolsTable);

    return symbol;
}

Symbol* GetSymbol(char* name) {
    if (!name) return NULL;

    for(int i = 0; i < MachineState.SymbolsTable->size; i++) {
        Symbol* s = MachineState.SymbolsTable->content[i];

        if (strcmp(s->token->lexeme, name) == 0) return s;
    }

    return NULL;
}

Symbol* CreateSymbol(Token* token, int length, int address, int resolved) {
    Symbol* symbol = calloc(1, sizeof(Symbol));

    if (!symbol) return NULL;

    symbol->token = token;
    symbol->address = address;
    symbol->resolved = resolved;

    return symbol;
}

IROpcode* CreateIROpcode(TokenType mnemonic) {
    IROpcode* opcode = calloc(1, sizeof(IROpcode));

    opcode->Mnemonic = mnemonic;

    return opcode;
}

void AdvanceParser(void) {
    if (ParserAtEnd()) return;

    CurrentToken++;
}

int ParserAtEnd(void) {
    if (CurrentToken >= TokensList->size) return 1;

    return 0;
}

Token* PeekToken(void) {
    return TokensList->content[CurrentToken];
}