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

#include "devices/system.h"
#include "devices/terminal.h"
#include "devices/clock.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 addresses */
#define SYSTEM_STATE 0xff00
#define SYSTEM_ERROR 0xff01
#define SYSTEM_S 0xff02
#define SYSTEM_U 0xff03
#define SYSTEM_PAGE 0xff04
#define SYSTEM_RANDOM 0xff05
#define SYSTEM_ON_ERROR 0xff06

/* Terminal device addresses */
#define TERMINAL_INPUT 0xff10
#define TERMINAL_ON_KEYPRESS 0xff12

#ifdef _WIN32
#include <conio.h>
int getch()
{
    return _getch();
}
#else
int getch()
{
    return getchar();
}
#endif
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)
{
    // Given $FF41, it should return $01
    return (xyw_byte)((addr - DEVICE_AREA_START) & 0x0F);
}

//// 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 int w=0, rx=0, ry=0; xyw_byte b = us_pop(), a = us_pop(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_X|opc: { const int w=0, rx=1, ry=0; xyw_byte b = *x, a = us_pop(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_Y|opc: { const int w=0, rx=0, ry=1; xyw_byte b = *y, a = us_pop(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_X|XYW_MODE_Y|opc: { const int w=0, rx=1, ry=1; xyw_byte a = *x, b = *y; body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_W|opc: { const int w=1, rx=0, ry=0; xyw_word b = us_popw(); xyw_word a = us_popw(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_X|XYW_MODE_W|opc: { const int w=1, rx=1, ry=0; xyw_word b = *xw; xyw_word a = us_popw(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_Y|XYW_MODE_W|opc: { const int w=1, rx=0, ry=1; xyw_word b = *yw; xyw_word a = us_popw(); body; (void)(rx); (void)(ry); } break; \
    case XYW_MODE_X|XYW_MODE_Y|XYW_MODE_W|opc: { const int w=1, rx=1, ry=1; xyw_word a = *xw; xyw_word b = *yw; body; (void)(rx); (void)(ry); } 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; \
}

// Create address based on direct page or full address
#define ADDR(ad) w ? ad : (DIRECT_PAGE | ad) 
// 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() : us_pop())
// Get argument from x register or stack
#define ARGX (rx ? (w ? *xw : *x) : (w ? us_popw() : us_pop()))
// Get argument from y register or stack
#define ARGY (ry ? (w ? *yw : *y) : (w ? us_popw() : us_pop()))
// Get argument from x or y register or stack
#define ARG (rx ? ARGX : (ry ? ARGY : POP()))
//#define ARG (rx ? (w ? *xw : *x) : (ry ? (w ? *yw : *y) : (w ? us_popw() : us_pop())))
// Set value to x register or push to stack
#define SETX(v) rx ? (w ? (*xw = v) : (*x = v)) : (w ? us_pushw(v) : us_push(v))
// Set value to y register or push to stack
#define SETY(v) ry ? (w ? (*yw = v) : (*y = v)) : (w ? us_pushw(v) : us_push(v))
// Set value to x or y register or push to stack
#define SET(v) do { if (rx) { SETX(v); } else if (ry) { SETY(v); } else { PUSH(v); } } while(0)

// Read a value from memory
#define READ(addr) (w ? readw(addr) : read(addr))
// Read top of the user stack
#define TOP READ(USER_STACK_START + (*u) - (w ? 2 : 1))

//// Initialize devices

xyw_device xyw_devices[XYW_TOTAL_DEVICES] = {
    // System device (not implemented yet)
    {&xyw_memory[0xff00], system_init, system_input, 0},
    // Terminal device
    {&xyw_memory[0xff10], 0, terminal_input, terminal_output},
    // Clock device
    {&xyw_memory[0xff20], 0, clock_input, 0},
    // File device
    {&xyw_memory[0xff30], 0, 0, 0},
    // ...
};
// clang-format on

