#include "core1.h"
#include <stdint.h>
#include <stdio.h>
#include <stdatomic.h>

#include "pico/stdlib.h"
#include "pico/multicore.h"
#include "hardware/irq.h"
#include "hardware/pio.h"

#include "core_comm.h"
#include "z80_bus.pio.h"
#include "text_mode.h"
#include "cvideo.h"

#define Z80_PIO          pio1
#define Z80_MONITOR_SM   0

#define Z80_OE_PIN       0
#define Z80_DIR_PIN      1
#define Z80_DATA_PIN     2
#define Z80_DATA_WIDTH   8
#define Z80_IOREQ_PIN    10
#define Z80_RD_PIN       11
#define Z80_CS_PIN       12
#define Z80_A0_PIN       13
#define Z80_A1_PIN       14
#define Z80_CLK_PIN      15
#define Z80_WAIT_PIN     16

#define Z80_RD_MASK      (1u << Z80_RD_PIN)
#define Z80_CS_MASK      (1u << Z80_CS_PIN)
#define Z80_ADDR_MASK    0x3u
#define Z80_DATA_MASK    0xffu
#define Z80_DIRTY_SHIFT  32u
#define Z80_DIRTY_MASK   (UINT64_C(0x0f) << Z80_DIRTY_SHIFT)

// Bits 0..31 contain registers 0..3. Bits 32..35 are their dirty flags.
static _Atomic uint64_t z80_registers;

uint16_t cursor_x = 0;
uint16_t cursor_y = 0;
uint16_t pixel_x = 0;
uint16_t pixel_y = 0;

#define VID_CMD_SET_TEXT_MODE 0
#define VID_CMD_SET_PIXEL_MODE 1
#define VID_CMD_CLR 2
#define VID_CMD_TEXT_POS_X 3
#define VID_CMD_PIXEL_POS_X 4
#define VID_CMD_TEXT_POS_Y 5
#define VID_CMD_PIXEL_POS_Y 6


typedef enum VID_MODE{
    VID_MODE_TEXT,
    VID_MODE_PIXEL,
    VID_MODE_TEXT_POS_X,
    VID_MODE_TEXT_POS_Y,
    VID_MODE_PIXEL_POS_X_1,
    VID_MODE_PIXEL_POS_X_2,
    VID_MODE_PIXEL_POS_Y_1,
    VID_MODE_PIXEL_POS_Y_2,
} vid_mode_t;
vid_mode_t vid_mode = VID_MODE_TEXT;


// --------------------------------
// CORE 1
// --------------------------------

static void core1_putc(char c){
    video_command_t cmd;

    switch((uint8_t)c){
        case '\r':
            cursor_x = 0;
            break;
        case '\n':
            cursor_y = (cursor_y<CHAR_LINES-1) ? cursor_y+1 : cursor_y;
            break;
        default:
            cmd.cmd = VIDEO_WRITE_TEXT,
            cmd.data[0] = (uint16_t)c,
            cmd.data[1] = cursor_x,
            cmd.data[2] = cursor_y,
            core1_to_core0_write_blocking(*((uint64_t*)&cmd));
            cursor_x = (cursor_x<CHARS_PER_LINE-1) ? cursor_x+1 : cursor_x;
    }
}

static void core1_puts(char * s){
    while(*s){
        core1_putc(*s);
        s++;
    }
}

static void z80_bus_irq_handler(void) {
    while (!pio_sm_is_rx_fifo_empty(Z80_PIO, Z80_MONITOR_SM)) {
        uint32_t bus_sample = pio_sm_get(Z80_PIO, Z80_MONITOR_SM);
        uint8_t addr = (bus_sample >> Z80_A0_PIN) & Z80_ADDR_MASK;
        uint32_t shift = (uint32_t)addr * 8u;

        uint64_t registers = atomic_load_explicit(&z80_registers,
                                                  memory_order_relaxed);

        if ((bus_sample & Z80_RD_MASK) != 0) {
            uint8_t data = (bus_sample >> Z80_DATA_PIN) & Z80_DATA_MASK;
            registers &= ~(UINT64_C(0xff) << shift);
            registers |= (uint64_t)data << shift;
            registers |= UINT64_C(1) << (Z80_DIRTY_SHIFT + addr);
            atomic_store_explicit(&z80_registers, registers,
                                  memory_order_release);
        } else {
            uint8_t data = (registers >> shift) & Z80_DATA_MASK;
            pio_sm_put(Z80_PIO, Z80_MONITOR_SM, data);
        }
    }
}

