#include <stdio.h>
#include "xyw.h"

#define TOKEN_MAX_LENGTH 32
#define DIRECTIVE_CONTENT_MAX_LENGTH 256
#define STRING_MAX_LENGTH 256
#define MAX_SYMBOLS 0x500
#define MAX_REFS 0x1000

#define error(msg) fprintf(stderr, "Error - %s [%s]\n", msg, ctx->file)
#define token_error(msg) fprintf(stderr, "Error - %s: %s [%s - Ln %d, Col %d]\n", token, msg, ctx->file, ctx->line, ctx->column)
#define symbol_error(symbol, msg) fprintf(stderr, "Error - %s: %s [%s - Ln %d, Col %d]\n", symbol, msg, ctx->file, ctx->line, ctx->column)

#define PSH 0x01

typedef struct
{
    enum
    {
        SRC_FILE,
        SRC_STRING
    } src;
    int line;
    int column;
    char *file;
    FILE *fp;
    const char *sp;
} xyw_context;

typedef enum
{
    SYMBOL_TYPE_LABEL,
    SYMBOL_TYPE_MACRO,
    SYMBOL_TYPE_REF
} xyw_symbol_type;

typedef struct
{
    char *name;
    xyw_symbol_type type;
    xyw_word address;
    char *contents;
} xyw_symbol;

typedef enum
{
    TOKEN_TYPE_UNKNOWN,
    TOKEN_TYPE_COMMENT,
    TOKEN_TYPE_DIRECTIVE,
    TOKEN_TYPE_NUMBER,
    TOKEN_TYPE_INSTRUCTION,
    TOKEN_TYPE_STRING,
    TOKEN_TYPE_SYMBOL
} xyw_token_type;

typedef enum
{
    TOKEN_END_EOF = 0,
    TOKEN_END_EOL,
    TOKEN_END_SPACE,
    TOKEN_END_INVALID
} xyw_token_terminator;

// The token being processed
static char token[TOKEN_MAX_LENGTH];
static unsigned int token_length = 0;
static xyw_token_type token_type = TOKEN_TYPE_UNKNOWN;
static xyw_word token_value = 0;

// Tracks whether we are inside a directive or not
static unsigned int directive = 0;

// Bytecode to be written
xyw_byte data[0xffff] = {};

// Current write pointer
xyw_word ptr = 0;

// Flat single-string "dictionary" for symbols (macros and labels)
// Removing 768 bytes to leave room for devices and stacks
// Actually following Uxn design here ;)
char dict[0x8000 - 0x300];
char *dictnext = dict;

// List of symbols (macros and labels)
xyw_symbol symbols[MAX_SYMBOLS];
xyw_word symcount = 0;

// List of references to symbols
xyw_symbol refs[MAX_REFS];
xyw_word refcount = 0;

//// Dictionary management

// Stores a new string in the dictionary and returns its address
static char *dict_store(char *str)
{
    char *o = dictnext;
    while (*str)
    {
        *dictnext++ = *str++;
    }
    *dictnext++ = '\0';
    return o;
}

// Checks if str is already in the dictionary
// if so, returns its address, else it stores it and returns the new address
static char *dict_add(char *str)
{
    char *dict_cursor = dict;

    while (dict_cursor < dictnext)
    {
        char *existing_str = dict_cursor;
        char *input_cursor = str;
        char *existing_cursor = existing_str;

        // Compare the two strings character by character
        while (*existing_cursor && *input_cursor && (*existing_cursor == *input_cursor))
        {
            existing_cursor++;
            input_cursor++;
        }

        // If both ended at the same time, we found an exact match
        if (*existing_cursor == '\0' && *input_cursor == '\0')
        {
            return existing_str;
        }

        // Move dict_cursor to the start of the next string in the dictionary
        while (*existing_cursor != '\0')
        {
            existing_cursor++;
        }
        dict_cursor = existing_cursor + 1; // skip the null terminator
    }

    // Not found, store the new string and return its address
    return dict_store(str);
}

