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

#include "devices/console.h"

#define MEMORY_SIZE 0xffff
#define STACK_SIZE 0xff

#define USER_MEMORY_START 0x000
#define SYSTEM_MEMORY_START 0xfd00
#define USER_STACK_START 0xfd00
#define SYSTEM_STACK_START 0xfe00
#define DEVICE_AREA_START 0xff00

/* system device */
#define SYSTEM_STATE 0xff00
#define SYSTEM_ERROR 0xff01
#define SYSTEM_S 0xff02
#define SYSTEM_U 0xff03
#define SYSTEM_PAGE 0xff04
#define SYSTEM_ON_ERROR 0xff05

typedef enum
{
    XYW_ERROR_NONE = 0,
    XYW_ERROR_INVALID_INSTRUCTION,
    XYW_ERROR_STACK_UNDERFLOW,
    XYW_ERROR_STACK_OVERFLOW,
    XYW_ERROR_DEVICE_UNAVAILABLE,
    XYW_ERROR_INVALID_DEVICE_PORT
} xyw_error;

/// Main memory and stacks
xyw_byte xyw_memory[MEMORY_SIZE];
xyw_byte system_stack[STACK_SIZE];
xyw_byte user_stack[STACK_SIZE];

#define DIRECT_PAGE (xyw_memory[SYSTEM_PAGE] << 8)

//// 1-byte pointers and registers
xyw_byte *s = &xyw_memory[SYSTEM_S];
xyw_byte *u = &xyw_memory[SYSTEM_U];
static xyw_byte x_val = 0; // X register storage
static xyw_byte y_val = 0; // Y register storage
xyw_byte *x = &x_val;
xyw_byte *y = &y_val;
xyw_byte *error = &xyw_memory[SYSTEM_ERROR];

//// 2-byte registers
static xyw_word pc_storage = 0; // Program counter storage
static xyw_word xw_storage = 0; // 16-bit X register storage
static xyw_word yw_storage = 0; // 16-bit Y register storage
xyw_word *pc = &pc_storage;
xyw_word *xw = &xw_storage;
xyw_word *yw = &yw_storage;

//// Helpers to retrieve device number and address within the device
static inline int get_device(xyw_word addr)
{
    return (addr - DEVICE_AREA_START) / 16;
}

static inline xyw_byte *get_device_data(xyw_word addr)
{
    return xyw_devices[get_device(addr)].data;
}

static inline xyw_byte get_device_address(xyw_word addr)
{
    // Should return a number between 0 and 15, add validation
    xyw_byte port = (addr - DEVICE_AREA_START) % 16;
    if (port >= 16)
    {
        *error = XYW_ERROR_DEVICE_UNAVAILABLE;
        return -1;
    }
    return port;
}

//// Read/write memory helpers

static inline xyw_word readw(xyw_word addr)
{
    int dev_num = get_device(addr);
    if (dev_num >= 0 && xyw_devices[dev_num].input)
    {
        xyw_byte dev_addr = get_device_address(addr);
        // Read two bytes from device
        xyw_byte high = xyw_devices[dev_num].input(get_device_data(addr), dev_addr, error);
        xyw_byte low = xyw_devices[dev_num].input(get_device_data(addr), dev_addr + 1, error);
        return (xyw_word)((high << 8) | low);
    }
    return XYW_PEEKW(&xyw_memory[addr]);
}

static inline void writew(xyw_word addr, xyw_word val)
{
    XYW_POKEW(&xyw_memory[addr], val);
    int dev_num = get_device(addr);
    if (dev_num >= 0 && xyw_devices[dev_num].output)
    {
        xyw_byte dev_addr = get_device_address(addr);
        xyw_devices[dev_num].output(get_device_data(addr), dev_addr, error);
    }
}

static inline xyw_byte read(xyw_word addr)
{
    int dev_num = get_device(addr);
    if (dev_num >= 0 && xyw_devices[dev_num].input)
    {
        xyw_byte dev_addr = get_device_address(addr);
        return xyw_devices[dev_num].input(get_device_data(addr), dev_addr, error);
    }
    return xyw_memory[addr];
}