static void z80_bus_init(void) {
    uint offset = pio_add_program(Z80_PIO, &z80_bus_program);

    // Keep the bus isolated while selecting B (Z80) -> A (Pico).
    gpio_init(Z80_OE_PIN);
    gpio_put(Z80_OE_PIN, 1);
    gpio_set_dir(Z80_OE_PIN, GPIO_OUT);
    gpio_init(Z80_DIR_PIN);
    gpio_put(Z80_DIR_PIN, 0);
    gpio_set_dir(Z80_DIR_PIN, GPIO_OUT);
    gpio_init(Z80_WAIT_PIN);
    gpio_put(Z80_WAIT_PIN, 1);
    gpio_set_dir(Z80_WAIT_PIN, GPIO_OUT);

    for (uint pin = Z80_OE_PIN; pin <= Z80_CLK_PIN; pin++) {
        pio_gpio_init(Z80_PIO, pin);
    }
    pio_gpio_init(Z80_PIO, Z80_WAIT_PIN);

    pio_sm_config config = z80_bus_program_get_default_config(offset);
    sm_config_set_in_pins(&config, Z80_OE_PIN);
    sm_config_set_in_shift(&config, false, false, 32);
    sm_config_set_out_pins(&config, Z80_DATA_PIN, Z80_DATA_WIDTH);
    sm_config_set_out_shift(&config, true, false, 32);
    sm_config_set_set_pins(&config, Z80_WAIT_PIN, 1);
    sm_config_set_sideset_pins(&config, Z80_OE_PIN);
    pio_sm_init(Z80_PIO, Z80_MONITOR_SM, offset, &config);

    uint32_t control_mask = (1u << Z80_OE_PIN) |
                            (1u << Z80_DIR_PIN) |
                            (1u << Z80_WAIT_PIN);
    uint32_t initial_levels = (1u << Z80_WAIT_PIN);
    pio_sm_set_pins_with_mask(Z80_PIO, Z80_MONITOR_SM,
                              initial_levels, control_mask);
    pio_sm_set_pindirs_with_mask(Z80_PIO, Z80_MONITOR_SM,
                                 control_mask, control_mask);
    pio_sm_set_consecutive_pindirs(Z80_PIO, Z80_MONITOR_SM,
                                   Z80_DATA_PIN,
                                   Z80_CLK_PIN - Z80_DATA_PIN + 1,
                                   false);

    // IRQ configuration is per-core. This function runs on core 1, so bus
    // events preempt diagnostics on core 1 and never involve core 0.
    pio_set_irq0_source_enabled(Z80_PIO,
                                pis_sm0_rx_fifo_not_empty,
                                true);
    irq_set_exclusive_handler(PIO1_IRQ_0, z80_bus_irq_handler);
    irq_set_priority(PIO1_IRQ_0, PICO_HIGHEST_IRQ_PRIORITY);
    irq_set_enabled(PIO1_IRQ_0, true);

    pio_sm_set_enabled(Z80_PIO, Z80_MONITOR_SM, true);
}

