#include <stdio.h>
#include "xyw.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

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

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

/* console device */
#define CONSOLE_BASE 0xff10
#define CONSOLE_INPUT 0xff10
#define CONSOLE_OUTPUT 0xff11
#define CONSOLE_ERROR 0xff12
#define CONSOLE_ON_KEYPRESS 0xff13

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

//// 1-byte pointers and registers
xyw_byte *ssp = &xyw_memory[SYSTEM_SSP];
xyw_byte *usp = &xyw_memory[SYSTEM_USP];
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;

//// 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 manage 2-byte values to memory (big-endian)

static inline xyw_word mget(xyw_word addr)
{
    return ((xyw_word)xyw_memory[addr] << 8) | xyw_memory[addr + 1];
}

static inline void mset(xyw_word addr, xyw_word val)
{
    xyw_memory[addr] = (val >> 8) & 0xFF;
    xyw_memory[addr + 1] = val & 0xFF;
}

static inline void madd(xyw_word addr, int delta)
{
    xyw_word val = mget(addr);
    mset(addr, (xyw_word)(val + delta));
}

//// System stack management

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

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

//// User stack management

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

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

//// Instructions

static void PSH()
{
    (*pc)++;
    xyw_byte val = xyw_memory[*pc];
    us_push(val);
}

static void PSHw()
{
    (*pc)++;
    xyw_word val = mget(*pc);
    us_pushw(val);
    (*pc)++;
}

static void PSHx() { us_push(*x); }
static void PSHy() { us_push(*y); }
static void PSHxw() { us_pushw(*xw); }
static void PSHyw() { us_pushw(*yw); }

static xyw_byte POP()
{
    if (*usp < 1)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return -1;
    }
    xyw_byte val = user_stack[--(*usp)];
    user_stack[*usp] = 0;
    return val;
}

static xyw_word POPw()
{
    if (*usp < 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return -1;
    }
    xyw_word val = mget(USER_STACK_START + (*usp) - 2);
    mset(USER_STACK_START + (*usp) - 2, 0);
    (*usp) -= 2;
    return val;
}

static void POPx()
{
    if (*usp < 1)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return;
    }
    *x = user_stack[--(*usp)];
    user_stack[*usp] = 0;
}

static void POPy()
{
    if (*usp < 1)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return;
    }
    *y = user_stack[--(*usp)];
    user_stack[*usp] = 0;
}

static void POPxw()
{
    if (*usp < 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return;
    }
    *xw = mget(USER_STACK_START + (*usp) - 2);
    mset(USER_STACK_START + (*usp) - 2, 0);
    (*usp) -= 2;
}

static void POPyw()
{
    if (*usp < 2)
    {
        xyw_memory[SYSTEM_ERROR] = XYW_ERROR_STACK_UNDERFLOW;
        return;
    }
    *yw = mget(USER_STACK_START + (*usp) - 2);
    mset(USER_STACK_START + (*usp) - 2, 0);
    (*usp) -= 2;
}

static void LDRx() { us_push(xyw_memory[*x]); }
static void LDRy() { us_push(xyw_memory[*y]); }
static void LDRxw() { us_pushw(mget(*xw)); }
static void LDRyw() { us_pushw(mget(*yw)); }

static void LDA() { us_push(xyw_memory[POPw()]); }
static void LDAw() { us_pushw(mget(POPw())); }

static void LDP()
{
    xyw_byte page = xyw_memory[SYSTEM_PAGE];
    xyw_word addr = (page << 8) | POP();
    us_push(xyw_memory[addr]);
}

static void LDPw()
{
    xyw_byte page = xyw_memory[SYSTEM_PAGE];
    xyw_word addr = (page << 8) | POPw();
    us_pushw(mget(addr));
}

static void STA(xyw_byte val)
{
    xyw_word addr = POPw();
    xyw_memory[addr] = val;
}

static void STAw(xyw_word val)
{
    xyw_word addr = POPw();
    mset(addr, val);
}

static void STP(xyw_byte val)
{
    xyw_word addr = (xyw_memory[SYSTEM_PAGE] << 8) | POP();
    xyw_memory[addr] = val;
}

static void STPw(xyw_word val)
{
    xyw_word addr = (xyw_memory[SYSTEM_PAGE] << 8) | POPw();
    mset(addr, val);
}