static inline void write(xyw_word addr, xyw_byte val)
{
    xyw_memory[addr] = val;
    int dev_num = get_device(addr);
    if (dev_num >= 0 && xyw_devices[dev_num].output)
    {
        xyw_byte dev_addr = get_device_address(addr);
        xyw_devices[dev_num].output(get_device_data(addr), dev_addr, error);
    }
}

//// Device dispatchers

//// System stack management

static xyw_word ss_popw()
{
    if (*s < 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return -1;
    }
    xyw_word val = readw(SYSTEM_STACK_START + (*s) - 2);
    writew(SYSTEM_STACK_START + (*s) - 2, 0);
    (*s) -= 2;
    return val;
}

static void ss_pushw(xyw_word val)
{
    if (*s > STACK_SIZE - 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_OVERFLOW;
        return;
    }
    writew(SYSTEM_STACK_START + *s, val);
    (*s) += 2;
}

//// User stack management

static void us_push(xyw_byte val)
{
    if (*u > STACK_SIZE - 1)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_OVERFLOW;
        return;
    }
    xyw_dbg(" => %02X\n", val);
    xyw_memory[USER_STACK_START + (*u)++] = val;
}

static void us_pushw(xyw_word val)
{
    if (*u > STACK_SIZE - 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_OVERFLOW;
        return;
    }
    writew(USER_STACK_START + *u, val);
    xyw_dbg(" => %04X\n", val);
    (*u) += 2;
}

static xyw_byte us_pop()
{
    if (*u < 1)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return -1;
    }
    (*u) -= 1;
    xyw_byte val = xyw_memory[USER_STACK_START + (*u)];
    xyw_memory[USER_STACK_START + (*u)] = 0;
    return val;
}

static xyw_word us_popw()
{
    if (*u < 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return -1;
    }
    xyw_word val = readw(USER_STACK_START + (*u) - 2);
    writew(USER_STACK_START + (*u) - 2, 0);
    (*u) -= 2;
    return val;
}

// clang-format off

//// Operations that support stack and all registers
#define OPC_SR2(opc, body) { \
    case opc: { const xyw_byte w = 0, a = us_pop(), b = us_pop(); body; } break; \
    case XYW_MODE_X|opc: { const xyw_byte w = 0, a = *x, b = us_pop(); body; } break; \
    case XYW_MODE_Y|opc: { const xyw_byte w = 0, a = *y, b = us_pop(); body; } break; \
    case XYW_MODE_X|XYW_MODE_Y|opc: { const xyw_byte w = 0, a = *x, b = *y; body; } break; \
    case XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = us_popw(); xyw_word b = us_popw(); body; } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = *xw; xyw_word b = us_popw(); body; } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = *yw; xyw_word b = us_popw(); body; } break; \
    case XYW_MODE_X|XYW_MODE_Y|XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = *xw; xyw_word b = *yw; body; } break; \
}

//// Operations that support stack and one register
#define OPC_SR1(opc, init, body) { \
    case opc: { const int w=0, rx=0, ry=0; init; body; } break; \
    case XYW_MODE_X|opc: { const int w=0, rx=1, ry=0; init; body; } break; \
    case XYW_MODE_Y|opc: { const int w=0, rx=0, ry=1; init; body; } break; \
    case XYW_MODE_W|opc: { const int w=1, rx=0, ry=0; init; body; } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { const int w=1, rx=1, ry=0; init; body; } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { const int w=1, rx=0, ry=1; init; body; } break; \
}

//// Operations that support one register only and no stack
#define OPC_R1(opc, body) { \
    case XYW_MODE_X|opc: { const xyw_byte w = 0, a = *x; body; } break; \
    case XYW_MODE_Y|opc: { const xyw_byte w = 0, a = *y; body; } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = *xw; body; } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { const xyw_byte w = 1; xyw_word a = *yw; body; } break; \
}

//// Operations that support stack only
#define OPC_S(opc, body) { \
    case opc: { const xyw_byte w = 0; body; } break; \
    case XYW_MODE_W|opc: { const xyw_byte w = 1; body; } break; \
}

