#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 MODE_X 0x20
#define MODE_Y 0x40
#define MODE_W 0x80

#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_short 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;

// Instruction mnemonics
static char instructions[][4] = {
    "BRK", "PSH", "POP", "NOT",
    "LDR", "STR", "LDD", "STD",
    "ILD", "LDI", "DLD", "LDD",
    "ADD", "SUB", "MUL", "DIV",
    "EQU", "NEQ", "GTH", "LTH",
    "AND", "IOR", "XOR", "SFT",
    "SWP", "DUP", "OVR", "ROT",
    "JMP", "JCN", "JSR", "RTS"};

// 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_short 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_short ptr = 0x0000;

// Dictionary for symbols (macros and labels)
// Actually following Uxn design here ;)
char dict[0x8000];
char *dictnext = dict;

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

// List of references to symbols
xyw_symbol refs[MAX_REFS];
xyw_short 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, const char *text);

//// Utilities

static int find_symbol(char *name)
{
    for (xyw_short 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;
}

// Write raw bytecode data
static inline void write(xyw_short val)
{
    if (val > 0xff)
    {
        data[ptr++] = val & 0xFF;
        data[ptr++] = (val >> 8) & 0xFF;
        return;
    }
    data[ptr++] = (xyw_byte)val;
}

// Encode a push instruction
static inline void push(xyw_short val)
{
    (val > 0xff) ? write(PSH | MODE_W) : write(PSH);
    write(val);
}

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

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

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)
{
    char c;
    // Process single-line comment
    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++;
        }
        // Ignore all characters until end of line
        if (c == '\n')
        {
            ctx->line++;
            ctx->column = 0;
            break;
        }
        ctx->column++;
    }
    token_type = TOKEN_TYPE_COMMENT;
    printf("Comment:\t%s\n", token);
    return 0;
}

static int process_hex_number(xyw_context *ctx)
{
    // Token is a hex number starting with $
    xyw_short value = 0;
    for (unsigned int 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;
        }
    }
    // TODO: Remove
    printf("Hex number:\t%s\t(%04X)\n", token, value);
    token_type = TOKEN_TYPE_NUMBER;
    token_value = value;
    if (!directive)
    {
        push(value);
    }
    return 0;
}

static int process_binary_number(xyw_context *ctx)
{
    // Token is a binary number starting with %
    xyw_short value = 0;
    for (unsigned int i = 1; i < token_length; i++)
    {
        char c = token[i];
        value <<= 1;
        if (c == '0')
        {
            // do nothing
        }
        else if (c == '1')
        {
            value |= 1;
        }
        else
        {
            token_error("Invalid binary digit");
            return -1;
        }
    }
    // TODO: Remove
    printf("Binary number:\t%s\t(%04X)\n", token, value);
    token_type = TOKEN_TYPE_NUMBER;
    token_value = value;
    if (!directive)
    {
        push(value);
    }
    return 0;
}

static int process_decimal_number(xyw_context *ctx)
{
    // Token is a decimal number (may or may not start with -)
    xyw_short value = 0;
    unsigned int start_index = 0;
    int is_negative = 0;
    if (token[0] == '-')
    {
        is_negative = 1;
        start_index = 1;
    }
    for (unsigned int i = start_index; i < token_length; i++)
    {
        char c = token[i];
        if (c >= '0' && c <= '9')
        {
            value = value * 10 + (c - '0');
        }
        else
        {
            token_error("Invalid decimal digit");
            return -1;
        }
    }
    if (is_negative)
    {
        value = -value;
    }
    // TODO: Remove
    printf("Decimal number:\t%s\t(%04X)\n", token, value);
    token_type = TOKEN_TYPE_NUMBER;
    token_value = value;
    if (!directive)
    {
        push(value);
    }
    return 0;
}

static int process_number(xyw_context *ctx)
{
    if (token[0] == '$')
    {
        return process_hex_number(ctx);
    }
    else if (token[0] == '%')
    {
        return process_binary_number(ctx);
    }
    else
    {
        return process_decimal_number(ctx);
    }
}

