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

#define TOKEN_MAX_LENGTH 32
#define DIRECTIVE_CONTENT_MAX_LENGTH 256
#define COMMENT_CONTENT_MAX_LENGTH 256
#define STRING_MAX_LENGTH 256

#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
{
    int line;
    int column;
    char *file;
    FILE *fp;
} xyw_context;

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;

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

// Directive identifier
static char directive[TOKEN_MAX_LENGTH];
static unsigned int directive_length = 0;

// Directive data
static char directive_data[DIRECTIVE_CONTENT_MAX_LENGTH];
static unsigned int directive_data_length = 0;

// Bytecode to be written
xyw_byte data[0xffff - 0x0100];

// Current write pointer
xyw_short ptr = 0x0100;

//// Utilities

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 void reset_directive()
{
    directive_length = 0;
    directive[0] = '\0';
    directive_data_length = 0;
    directive_data[0] = '\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);
}

//// Processing functions

static int process_comment(xyw_context *ctx)
{
    char c;
    char comment_data[COMMENT_CONTENT_MAX_LENGTH];
    unsigned int comment_length = 0;
    reset_directive();
    // Process single-line comment
    while (fread(&c, 1, 1, ctx->fp) == 1)
    {
        if (c == '\n')
        {
            ctx->line++;
            ctx->column = 0;
            break;
        }
        ctx->column++;
        if (comment_length < COMMENT_CONTENT_MAX_LENGTH - 1)
        {
            comment_data[comment_length++] = c;
            comment_data[comment_length] = '\0';
        }
        else
        {
            token_error("Comment too long");
            return -1;
        }
    }
    // TODO: Remove
    printf("Comment:\t%s%s\n", token, comment_data);
    token_type = TOKEN_TYPE_COMMENT;
    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) [directive: %s]\n", token, value, directive_length > 0 ? directive : "N/A");
    token_type = TOKEN_TYPE_NUMBER;
    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) [directive: %s]\n", token, value, directive_length > 0 ? directive : "N/A");
    token_type = TOKEN_TYPE_NUMBER;
    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) [directive: %s]\n", token, value, directive_length > 0 ? directive : "N/A");
    token_type = TOKEN_TYPE_NUMBER;
    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 < 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) [directive: %s]\n", token, opcode, bits, directive_length > 0 ? directive : "N/A");
            write(opcode);
            token_type = TOKEN_TYPE_INSTRUCTION;
            return 0;
        }
    }
    return -1;
}

static int process_string(xyw_context *ctx)
{
    (void)ctx;
    // Token is a string enclosed in double quotes
    // TODO: Remove
    printf("String:\t\t%s [directive: %s]\n", token, directive_length > 0 ? directive : "N/A");
    // TODO: Store string data
    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 [directive: %s]\n", token, directive_length > 0 ? directive : "N/A");
    token_type = TOKEN_TYPE_SYMBOL;
    return 0;
}

// Returns: 1 token processed, 0 EOF, -1 error
static int next_token(xyw_context *ctx);

static int process_directive_label(xyw_context *ctx)
{
    int r = next_token(ctx);
    if (r != 1 || token_type != TOKEN_TYPE_SYMBOL)
    {
        token_error("Expected symbol after .label directive");
        return -1;
    }
    // TODO: store label info
    return 0;
}

static int process_directive(xyw_context *ctx)
{
    if (validate_symbol(ctx, token + 1, token_length - 1) != 0)
    {
        return -1;
    }
    directive_length = token_length;
    for (unsigned int i = 0; i < directive_length; i++)
    {
        directive[i] = token[i];
    }
    directive[directive_length] = '\0';
    // TODO: Remove
    printf("Directive:\t%s\n", token);
    token_type = TOKEN_TYPE_DIRECTIVE;
    // TODO: Process directive
    if (strcmpn(directive, ".label", 6) == 0)
    {
        return process_directive_label(ctx);
    }
    return 0;
}

// Process and validate tokens
static int process_token(xyw_context *ctx)
{
    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;
}

static int next_token(xyw_context *ctx)
{
    unsigned char c;
    token_length = 0;
    token[0] = '\0';
    int in_string = 0;

    while (fread(&c, 1, 1, ctx->fp) == 1)
    {
        if (!in_string && c == '\n')
        {
            if (token_length > 0)
            {
                if (process_token(ctx) != 0)
                    return -1;
                token_length = 0;
                token[0] = '\0';
                ctx->line++;
                ctx->column = 1;
                reset_directive();
                return 1;
            }
            ctx->line++;
            ctx->column = 1;
            reset_directive();
            continue;
        }
        if (!in_string && token_length == 0 && c == ';')
        {
            token[token_length++] = (char)c;
            token[token_length] = '\0';
            ctx->column++;
            if (process_token(ctx) != 0)
                return -1;
            token_length = 0;
            token[0] = '\0';
            return 1;
        }
        if (c == '"')
        {
            ctx->column++;
            if (!in_string)
            {
                in_string = 1;
                if (token_length < STRING_MAX_LENGTH - 1)
                {
                    token[token_length++] = (char)c;
                    token[token_length] = '\0';
                }
                else
                {
                    token_error("Token too long");
                    return -1;
                }
                continue;
            }
            else
            {
                if (token_length < STRING_MAX_LENGTH - 1)
                {
                    token[token_length++] = (char)c;
                    token[token_length] = '\0';
                }
                else
                {
                    token_error("Token too long");
                    return -1;
                }
                if (process_token(ctx) != 0)
                    return -1;
                token_length = 0;
                token[0] = '\0';
                in_string = 0;
                return 1;
            }
        }
        else if (in_string)
        {
            ctx->column++;
            if (token_length < STRING_MAX_LENGTH - 1)
            {
                token[token_length++] = (char)c;
                token[token_length] = '\0';
            }
            else
            {
                token_error("Token too long");
                return -1;
            }
            continue;
        }
        else if (c == ' ' || c == '\t')
        {
            if (token_length > 0)
            {
                if (process_token(ctx) != 0)
                    return -1;
                token_length = 0;
                token[0] = '\0';
                ctx->column++;
                return 1;
            }
            ctx->column++;
            continue;
        }
        else if (c > 0x20 && c != 0x7f)
        {
            ctx->column++;
            if (token_length < TOKEN_MAX_LENGTH - 1)
            {
                token[token_length++] = (char)c;
                token[token_length] = '\0';
                if (directive_length > 0)
                {
                    if (directive_data_length >= DIRECTIVE_CONTENT_MAX_LENGTH - 1)
                    {
                        token_error("Directive data too long");
                        return -1;
                    }
                    directive_data[directive_data_length++] = (char)c;
                    directive_data[directive_data_length] = '\0';
                }
            }
            else
            {
                token_error("Token too long");
                return -1;
            }
        }
        else
        {
            token_error("Invalid character encountered");
            return -1;
        }
    }
    if (token_length > 0)
    {
        if (process_token(ctx) != 0)
            return -1;
        token_length = 0;
        token[0] = '\0';
        return 1; // last token at EOF
    }
    return 0; // EOF no token
}

// Read a file and process its tokens
static int process_file(xyw_context *ctx)
{
    ctx->fp = fopen(ctx->file, "r");
    if (!ctx->fp)
    {
        error("Could not open file");
        return -1;
    }

    int r;
    while ((r = next_token(ctx)) == 1)
    {
        // tokens processed inline
    }
    fclose(ctx->fp);
    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 (process_file(&ctx) != 0)
    {
        return -1;
    }

    (void)output; // TODO: write to output file
    return 0;
}