//// Forward declarations
static xyw_token_terminator next_token(xyw_context *ctx);
static int assemble_string(xyw_context *ctx);
static int assemble_file(xyw_context *ctx);

//// Utilities

static int find_symbol(char *name)
{
    for (xyw_word i = 0; i < symcount; i++)
    {
        for (unsigned int j = 0; j < TOKEN_MAX_LENGTH; j++)
        {
            if (symbols[i].name[j] != name[j])
            {
                break;
            }
            if (name[j] == '\0')
            {
                return i;
            }
        }
    }
    return -1;
}

static int strcmpn(const char *s1, const char *s2, unsigned int n)
{
    for (unsigned int i = 0; i < n; i++)
    {
        if (s1[i] != s2[i])
        {
            return s1[i] - s2[i];
        }
    }
    return 0;
}

void byte_to_binary(xyw_byte byte, char *str)
{
    int i;
    for (i = 7; i >= 0; i--)
    {
        str[7 - i] = ((byte >> i) & 1) ? '1' : '0';
    }
    str[8] = '\0';
}

static int validate_symbol(xyw_context *ctx, char *str, unsigned int length)
{
    // symbols must start with a letter or underscore,
    // and they can contain letters, numbers, underscores, dashes, or dots.
    if (!((str[0] >= 'A' && str[0] <= 'Z') || (str[0] >= 'a' && str[0] <= 'z') || str[0] == '_'))
    {
        symbol_error(str, "symbol must start with a letter or underscore");
        return -1;
    }
    for (unsigned int i = 1; i < length; i++)
    {
        char c = str[i];
        if (!((c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_' || c == '-' || c == '.'))
        {
            symbol_error(str, "symbol can only contain letters, numbers, underscores, dashes, or dots");
            return -1;
        }
    }
    return 0;
}

static int validate_single_directive_argument(xyw_context *ctx, xyw_token_terminator r)
{
    while (r != TOKEN_END_EOL && r != TOKEN_END_EOF)
    {
        r = next_token(ctx);
        if (r == TOKEN_END_INVALID)
        {
            return -1;
        }
        if (r == TOKEN_END_EOL || r == TOKEN_END_EOF)
        {
            break;
        }
        token_error("Too many tokens after directive");
        return -1;
    }
    return 0;
}

// Append raw bytecode data
static inline void append(xyw_word val)
{
    if (val > 0xff)
    {
        XYW_POKEW(&data[ptr], val);
        ptr++;
    }
    else
    {
        data[ptr] = (xyw_byte)val;
    }
    ptr++;
}

// Encode a push instruction
static inline void push(xyw_word val)
{
    (val > 0xff) ? append(PSH | XYW_MODE_W) : append(PSH);
    append(val);
}

// Store macro definition in dictionary
static int macro(xyw_context *ctx, char *name, char *contents)
{
    if (symcount >= MAX_SYMBOLS)
    {
        error("Too many symbols defined");
        return -1;
    }
    if (find_symbol(name) != -1)
    {
        symbol_error(name, "Symbol already defined");
        return -1;
    }
    xyw_symbol *sym = &symbols[symcount++];
    sym->type = SYMBOL_TYPE_MACRO;
    sym->name = dict_add(name);
    sym->contents = dict_store(contents);
    return 0;
}

// Store label in dictionary
static int label(xyw_context *ctx, char *name)
{
    if (symcount >= MAX_SYMBOLS)
    {
        error("Too many symbols defined");
        return -1;
    }
    if (find_symbol(name) != -1)
    {
        symbol_error(name, "Symbol already defined");
        return -1;
    }
    xyw_symbol *sym = &symbols[symcount++];
    sym->type = SYMBOL_TYPE_LABEL;
    sym->name = dict_add(name);
    sym->address = ptr;
    return 0;
}

// Store reference in dictionary
static int ref(xyw_context *ctx, char *name)
{
    if (refcount >= MAX_REFS)
    {
        error("Too many symbol references");
        return -1;
    }
    xyw_symbol *sym = &refs[refcount++];
    sym->type = SYMBOL_TYPE_REF;
    sym->name = dict_add(name);
    sym->address = ptr;
    return 0;
}

//// Processing functions

static int process_comment(xyw_context *ctx)
{
    (void)ctx;
    token_type = TOKEN_TYPE_COMMENT;
    xyw_dbg("Comment:\t%s\n", token);
    return 0;
}

static int process_number(xyw_context *ctx)
{
    xyw_word value = 0;
    unsigned int i = 0;
    if (token[0] == '$')
    {
        for (i = 1; i < token_length; i++)
        {
            char c = token[i];
            value <<= 4;
            if (c >= '0' && c <= '9')
            {
                value |= (c - '0');
            }
            else if (c >= 'A' && c <= 'F')
            {
                value |= (c - 'A' + 10);
            }
            else if (c >= 'a' && c <= 'f')
            {
                value |= (c - 'a' + 10);
            }
            else
            {
                token_error("Invalid hex digit");
                return -1;
            }
        }
        xyw_dbg("Hex number:\t%s\t(%04X)\n", token, value);
    }
    else
    {
      token_error("Invalid number.");
      return -1;
    }
    token_type = TOKEN_TYPE_NUMBER;
    token_value = value;
    if (!directive)
    {
        push(value);
    }
    return 0;
}

static int process_instruction(xyw_context *ctx)
{
    (void)ctx;
    xyw_byte opcode = 0;
    // Match token against instruction mnemonics
    for (unsigned int i = 0; i < XYW_TOTAL_INSTRUCTIONS; i++)
    {
        if (strcmpn(token, xyw_instructions[i], 3) == 0)
        {
            opcode = (xyw_byte)i;
            // There could be 0 to 3 modes (x y z) after the mnemonic in any order
            if (token_length > 3)
            {
                for (unsigned int j = 3; j < token_length; j++)
                {
                    char mode_char = token[j];
                    int mode_found = 0;
                    if (mode_char == 'x' || mode_char == 'X')
                    {
                        opcode |= XYW_MODE_X;
                        mode_found = 1;
                    }
                    else if (mode_char == 'y' || mode_char == 'Y')
                    {
                        opcode |= XYW_MODE_Y;
                        mode_found = 1;
                    }
                    else if (mode_char == 'w' || mode_char == 'W')
                    {
                        opcode |= XYW_MODE_W;
                        mode_found = 1;
                    }
                    if (!mode_found)
                    {
                        return -1;
                    }
                }
            }
            char bits[9];
            byte_to_binary(opcode, bits);
            xyw_dbg("Instruction:\t%s\t(%02X - %s)\n", token, opcode, bits);
            token_type = TOKEN_TYPE_INSTRUCTION;
            token_value = opcode;
            append(opcode);
            return 0;
        }
    }
    return -1;
}

static int process_string(xyw_context *ctx)
{
    (void)ctx;
    token_type = TOKEN_TYPE_STRING;
    return 0;
}

static int process_symbol(xyw_context *ctx)
{
    if (validate_symbol(ctx, token, token_length) != 0)
    {
        return -1;
    }
    token_type = TOKEN_TYPE_SYMBOL;
    if (!directive)
    {
        int i = find_symbol(token);
        if (i != -1)
        {
            xyw_symbol *sym = &symbols[i];
            if (sym->type == SYMBOL_TYPE_MACRO)
            {
                xyw_dbg("Macro:\t\t%s -> %s\n", token, sym->contents);
                ctx->sp = sym->contents;
                assemble_string(ctx);
                ctx->src = SRC_FILE;
                ctx->sp = NULL;
            }
            else if (sym->type == SYMBOL_TYPE_LABEL)
            {
                xyw_dbg("Label Ref:\t%s -> %04X\n", token, sym->address);
                // Push label address
                append(PSH | XYW_MODE_W);
                XYW_POKEW(&data[ptr], sym->address);
                ptr += 2;
            }
            else if (sym->type == SYMBOL_TYPE_REF)
            {
                // Should not happen
                token_error("Internal error: symbol is a reference");
                return -1;
            }
        }
        else
        {
            // Store reference to be resolved later
            xyw_dbg("Reference:\t%s\n", token);
            // Push placeholder value $0000 for now
            append(PSH | XYW_MODE_W);
            if (ref(ctx, token) != 0)
            {
                return -1;
            }
            append(0x00);
            append(0x00);
        }
    }
    return 0;
}

//// Directive processing

static int process_directive_label(xyw_context *ctx)
{
    xyw_token_terminator r = next_token(ctx);
    if (r == TOKEN_END_INVALID)
    {
        return -1;
    }
    if (process_symbol(ctx) != 0)
    {
        return -1;
    }
    xyw_dbg("Directive:\t.label %s [$%04X]\n", token, ptr);
    label(ctx, token);
    return validate_single_directive_argument(ctx, r);
}

static int process_directive_include(xyw_context *ctx)
{
    xyw_token_terminator r = next_token(ctx);
    if (r == TOKEN_END_INVALID)
    {
        return -1;
    }
    process_string(ctx);
    xyw_dbg("Directive:\t.include %s\n", token);
    xyw_context include_ctx;
    include_ctx.src = SRC_FILE;
    // Ignore token quotes for file name
    token[token_length - 1] = '\0';
    include_ctx.file = dict_store(token + 1);
    include_ctx.line = 1;
    include_ctx.column = 1;
    if (validate_single_directive_argument(ctx, r) != 0)
    {
        return -1;
    }
    xyw_dbg("-> Accessing:\t%s\n", include_ctx.file);
    int result = assemble_file(&include_ctx);
    xyw_dbg("<- Continuing:\t%s\n", ctx->file);
    return result;
}

static int process_directive_string(xyw_context *ctx)
{
    xyw_token_terminator r = next_token(ctx);
    if (r == TOKEN_END_INVALID)
    {
        return -1;
    }
    process_string(ctx);
    xyw_dbg("Directive:\t.string %s\n", token);
    // Store string in memory (without quotes)
    for (unsigned int i = 1; i < token_length - 1; i++)
    {
        data[ptr++] = (xyw_byte)token[i];
    }
    data[ptr++] = 0;
    return validate_single_directive_argument(ctx, r);
}

static int process_directive_byte(xyw_context *ctx)
{
    xyw_token_terminator r;
    xyw_dbg("Directive:\t.byte\n");
    while (1)
    {
        r = next_token(ctx);
        if (r == TOKEN_END_INVALID)
        {
            return -1;
        }
        if (process_number(ctx) != 0)
        {
            return -1;
        }
        if (token_value > 0xFF)
        {
            token_error("Byte value out of range (0-255)");
            return -1;
        }
        if (r == TOKEN_END_EOL || r == TOKEN_END_EOF)
        {
            break;
        }
        data[ptr++] = (xyw_byte)token_value;
    }
    return 0;
}

static int process_directive_macro(xyw_context *ctx)
{
    xyw_token_terminator r = next_token(ctx);
    if (r == TOKEN_END_INVALID)
    {
        return -1;
    }
    if (process_symbol(ctx) != 0)
    {
        return -1;
    }
    xyw_dbg("Directive:\t.macro %s\n", token);
    char c;
    char contents[DIRECTIVE_CONTENT_MAX_LENGTH];
    int length = 0;
    while (1)
    {
        if (ctx->src == SRC_FILE)
        {
            if (fread(&c, 1, 1, ctx->fp) != 1)
            {
                break;
            }
        }
        else
        {
            if (!ctx->sp || *ctx->sp == '\0')
            {
                break;
            }
            c = *ctx->sp++;
        }
        if (c == '\n')
        {
            ctx->line++;
            ctx->column = 1;
            break;
        }
        if (c == ';')
        {
            break;
        }
        if (length > DIRECTIVE_CONTENT_MAX_LENGTH)
        {

            token_error("Macro contents too long");
            return -1;
        }
        ctx->column++;
        contents[length++] = c;
        contents[length] = '\0';
    }
    macro(ctx, token, contents);
    return 0;
}

static int process_directive(xyw_context *ctx)
{
    if (validate_symbol(ctx, token + 1, token_length - 1) != 0)
    {
        return -1;
    }
    directive = 1;
    token_type = TOKEN_TYPE_DIRECTIVE;
    int r;
    if (strcmpn(token, ".label", 6) == 0)
    {
        r = process_directive_label(ctx);
    }
    else if (strcmpn(token, ".include", 8) == 0)
    {
        r = process_directive_include(ctx);
    }
    else if (strcmpn(token, ".string", 7) == 0)
    {
        r = process_directive_string(ctx);
    }
    else if (strcmpn(token, ".byte", 5) == 0)
    {
        r = process_directive_byte(ctx);
    }
    else if (strcmpn(token, ".macro", 6) == 0)
    {
        r = process_directive_macro(ctx);
    }
    else
    {
        token_error("Unknown directive");
        r = -1;
    }
    directive = 0;
    return r;
}

// Process and validate a token
static int process_token(xyw_context *ctx)
{
    if (token_length == 0)
    {
        return 0;
    }
    if (token[0] == ';')
    {
        return process_comment(ctx);
    }
    else if (token[0] == '.')
    {
        return process_directive(ctx);
    }
    else if (token[0] == '$' || token[0] == '"')
    {
        if (token[0] == '"')
        {
            return process_string(ctx);
        }
        return process_number(ctx);
    }
    else if ((token[0] >= '0' && token[0] <= '9') || token[0] == '-')
    {
        return process_number(ctx);
    }
    else if ((token[0] >= 'A' && token[0] <= 'Z') || (token[0] >= 'a' && token[0] <= 'z') || token[0] == '_')
    {
        if (process_instruction(ctx) == 0)
        {
            return 0;
        }
        return process_symbol(ctx);
    }
    token_error("Unrecognized token");
    return -1;
}

// Get the next token from the input file (without processing it)
static xyw_token_terminator next_token(xyw_context *ctx)
{
    char c;
    int in_string = 0;
    int in_comment = 0;
    token_length = 0;
    token[0] = '\0';
    while (1)
    {
        // Read next character from active source
        if (ctx->src == SRC_FILE)
        {
            if (fread(&c, 1, 1, ctx->fp) != 1)
            {
                break;
            }
        }
        else
        {
            if (!ctx->sp || *ctx->sp == '\0')
            {
                if (token_length > 0)
                {
                    // Force token finalization on string end
                    c = '\n';
                }
                else
                {
                    break;
                }
            }
            else
            {
                c = *ctx->sp++;
            }
        }
        // printf("Char read: '%c' (0x%02X) [Ln %d, Col %d] - Token: %s\n", c, (unsigned char)c, ctx->line, ctx->column, token);
        //  Handle character
        if (c == '\n')
        {
            if (in_string)
            {
                token_error("Unterminated string literal");
                return TOKEN_END_INVALID;
            }
            if (in_comment)
            {
                in_comment = 0;
            }
            ctx->line++;
            ctx->column = 1;
            return TOKEN_END_EOL;
        }
        else if (in_comment)
        {
            ctx->column++;
            // Not storing comment context in token for now
        }
        else if (c == '"')
        {
            ctx->column++;
            if (token_length > STRING_MAX_LENGTH)
            {
                token_error("Token too long");
                return TOKEN_END_INVALID;
            }
            token[token_length++] = (char)c;
            token[token_length] = '\0';
            if (in_string)
            {
                in_string = 0;
            }
            else
            {
                in_string = 1;
            }
        }
        else if (in_string)
        {
            if (c == '\\')
            {
                // Handle escape sequences
                char next_char;
                if (ctx->src == SRC_FILE)
                {
                    if (fread(&next_char, 1, 1, ctx->fp) != 1)
                    {
                        token_error("Unterminated escape sequence in string literal");
                        return TOKEN_END_INVALID;
                    }
                }
                else
                {
                    if (!ctx->sp || *ctx->sp == '\0')
                    {
                        token_error("Unterminated escape sequence in string literal");
                        return TOKEN_END_INVALID;
                    }
                    next_char = *ctx->sp++;
                }
                ctx->column++;
                switch (next_char)
                {
                case 'n':
                    c = '\n';
                    break;
                case 't':
                    c = '\t';
                    break;
                case '"':
                    c = '"';
                    break;
                case 'r':
                    c = '\r';
                    break;
                case '\'':
                    c = '\'';
                    break;
                case '\\':
                    c = '\\';
                    break;
                default:
                    token_error("Unknown escape sequence in string literal");
                    return TOKEN_END_INVALID;
                }
            }
            ctx->column++;
            if (token_length > STRING_MAX_LENGTH)
            {
                token_error("Token too long");
                return TOKEN_END_INVALID;
            }
            token[token_length++] = (char)c;
            token[token_length] = '\0';
        }
        else if (c == ';' && !in_string && !in_comment)
        {
            in_comment = 1;
            ctx->column++;
            // Start of comment token
            token[token_length++] = (char)c;
            token[token_length] = '\0';
        }
        // Whitespace characters
        else if (c == ' ' || c == '\t')
        {
            ctx->column++;
            if (token_length > 0)
            {
                return TOKEN_END_SPACE;
            }
        }
        // Non-whitespace, non-control characters
        else if (c > 0x20 && c != 0x7f)
        {
            ctx->column++;
            if (token_length > TOKEN_MAX_LENGTH)
            {
                token_error("Token too long");
                return TOKEN_END_INVALID;
            }
            token[token_length++] = (char)c;
            token[token_length] = '\0';
        }
        else
        {
            token_error("Invalid character encountered");
            return TOKEN_END_INVALID;
        }
    }
    return TOKEN_END_EOF;
}

// Read a file and process its tokens
static int assemble_file(xyw_context *ctx)
{
    ctx->fp = fopen(ctx->file, "r");
    if (!ctx->fp)
    {
        error("Could not open file");
        return -1;
    }
    // Set source mode to file
    ctx->src = SRC_FILE;
    ctx->sp = NULL;
    int r;
    while (1)
    {
        r = next_token(ctx);
        if (r == TOKEN_END_EOF || r == TOKEN_END_INVALID)
        {
            break;
        }
        if (process_token(ctx) != 0)
        {
            break;
        }
    }
    fclose(ctx->fp);
    return r;
}

// Process tokens from an in-memory string buffer (mainly for macro body)
static int assemble_string(xyw_context *ctx)
{
    ctx->src = SRC_STRING;
    int r;
    while (1)
    {
        r = next_token(ctx);
        if (r == TOKEN_END_EOF || r == TOKEN_END_INVALID)
        {
            break;
        }
        if (process_token(ctx) != 0)
        {
            break;
        }
    }
    return r;
}

///// Public API

int xyw_assemble(const char *input, const char *output)
{
    xyw_context ctx;
    ctx.file = (char *)input;
    ctx.line = 1;
    ctx.column = 1;
    if (assemble_file(&ctx) != 0)
    {
        return -1;
    }
    FILE *fp = fopen(output, "wb");
    if (!fp)
    {
        fprintf(stderr, "Could not open output file for writing");
        return -1;
    }
    // Patch labels
    for (xyw_word i = 0; i < refcount; i++)
    {
        xyw_symbol *refsym = &refs[i];
        int symindex = find_symbol(refsym->name);
        if (symindex == -1)
        {
            fprintf(stderr, "Undefined symbol: %s\n", refsym->name);
            fclose(fp);
            return -1;
        }
        xyw_symbol *defsym = &symbols[symindex];
        if (defsym->type != SYMBOL_TYPE_LABEL)
        {
            fprintf(stderr, "Invalid reference to non-label symbol: %s\n", refsym->name);
            fclose(fp);
            return -1;
        }
        // Patch the data at the reference address with the label address
        xyw_dbg("Patch:\t\t%s -> $%04X @ $%04X\n", defsym->name, defsym->address, refsym->address);
        XYW_POKEW(&data[refsym->address], defsym->address);
    }
    fwrite(data, 1, ptr, fp);
    fclose(fp);
    printf("File '%s' assembled successfully [%d bytes - %d symbols - %d refs]\n", input, ptr, symcount, refcount);
    return 0;
}