static void SWP(xyw_byte a, xyw_byte b)
{
    us_push(a);
    us_push(b);
}

static void SWPw(xyw_word a, xyw_word b)
{
    us_pushw(a);
    us_pushw(b);
}

static void DUP(xyw_byte val)
{
    us_push(val);
    us_push(val);
}

static void DUPw(xyw_word val)
{
    us_pushw(val);
    us_pushw(val);
}

static void OVR(xyw_byte a, xyw_byte b)
{
    us_push(b);
    us_push(a);
    us_push(b);
}

static void OVRw(xyw_word a, xyw_word b)
{
    us_pushw(b);
    us_pushw(a);
    us_pushw(b);
}

static void ROT(xyw_byte a, xyw_byte b)
{
    us_push(b);
    us_push(POP());
    us_push(a);
}

static void ROTw(xyw_word a, xyw_word b)
{
    us_pushw(b);
    us_pushw(POPw());
    us_pushw(a);
}

static void JMP() { *pc = POPw(); }
static void JCN(xyw_byte a)
{
    if (a)
        *pc = POPw();
}
static void JCNw(xyw_word a)
{
    if (a)
        *pc = POPw();
}
static void JSR()
{
    xyw_word addr = POPw();
    ss_pushw(*pc);
    *pc = addr;
}
static void RTS()
{
    *pc = ss_popw();
}

//// Macros for instruction implementations

// clang-format off
// Infix operations
#define InOP(opcode, op)  \
    case opcode: us_push(POP() op POP()); break; \
    case XYW_MODE_X|opcode: us_push(*x op POP()); break; \
    case XYW_MODE_Y|opcode: us_push(*y op POP()); break; \
    case XYW_MODE_W|opcode: us_pushw(POPw() op POPw()); break; \
    case XYW_MODE_X|XYW_MODE_W|opcode: us_pushw(*xw op POPw()); break; \
    case XYW_MODE_Y|XYW_MODE_W|opcode: us_pushw(*yw op POPw()); break; \
    case XYW_MODE_X|XYW_MODE_Y|opcode: us_push(*x op *y); break; \
    case XYW_MODE_X|XYW_MODE_Y|XYW_MODE_W|opcode: us_pushw(*xw op *yw); break;

// Unary operations
#define UnOP(opcode, op)  \
    case opcode: us_push(op(POP())); break; \
    case XYW_MODE_X|opcode: us_push(op(*x)); break; \
    case XYW_MODE_Y|opcode: us_push(op(*y)); break; \
    case XYW_MODE_W|opcode: us_pushw(op(POPw())); break; \
    case XYW_MODE_X|XYW_MODE_W|opcode: us_pushw(op(*xw)); break; \
    case XYW_MODE_Y|XYW_MODE_W|opcode: us_pushw(op(*yw)); break;

// Postfix-ish operations (INC/DEC)
#define PoOP(opcode, op) \
    case opcode: us_push(POP() op); break; \
    case XYW_MODE_X|opcode: us_push(*x op); break; \
    case XYW_MODE_Y|opcode: us_push(*y op); break; \
    case XYW_MODE_W|opcode: us_pushw(POPw() op); break; \
    case XYW_MODE_X|XYW_MODE_W|opcode: us_pushw(*xw op); break; \
    case XYW_MODE_Y|XYW_MODE_W|opcode: us_pushw(*yw op); break;

// Prefix operations with two operands
#define PrOP2(opcode, op) \
    case opcode: op(POP(), POP()); break; \
    case XYW_MODE_X|opcode: op(*x, POP()); break; \
    case XYW_MODE_Y|opcode: op(*y, POP()); break; \
    case XYW_MODE_X|XYW_MODE_Y|opcode: op(*x, *y); break; \

// Prefix operations with two 2-byte operands
#define PrOP2w(opcode, op) \
    case XYW_MODE_W|opcode: op(POPw(), POPw()); break; \
    case XYW_MODE_X|XYW_MODE_W|opcode: op(*xw, POPw()); break; \
    case XYW_MODE_Y|XYW_MODE_W|opcode: op(*yw, POPw()); break; \
    case XYW_MODE_X|XYW_MODE_Y|XYW_MODE_W|opcode: op(*xw, *yw); break;