static int process_instruction(xyw_context *ctx)
{
    (void)ctx;
    xyw_byte opcode = 0;
    // Match token against instruction mnemonics
    for (unsigned int i = 0; i < sizeof(instructions) / sizeof(instructions[0]); i++)
    {
        if (strcmpn(token, 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 |= MODE_X;
                        mode_found = 1;
                    }
                    else if (mode_char == 'y' || mode_char == 'Y')
                    {
                        opcode |= MODE_Y;
                        mode_found = 1;
                    }
                    else if (mode_char == 'w' || mode_char == 'W')
                    {
                        opcode |= MODE_W;
                        mode_found = 1;
                    }
                    if (!mode_found)
                    {
                        return -1;
                    }
                }
            }
            // TODO: Remove
            char bits[9];
            byte_to_binary(opcode, bits);
            printf("Instruction:\t%s\t(%02X - %s)\n", token, opcode, bits);
            write(opcode);
            token_type = TOKEN_TYPE_INSTRUCTION;
            token_value = opcode;
            if (!directive)
            {
                push(opcode);
            }
            return 0;
        }
    }
    return -1;
}

static int process_string(xyw_context *ctx)
{
    (void)ctx;
    // TODO: Remove
    printf("String:\t\t%s\n", token);
    token_type = TOKEN_TYPE_STRING;
    return 0;
}

static int process_symbol(xyw_context *ctx)
{
    if (validate_symbol(ctx, token, token_length) != 0)
    {
        return -1;
    }
    // TODO: Remove
    printf("Symbol:\t\t%s\n", token);
    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)
            {
                printf("-> Macro:\t%s\n", sym->contents);
                assemble_string(ctx, sym->contents);
                ctx->src = SRC_FILE;
                ctx->sp = NULL;
            }
            else if (sym->type == SYMBOL_TYPE_LABEL)
            {
                // Push label address
                push(sym->address);
            }
            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
            printf("-> Reference:\t%s\n", token);
            if (ref(ctx, token) != 0)
            {
                return -1;
            }
        }
    }
    return 0;
}

//// Directive processing

static int process_directive_org(xyw_context *ctx)
{
    xyw_token_terminator r = next_token(ctx);
    if (r == TOKEN_END_INVALID)
    {
        return -1;
    }
    if (process_number(ctx) != 0)
    {
        return -1;
    }
    printf("Directive:\t.org %s\n", token);
    // Move the pointer to the specified address
    ptr = (xyw_short)token_value;
    return validate_single_directive_argument(ctx, r);
}

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;
    }
    printf("Directive:\t.label %s\n", token);
    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);
    printf("Directive:\t.include %s\n", token);
    // TODO: process include
    return validate_single_directive_argument(ctx, r);
}

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);
    printf("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;
    printf("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;
    }
    printf("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, ".org", 4) == 0)
    {
        r = process_directive_org(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] == '$')
    {
        return process_hex_number(ctx);
    }
    else if (token[0] == '%')
    {
        return process_binary_number(ctx);
    }
    else if (token[0] == '\"')
    {
        return process_string(ctx);
    }
    else if ((token[0] >= '0' && token[0] <= '9') || token[0] == '-')
    {
        return process_decimal_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;
    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')
            {
                break;
            }
            c = *ctx->sp++;
        }
        if (c == '\n')
        {
            if (in_string)
            {
                token_error("Unterminated string literal");
                return TOKEN_END_INVALID;
            }
            ctx->line++;
            ctx->column = 1;
            return TOKEN_END_EOL;
        }
        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)
        {
            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';
        }
        // 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, const char *text)
{
    ctx->src = SRC_STRING;
    ctx->sp = text;
    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;
    }
    fwrite(data, 1, ptr, fp);
    fclose(fp);
    printf("Assembled %s successfully (%d bytes - %d symbols - %d refs).\n", input, ptr, symcount, refcount);
    return 0;
}