void core1_entry() {
    z80_bus_init();
    core1_puts("core 1: Z80 read/write bus started > ");

    while (true) {
        // Take a coherent snapshot and consume all four dirty flags. A
        // concurrent later write sets its register's flag again.
        uint64_t registers_i = atomic_fetch_and_explicit(&z80_registers, ~Z80_DIRTY_MASK, memory_order_acq_rel);
        registers_t regs = *((registers_t*)&registers_i);
        if((regs.flags & 0x01) == 0x01){
            // Write to vid_cmd happened
            switch(regs.vid_cmd){
                case VID_CMD_SET_TEXT_MODE: {
                    video_command_t cmd = {
                        .cmd = VIDEO_COMMAND_MODE_SET,
                        .data[0] = VIDEO_MODE_TEXT,
                    };
                    core1_to_core0_write_blocking(*((uint64_t*)&cmd));
                    vid_mode = VID_MODE_TEXT;
                } break;
                case VID_CMD_SET_PIXEL_MODE: {
                    video_command_t cmd = {
                        .cmd = VIDEO_COMMAND_MODE_SET,
                        .data[0] = VIDEO_MODE_PIXEL,
                    };
                    core1_to_core0_write_blocking(*((uint64_t*)&cmd));
                    vid_mode = VID_MODE_PIXEL;
                } break;
                case VID_CMD_CLR: {
                    if(vid_mode == VIDEO_MODE_TEXT){
                        video_command_t cmd = {
                            .cmd = VIDEO_CLR_TEXT
                        };
                        core1_to_core0_write_blocking(*((uint64_t*)&cmd));
                    }else if(vid_mode == VIDEO_MODE_PIXEL){
                        video_command_t cmd = {
                            .cmd = VIDEO_CLR_PIXEL
                        };
                        core1_to_core0_write_blocking(*((uint64_t*)&cmd));
                    }
                } break;
                case VID_CMD_TEXT_POS_X: {
                    vid_mode = VID_MODE_TEXT_POS_X;
                } break;
                case VID_CMD_PIXEL_POS_X: {
                    vid_mode = VID_MODE_PIXEL_POS_X_1;
                } break;
                case VID_CMD_TEXT_POS_Y: {
                    vid_mode = VID_MODE_TEXT_POS_Y;
                } break;
                case VID_CMD_PIXEL_POS_Y: {
                    vid_mode = VID_MODE_PIXEL_POS_Y_1;
                } break;
                default:
            }
        }
        if((regs.flags & 0x02) == 0x02){
            // Write to vid_dat happened
            switch(vid_mode){
                case VID_MODE_TEXT_POS_X:
                    cursor_x = (regs.vid_data<CHARS_PER_LINE) ? regs.vid_data : cursor_x;
                    break;
                case VID_MODE_TEXT_POS_Y:
                    cursor_y = (regs.vid_data<CHAR_LINES) ? regs.vid_data : cursor_y;
                    break;

                case VID_MODE_PIXEL_POS_X_1:
                    pixel_x = pixel_x&0xff00 | regs.vid_data;
                    vid_mode = VID_MODE_PIXEL_POS_X_2;
                    break;
                case VID_MODE_PIXEL_POS_X_2:
                    pixel_x = pixel_x&0x00ff | regs.vid_data<<8;
                    vid_mode = VID_MODE_PIXEL_POS_X_1;
                    pixel_x = (pixel_x<VIDEO_WIDTH) ? pixel_x : 0;
                    break;
                case VID_MODE_PIXEL_POS_Y_1:
                    pixel_y = pixel_y&0xff00 | regs.vid_data;
                    vid_mode = VID_MODE_PIXEL_POS_Y_2;
                    break;
                case VID_MODE_PIXEL_POS_Y_2:
                    pixel_y = pixel_y&0x00ff | regs.vid_data<<8;
                    vid_mode = VID_MODE_PIXEL_POS_Y_1;
                    pixel_y = (pixel_y<VIDEO_HEIGHT) ? pixel_y : 0;
                    break;

                case VID_MODE_PIXEL:{
                    video_command_t cmd = {
                        .cmd = VIDEO_WRITE_PIXEL,
                        .data[0] = regs.vid_data,
                        .data[1] = pixel_x,
                        .data[2] = pixel_y,
                    };
                    pixel_x = (pixel_x<VIDEO_WIDTH-8) ? pixel_x+8 : pixel_x;
                    core1_to_core0_write_blocking(*((uint64_t*)&cmd));
                } break;

                case VID_MODE_TEXT:
                default:
                    core1_putc(regs.vid_data);
            }
        }

        tight_loop_contents();
    }
}