#include "lexer.h"
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <string.h>

void addTokenToLine(TokenLine* line, Token tok) {
    line->count++;
    Token* newTokens = realloc(line->tokens, line->count * sizeof(Token));
    if (!newTokens) {
        perror("Failed to allocate token");
        exit(EXIT_FAILURE);
    }
    line->tokens = newTokens;
    line->tokens[line->count - 1] = tok;
}

void addLineToLexed(LexedFile* lf, TokenLine line) {
    lf->lineCount++;
    TokenLine* newLines = realloc(lf->lines, lf->lineCount * sizeof(TokenLine));
    if (!newLines) {
        perror("Failed to allocate line");
        exit(EXIT_FAILURE);
    }
    lf->lines = newLines;
    lf->lines[lf->lineCount - 1] = line;
}

LexedFile lexFile(const char* fileContents) {
    LexedFile result = {0};
    result.lines = NULL;
    result.lineCount = 0;
    
    TokenLine currentLine = {0};
    currentLine.tokens = NULL;
    currentLine.count = 0;
    
    char buf[1024] = {0};
    size_t bufLen = 0;
    bool inString = false;
    bool inChar = false;
    bool isComment = false;
    
    for (size_t i = 0; fileContents[i] != '\0'; i++) {
        char c = fileContents[i];
        
        switch (c) {
            case '"':
                if (!isComment) {
                    if (inChar) {
                        buf[bufLen++] = c;
                    } else {
                        inString = !inString;
                        buf[bufLen++] = c;
                    }
                }
                break;
                
            case '\'':
                if (!isComment) {
                    if (inString) {
                        buf[bufLen++] = c;
                    } else {
                        inChar = !inChar;
                        buf[bufLen++] = c;
                    }
                }
                break;
                
            case '\n':
                if (!inString && !inChar) {
                    // Add current token to line if exists
                    if (bufLen > 0) {
                        buf[bufLen] = '\0';
                        Token tok;
                        tok.text = strdup(buf);
                        // Add tok to currentLine (need helper function)
                        addTokenToLine(&currentLine, tok);
                        bufLen = 0;
                    }
                    // Add line to result (need helper function)
                    addLineToLexed(&result, currentLine);
                    // Reset for next line
                    currentLine.tokens = NULL;
                    currentLine.count = 0;
                    isComment = false;
                } else if (!isComment) {
                    buf[bufLen++] = c;
                }
                break;
                
            case '#':
                if (!inString && !inChar) {
                    isComment = true;
                    if (bufLen > 0) {
                        buf[bufLen] = '\0';
                        Token tok;
                        tok.text = strdup(buf);
                        addTokenToLine(&currentLine, tok);
                        bufLen = 0;
                    }
                    addLineToLexed(&result, currentLine);
                    currentLine.tokens = NULL;
                    currentLine.count = 0;
                } else {
                    buf[bufLen++] = c;
                }
                break;
                
            case ' ':
                if (!inString && !inChar) {
                    if (bufLen > 0 && !isComment) {
                        buf[bufLen] = '\0';
                        Token tok;
                        tok.text = strdup(buf);
                        addTokenToLine(&currentLine, tok);
                        bufLen = 0;
                    }
                } else {
                    buf[bufLen++] = c;
                }
                break;
                
            default:
                if (!isComment) {
                    buf[bufLen++] = c;
                }
                break;
        }
    }
    
    // Handle any remaining content
    if (bufLen > 0) {
        buf[bufLen] = '\0';
        Token tok;
        tok.text = strdup(buf);
        addTokenToLine(&currentLine, tok);
    }
    if (currentLine.count > 0) {
        addLineToLexed(&result, currentLine);
    }
    
    return result;
}