// Prefix operations with one operand
#define PrOP(opcode, op) \
    case opcode: op(POP()); break; \
    case XYW_MODE_X|opcode: op(*x); break; \
    case XYW_MODE_Y|opcode: op(*y); break;

// Prefix operations with one 2-byte operand
#define PrOPw(opcode, op) \
    case XYW_MODE_W|opcode: op(POPw()); break; \
    case XYW_MODE_X|XYW_MODE_W|opcode: op(*xw); break; \
    case XYW_MODE_Y|XYW_MODE_W|opcode: op(*yw); break;

// clang-format on

int 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 */ case 0x01: PSH(); break;
            /* PSHw */ case XYW_MODE_W|0x01: PSHw(); break;
            /* PSHx */ case XYW_MODE_X|0x01: PSHx(); break;
            /* PSHy */ case XYW_MODE_Y|0x01: PSHy(); break;
            /* PSHxw */ case XYW_MODE_X|XYW_MODE_W|0x01: PSHxw(); break;
            /* PSHyw */ case XYW_MODE_Y|XYW_MODE_W|0x01: PSHyw(); break;
            /* POP */ case 0x02: POP(); break;
            /* POPw */ case XYW_MODE_W|0x02: POPw(); break;
            /* POPx */ case XYW_MODE_X|0x02: POPx(); break;
            /* POPy */ case XYW_MODE_Y|0x02: POPy(); break;
            /* POPxw */ case XYW_MODE_X|XYW_MODE_W|0x02: POPxw(); break;
            /* POPyw */ case XYW_MODE_Y|XYW_MODE_W|0x02: POPyw(); break;
            /* LDRx */ case XYW_MODE_X|0x03: LDRx(); break;
            /* LDRy */ case XYW_MODE_Y|0x03: LDRy(); break;
            /* LDRxw */ case XYW_MODE_X|XYW_MODE_W|0x03: LDRxw(); break;
            /* LDRyw */ case XYW_MODE_Y|XYW_MODE_W|0x03: LDRyw(); break;
            /* LDA */ case 0x04: LDA(); break;
            /* LDAw */ case XYW_MODE_W|0x04: LDAw(); break;
            /* STA */ PrOP(0x05, STA);
            /* STAw */ PrOPw(0x05, STAw);
            /* LDP */ case 0x06: LDP(); break;
            /* LDPw */ case XYW_MODE_W|0x06: LDPw(); break;
            /* STP */ PrOP(0x07, STP); 
            /* STPw */ PrOPw(0x07, STPw);
            /* INC */ PoOP(0x08, + 1);
            /* DEC */ PoOP(0x09, - 1);
            /* SHL */ InOP(0x0A, <<);
            /* SHR */ InOP(0x0B, >>);
            /* ADD */ InOP(0x0C, +);
            /* SUB */ InOP(0x0D, -);
            /* MUL */ InOP(0x0E, *);
            /* DIV */ InOP(0x0F, /);
            /* EQU */ InOP(0x10, ==);
            /* NEQ */ InOP(0x11, !=);
            /* GTH */ InOP(0x12, >);
            /* LTH */ InOP(0x13, <);
            /* AND */ InOP(0x14, &);
            /* IOR */ InOP(0x15, |);
            /* XOR */ InOP(0x16, ^);
            /* NOT */ UnOP(0x17, ~);
            /* SWP  */ PrOP2(0x18, SWP);
            /* SWPw */ PrOP2w(0x18, SWPw);
            /* DUP */ PrOP(0x19, DUP);
            /* DUPw */ PrOPw(0x19, DUPw);
            /* OVR */ PrOP2(0x1A, OVR);
            /* OVRw */ PrOP2w(0x1A, OVRw);
            /* ROT */ PrOP2(0x1B, ROT);
            /* ROTw */ PrOP2w(0x1B, ROTw);
            /* JMP */ case 0x1C: JMP(); break;
            /* JCN */ PrOP(0x1D, JCN);
            /* JCNw */ PrOPw(0x1D, JCNw);
            /* JSR */ case 0x1E: JSR(); break;
            /* RTS */ case 0x1F: RTS(); break;
            // clang-format on
        }
        // TODO: error handling etc.
        (*pc)++;
    }

    return 0;
}