//// POP operations 
#define OPC_POP(opc) { \
    case opc: { us_pop(); } break; \
    case XYW_MODE_X|opc: { *x = us_pop(); } break; \
    case XYW_MODE_Y|opc: { *y = us_pop(); } break; \
    case XYW_MODE_W|opc: { us_pop(); } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { *xw = us_popw(); } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { *yw = us_popw(); } break; \
}

//// PSH operations
#define OPC_PSH(opc) { \
    case opc: { (*pc)++; us_push(xyw_memory[*pc]); } break; \
    case XYW_MODE_X|opc: { us_push(*x); } break; \
    case XYW_MODE_Y|opc: { us_push(*y); } break; \
    case XYW_MODE_W|opc: { (*pc)++; us_pushw(readw(*pc)); (*pc)++; } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { us_pushw(*xw); } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { us_pushw(*yw); } break; \
}
//// Increment/decrement operations
#define OPC_IDC(opc, delta) { \
    case opc: { us_push(us_pop() + delta);} break; \
    case XYW_MODE_X|opc: { *x = *x + delta; } break; \
    case XYW_MODE_Y|opc: { *y = *y + delta; } break; \
    case XYW_MODE_W|opc: { us_pushw(us_popw() + delta); } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { *xw = *xw + delta; } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { *yw = *yw + delta; } break; \
}

//// SWP operations
#define OPC_SWP(opc) { \
    case opc: { xyw_byte a = us_pop(); xyw_byte b = us_pop(); us_push(a); us_push(b); } break; \
    case XYW_MODE_W|opc: { xyw_word a = us_popw(); xyw_word b = us_popw(); us_pushw(a); us_pushw(b); } break; \
    case XYW_MODE_X|opc: { xyw_byte a = us_pop(); xyw_byte b = *x; *x = a; us_push(b); } break; \
    case XYW_MODE_Y|opc: { xyw_byte a = us_pop(); xyw_byte b = *y; *y = a; us_push(b); } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { xyw_word a = us_popw(); xyw_word b = *xw; *xw = a; us_pushw(b); } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { xyw_word a = us_popw(); xyw_word b = *yw; *yw = a; us_pushw(b); } break; \
    case XYW_MODE_X|XYW_MODE_Y|opc: { xyw_byte a = *x; xyw_byte b = *y; *x = b; *y = a; } break; \
    case XYW_MODE_X|XYW_MODE_Y|XYW_MODE_W|opc: { xyw_word a = *xw; xyw_word b = *yw; *xw = b; *yw = a; } break; \
}

//// "Microcode" similar to Uxn's, to be used with above macros
// Get argument from x or stack
#define ARGX rx ? (w ? *xw : *x) : (w ? us_popw() : us_pop())
// Get argument from y or stack
#define ARGY ry ? (w ? *yw : *y) : (w ? us_popw() : us_pop())
// Get arguments from x/y or stack
#define ARGXY rx ? (rx ? (w ? *xw : *x) : (w ? us_popw() : us_pop())) : (ry ? (w ? *yw : *y) : (w ? us_popw() : us_pop()))
// Push a value to stack
#define PUSH(n) do { if (w) { us_pushw(n); } else { us_push(n); } } while(0)
// Pop value from stack
#define POP() (w ? us_popw() : (xyw_word)us_pop())
// Read value from address
#define READ(addr) (w ? readw(addr) : (xyw_byte)xyw_memory[addr])
// Write value as word or byte depending on width flag 'w'
#define WRITE(addr, val) w ? writew(addr, val) : write(addr, (xyw_byte)(val))

//// Initialize devices

xyw_device xyw_devices[XYW_TOTAL_DEVICES] = {
    // System device (not implemented yet)
    {&xyw_memory[0xff00], 0, 0},
    // Console device
    {&xyw_memory[0xff10], 0, console_output},
    // ...
};