//// Main evaluation function - runs from current pc until BRK
//// Returns 1 on success, 0 on error
int xyw_eval(xyw_word start_pc)
{
    *pc = start_pc;
    while (xyw_memory[SYSTEM_STATE] && *pc < SYSTEM_MEMORY_START)
    {
        xyw_dbg("PC: %04X -> %02X (%s) [U:%02X|S:%02X|X:%02X|Y:%02X|XW:%04X|YW:%04X]\n", *pc, xyw_memory[*pc], xyw_instructions[xyw_memory[*pc] & 0x1F], *u, *s, *x, *y, *xw, *yw);
        switch (xyw_memory[*pc])
        {
            // clang-format off
            /* BRK */ case 0x00: return 1; // Simply return from evaluation
            /* PSH */ OPC_SR1(0x01, , if (rx||ry) { PUSH(ARG); } else { (*pc)++; PUSH(READ(*pc)); if (w) { (*pc)++; } })
            /* POP */ OPC_SR1(0x02, , if (rx||ry) { SET(POP()); } else { POP(); } );
            /* CLR */ OPC_SR1(0x03, , SET(0));
            /* LDB */ OPC_SR1(0x04, , us_push(read(ADDR(ARG))));
            /* LDW */ OPC_SR1(0x05, , us_pushw(readw(ADDR(ARG)) ));
            /* STB */ OPC_SR1(0x06, xyw_word addr = ADDR(ARG), write( addr, us_pop() ));
            /* STW */ OPC_SR1(0x07, xyw_word addr = ADDR(ARG), writew( addr, us_popw() ));
            /* INC */ OPC_SR1(0x08, , SET(ARG + 1));
            /* DEC */ OPC_SR1(0x09, , SET(ARG - 1));
            /* SHL */ OPC_SR1(0x0A, xyw_byte b = us_pop(), PUSH(ARG << b));
            /* SHR */ OPC_SR1(0x0B, xyw_byte b = us_pop(), PUSH(ARG >> b));
            /* ADD */ OPC_SR2(0x0C, PUSH(a + b));
            /* SUB */ OPC_SR2(0x0D, PUSH(a - b));
            /* MUL */ OPC_SR2(0x0E, PUSH(a * b));
            /* DIV */ OPC_SR2(0x0F, PUSH(a % b); PUSH(a / b));
            /* EQU */ OPC_SR2(0x10, PUSH(a == b));
            /* NEQ */ OPC_SR2(0x11, PUSH(a != b));
            /* GTH */ OPC_SR2(0x12, PUSH(a > b));
            /* LTH */ OPC_SR2(0x13, PUSH(a < b));
            /* AND */ OPC_SR2(0x14, PUSH(a & b));
            /* IOR */ OPC_SR2(0x15, PUSH(a | b));
            /* XOR */ OPC_SR2(0x16, PUSH(a ^ b));
            /* NOT */ OPC_SR1(0x17, , SET(~ARG));
            /* SWP */ OPC_SR2(0x18, if ((rx||ry)) { SET(b); } else { PUSH(a); PUSH(b); } );
            /* DUP */ OPC_SR1(0x19, , PUSH(ARG));
            /* OVR */ OPC_SR2(0x1A, PUSH(b); PUSH(a); PUSH(b));
            /* ROT */ OPC_SR2(0x1B, PUSH(b); PUSH(POP()); PUSH(a));
            /* JMP */ OPC_SR1(0x1C, , *pc = ADDR(ARG)-1;);
            /* JCN */ OPC_SR1(0x1D, xyw_word addr = ADDR(ARG), if (us_pop()) { *pc = addr-1; });
            /* JSR */ OPC_SR1(0x1E, xyw_word addr = ADDR(ARG),  ss_pushw(*pc+1); *pc = addr-1;);
            /* RTS */ case 0x1F: *pc = ss_popw()-1; break;
            // clang-format on
        }
        (*pc)++;
    }
    return 0;
}

//// Main run function
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;
    }
    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;
    // Initialize devices
    for (int i = 0; i < XYW_TOTAL_DEVICES; i++)
    {
        if (xyw_devices[i].init)
        {
            xyw_devices[i].init(xyw_devices[i].data);
        }
    }
    // Set direct page to device page by default
    xyw_memory[SYSTEM_PAGE] = 0xff;

    int eval_result = xyw_eval(0x0000);
    if (eval_result)
    {
        xyw_word terminal_vector = XYW_PEEKW(&xyw_memory[TERMINAL_ON_KEYPRESS]);

        // If terminal vector is set, enter stdin reading loop
        if (terminal_vector)
        {
            xyw_dbg("terminal.on_keypress vector set to %04X, entering input loop\n", terminal_vector);
            while (xyw_memory[SYSTEM_STATE])
            {
                int c = getch();
                // if CTRL+C or CTRL+D is pressed, exit
                if (c == 3 || c == 4)
                {
                    return 0;
                }
                write(TERMINAL_INPUT, (xyw_byte)c);
            }
        }
    }

    return 0;
}