// clang-format on
int xyw_run(const char *input)
{
    // Load bytecode into memory
    FILE *fp = fopen(input, "rb");
    if (!fp)
    {
        fprintf(stderr, "Error - Could not open file [%s]\n", input);
        return -1;
    }
    // Check header
    xyw_byte header[6];
    xyw_dbg("Reading header from file [%s]\n", input);
    fread(header, 1, 6, fp);
    if (header[0] != 0x0F || header[1] != 'x' || header[2] != 'y' || header[3] != 'w' || header[4] != 0x00 || header[5] != 0xF0)
    {
        fprintf(stderr, "Error - Invalid file header [%s]\n", input);
        fclose(fp);
        return -1;
    }
    xyw_dbg("File header OK\n");
    // Load the rest of the file into memory
    size_t bytes_read = fread(&xyw_memory[USER_MEMORY_START], 1, MEMORY_SIZE - USER_MEMORY_START, fp);
    if (bytes_read == 0)
    {
        fprintf(stderr, "Error - Could not read file contents [%s]\n", input);
        fclose(fp);
        return -1;
    }
    xyw_dbg("Loaded %zu bytes into memory\n", bytes_read);
    fclose(fp);
    xyw_memory[SYSTEM_STATE] = 1;
    while (xyw_memory[SYSTEM_STATE] && *pc < SYSTEM_MEMORY_START)
    {
        xyw_dbg("PC: %04X -> %02X (%s)\n", *pc, xyw_memory[*pc], xyw_instructions[xyw_memory[*pc] & 0x1F]);
        switch (xyw_memory[*pc])
        {
            // clang-format off
            /* BRK */ case 0x00: break;
            /* PSH */ OPC_PSH(0x01);
            /* POP */ OPC_POP(0x02);
            /* LDR */ OPC_R1(0x03, PUSH(READ(a)));
            /* LDA */ OPC_S(0x04, READ(us_popw()));
            /* STA */ OPC_SR1(0x05, xyw_word a = us_popw(), WRITE(a, ARGXY));
            /* LDP */ OPC_S(0x06, READ(DIRECT_PAGE | us_pop()));
            /* STP */ OPC_SR1(0x07, xyw_word a = DIRECT_PAGE | us_pop(), WRITE(a, ARGXY));
            /* INC */ OPC_IDC(0x08, +1);
            /* DEC */ OPC_IDC(0x09, -1);
            /* SHL */ OPC_SR2(0x0A, PUSH(b << a));
            /* SHR */ OPC_SR2(0x0B, PUSH(b >> a));
            /* ADD */ OPC_SR2(0x0C, PUSH(a + b));
            /* SUB */ OPC_SR2(0x0D, PUSH(b - a));
            /* MUL */ OPC_SR2(0x0E, PUSH(a * b));
            /* DIV */ OPC_SR2(0x0F, PUSH(b / a));
            /* EQU */ OPC_SR2(0x10, PUSH(b == a));
            /* NEQ */ OPC_SR2(0x11, PUSH(b != a));
            /* GTH */ OPC_SR2(0x12, PUSH(b > a));
            /* LTH */ OPC_SR2(0x13, PUSH(b < a));
            /* AND */ OPC_SR2(0x14, PUSH(b & a));
            /* IOR */ OPC_SR2(0x15, PUSH(b | a));
            /* XOR */ OPC_SR2(0x16, PUSH(b ^ a));
            /* NOT */ OPC_SR1(0x17, , PUSH(~ARGXY));
            /* SWP */ OPC_SWP(0x18);
            /* DUP */ OPC_SR1(0x19, , PUSH(ARGXY));
            /* OVR */ OPC_SR2(0x1A, PUSH(b); PUSH(a); PUSH(b));
            /* ROT */ OPC_SR2(0x1B, PUSH(b); PUSH(POP()); PUSH(a));
            /* JMP */ case 0x1C: *pc = us_popw(); break;
            /* JCN */ OPC_SR1(0x1D, , if (ARGXY) { *pc = POP(); } );
            /* JSR */ case 0x1E: xyw_word addr = us_popw(); ss_pushw(*pc); *pc = addr; break;
            /* RTS */ case 0x1F: *pc = ss_popw(); break;
            // clang-format on
        }
        // TODO: error handling etc.
        (*pc)++;
    }
    return 0;
}