// OMSE: One More Spectrum Emulator (Mini Version)
// (c) 2025 Melissa O'Neill
//
// A Simple ZX Spectrum Emulator

#include <SDL.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <queue>
#include <memory>
#include <functional>
#include <chrono>
#include <cassert>
#include <atomic>
#include <optional>
#define CHIPS_IMPL
#define CHIPS_UTIL_IMPL
#include "z80.h"  // Andre Weissflog (Floooh)'s CHIPS Z80 emulator

// Timing constants (all in T-states)
constexpr uint32_t T_STATES_PER_LINE = 224;
constexpr uint32_t T_STATES_PER_FRAME = 69888;  // (64+192+56)*224
constexpr uint64_t CLOCK_RATE = 3'500'000;      // 3.5MHz

// Basic memory system - just implement what we need for display
class Memory {
 private:
    std::vector<uint8_t> ram_;

 public:
    Memory();

    uint8_t read(uint16_t address) const {
        return ram_[address];
    }

    void write(uint16_t address, uint8_t value) {
        if (address < 0x4000) {
            return;  // Ignore writes to ROM
        }
        ram_[address] = value;
    }

    void loadFromFile(const std::string& filename, uint16_t addr,
                      uint16_t size);
    void loadFromStream(std::istream& stream, uint16_t addr, uint16_t size);
};

Memory::Memory() : ram_(0x10000, 0) {  // 64K of RAM
    // Create a more recognizable pattern in screen memory
    for (uint16_t y = 0; y < 192; y++) {
        for (uint16_t x = 0; x < 32; x++) {
            uint16_t addr = 0x4000 + (y * 32) + x;
            // Create diagonal stripes
            ram_[addr] = ((x + (y / 8)) & 0x07) ? 0xAA : 0x55;
        }
    }

    // Set attributes to alternate colors
    for (uint16_t y = 0; y < 24; y++) {
        for (uint16_t x = 0; x < 32; x++) {
            uint16_t attr_addr = 0x5800 + (y * 32) + x;
            // Alternate between cyan on black and yellow on blue
            ram_[attr_addr] = ((x + y) & 1) ? 0x45 : 0x16;
        }
    }
}

void Memory::loadFromFile(const std::string& filename, uint16_t addr,
                          uint16_t size) {
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Could not open file: " + filename);
    }

    // Get file size
    file.seekg(0, std::ios::end);
    size_t fileSize = file.tellg();
    file.seekg(0, std::ios::beg);

    // Check if we have enough data after the offset
    if (fileSize < size) {
        throw std::runtime_error(
            "File too small: need at least 6912 bytes after offset");
    }

    // Read data into memory
    loadFromStream(file, addr, size);
}

void Memory::loadFromStream(std::istream& stream, uint16_t addr,
                            uint16_t size) {
    // Read data into memory
    stream.read(reinterpret_cast<char*>(&ram_[addr]), size);
}

// CRT display using SDL
class CRT {
 private:
    SDL_Window* window_;
    SDL_Renderer* renderer_;
    SDL_Texture* screenTexture_;
    std::vector<uint32_t> pixels_;
    bool odd_field_ = false;  // For interlacing
    bool flash_inverted_ = false;

 public:
    // Physical display characteristics
    static constexpr int TOTAL_WIDTH = 352;          // 352 pixels
    static constexpr int COLUMNS = TOTAL_WIDTH / 8;  // 352/8 columns
    static constexpr int FIELD_LINES = 312;    // Total PAL lines per field
    static constexpr int TOP_BLANKING = 16;    // Lines before visible area
    static constexpr int BOTTOM_BLANKING = 4;  // Lines after visible area
    static constexpr int VISIBLE_LINES =
        FIELD_LINES - TOP_BLANKING - BOTTOM_BLANKING;
    static constexpr int CRT_LINES = VISIBLE_LINES * 2;  // Two fields

 public:
    CRT() : window_(nullptr), renderer_(nullptr), screenTexture_(nullptr) {
        if (SDL_Init(SDL_INIT_VIDEO) < 0) {
            throw std::runtime_error("SDL initialization failed");
        }

        // Scale up 2x for better visibility
        window_ =
            SDL_CreateWindow("OMSE — One More Spectrum Emulator",
                             SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
                             TOTAL_WIDTH * 2, CRT_LINES, SDL_WINDOW_SHOWN);
        if (!window_) {
            throw std::runtime_error("Window creation failed");
        }

        renderer_ = SDL_CreateRenderer(window_, -1, SDL_RENDERER_ACCELERATED);
        if (!renderer_) {
            throw std::runtime_error("Renderer creation failed");
        }

        screenTexture_ = SDL_CreateTexture(renderer_, SDL_PIXELFORMAT_RGBA8888,
                                           SDL_TEXTUREACCESS_STREAMING,
                                           TOTAL_WIDTH, CRT_LINES);
        if (!screenTexture_) {
            throw std::runtime_error("Texture creation failed");
        }

        pixels_.resize(TOTAL_WIDTH * CRT_LINES, 0);

        // Set up 2x scaling horizontally, 1x vertically
        SDL_RenderSetScale(renderer_, 2.0f, 1.0f);
    }

    ~CRT() {
        if (screenTexture_) SDL_DestroyTexture(screenTexture_);
        if (renderer_) SDL_DestroyRenderer(renderer_);
        if (window_) SDL_DestroyWindow(window_);
        SDL_Quit();
    }

    // Update a group of 8 pixels at the specified location
    void updatePixels(uint32_t line, uint32_t column, uint8_t display_byte,
                      uint8_t attr_byte) {
        // Assertions
        assert(line >= 0 && line < FIELD_LINES);
        assert(column >= 0 && column < COLUMNS);

        // Ignore updates in blanking intervals
        if (line < TOP_BLANKING || line >= (TOP_BLANKING + VISIBLE_LINES)) {
            return;
        }
        line -= TOP_BLANKING;  // Adjust for top blanking

        // Interlace fields
        line = line * 2 + (odd_field_ ? 1 : 0);

        // Update 8 pixels at once
        uint32_t offset = (line * TOTAL_WIDTH) + (column * 8);
        // To bleed into the other line
        uint32_t bleed_offset =
            offset + (odd_field_ ? -TOTAL_WIDTH : TOTAL_WIDTH);

        // Convert attribute byte
        bool flash = (attr_byte & 0x80) != 0;
        bool bright = (attr_byte & 0x40) != 0;
        uint8_t paper = (attr_byte >> 3) & 0x07;
        uint8_t ink = attr_byte & 0x07;
        if (flash && flash_inverted_) {
            std::swap(paper, ink);
        }

        // Create RGB colors
        uint32_t paper_color = rgba_color_table[paper];
        uint32_t ink_color = rgba_color_table[ink];

        // Update 8 pixels - note MSB is leftmost pixel
        for (int bit = 7; bit >= 0; bit--) {
            bool pixel_set = (display_byte & (1 << bit)) != 0;
            uint32_t color = pixel_set ? ink_color : paper_color;
            // For main scanline, set the pixel with phosphor fade from previous
            pixels_[offset + (7 - bit)] =
                ((pixels_[offset + (7 - bit)] >> 2) & 0x3f3f3f3f) | color;

            // Other scanline is less bright, but bleeds through. Bright colors
            // bleed through more to create the bright effect.
            if (!bright) {
                color = ((color >> 1) & 0x7f7f7f7f) | 0xff;  // 50% brightness
            } else {
                color =
                    ((color >> 3) & 0x07070707) * 27 | 0xff;  // 84% brightness
            }
            pixels_[bleed_offset + (7 - bit)] =
                ((pixels_[bleed_offset + (7 - bit)] >> 2) & 0x3f3f3f3f) | color;
        }
    }

    void refresh() {
        SDL_UpdateTexture(screenTexture_, nullptr, pixels_.data(),
                          TOTAL_WIDTH * sizeof(uint32_t));
        SDL_RenderClear(renderer_);
        SDL_RenderCopy(renderer_, screenTexture_, nullptr, nullptr);
        SDL_RenderPresent(renderer_);
    }

    void toggleFlash() {
        flash_inverted_ = !flash_inverted_;
    }

 private:
    uint32_t rgba_color_table[8] = {
        0x00000000,  // Black
        0x0000FFFF,  // Blue
        0xFF000000,  // Red
        0xFF00FFFF,  // Magenta
        0x00FF0000,  // Green
        0x00FFFFFF,  // Cyan
        0xFFFF0000,  // Yellow
        0xFFFFFFFF   // White
    };
};

// Abstact I/O device class, supports read and write
class IODevice {
 public:
    virtual uint8_t read(uint16_t addr) = 0;
    virtual void write(uint16_t addr, uint8_t value) = 0;
};

// A bus for I/O devices
class IODeviceBus {
 private:
    using dev_mask_t = uint16_t;
    using dev_entry_t = std::pair<dev_mask_t, IODevice*>;
    std::vector<dev_entry_t> devices_;

 public:
    void addDevice(dev_mask_t mask, IODevice* device) {
        devices_.emplace_back(mask, device);
    }

    uint8_t read(uint16_t addr) {
        for (const auto& [mask, device] : devices_) {
            if (((~addr) & mask) == mask) {
                return device->read(addr);
            }
        }
        return 0xff;  // Default to all bits set
    }

    void write(uint16_t addr, uint8_t value) {
        for (const auto& [mask, device] : devices_) {
            if (((~addr) & mask) == mask) {
                device->write(addr, value);
                return;
            }
        }
    }
};

class CPU : public z80_t {
 public:
    CPU(Memory& memory, IODeviceBus& bus) : memory_(memory), bus_(bus) {
        pins_ = z80_init(this);
    }

    void tick() {
        pins_ = z80_tick(this, pins_);
    }

    void transact() {
        if ((pins_ & Z80_MREQ)) {
            const uint16_t addr = Z80_GET_ADDR(pins_);
            if (pins_ & Z80_RD) {
                Z80_SET_DATA(pins_, memory_.read(addr));
            } else if (pins_ & Z80_WR) {
                uint8_t data = Z80_GET_DATA(pins_);
                memory_.write(addr, data);
            }
        } else if (pins & Z80_IORQ) {
            if (pins & Z80_M1) {
                // Interrupt acknowledge
                Z80_SET_DATA(pins_, 0xff);
            } else {
                // I/O request
                const uint16_t addr = Z80_GET_ADDR(pins);
                if (pins & Z80_RD) {
                    Z80_SET_DATA(pins_, bus_.read(addr));
                } else if (pins & Z80_WR) {
                    bus_.write(addr, Z80_GET_DATA(pins_));
                }
            }
        }
    }

    void interrupt(bool status = true) {
        if (status) {
            pins_ |= Z80_INT;
        } else {
            pins_ &= ~Z80_INT;
        }
    }

    void setPC(uint16_t addr) {
        pins_ = z80_prefetch(this, addr);
    }

    uint64_t readPins() const {
        return pins_;
    }

 private:
    uint64_t pins_;
    Memory& memory_;
    IODeviceBus& bus_;
};

class System;
class ULA final : public IODevice {
 private:
    Memory& memory_;
    CRT& crt_;
    CPU& cpu_;
    uint8_t border_color_;
    uint8_t flash_flipper_ = FLASH_RATE;

    // Current position tracking
    uint32_t line_ = 0;  // Current scanline (0-311)
    uint32_t line_cycle_ =
        BORDER_T_STATES;  // Current cycle within line (0-223)
    uint32_t current_column_ = 0;

 public:
    ULA(Memory& mem, CPU& processor, CRT& display)
        : memory_(mem), crt_(display), cpu_(processor), border_color_(0) {
    }

    // Read and write to I/O ports
    uint8_t read(uint16_t addr) override {
        return 0xff;
    }

    void write(uint16_t addr, uint8_t value) override {
        setBorderColor(value);
    }

    void tick() {
        cpu_.tick();

        // Check if we're in the visible (non-blanking) area
        bool visible = (line_ >= CRT::TOP_BLANKING
                        && line_ < (CRT::FIELD_LINES - CRT::BOTTOM_BLANKING));

        if (visible
            && line_cycle_
                   < (CRT::COLUMNS
                      * 4)) {  // Only process during active display time
            // Every 4 cycles we output 8 pixels
            bool in_screen_line =
                (line_ >= SCREEN_START_LINE
                 && line_ < (SCREEN_START_LINE + SCREEN_HEIGHT));

            if (line_cycle_ % 4 == 0) {
                current_column_ = line_cycle_ / 4;
                bool in_screen_col =
                    (current_column_ >= SCREEN_START_COLUMN
                     && current_column_
                            < (SCREEN_START_COLUMN + SCREEN_WIDTH_BYTES));

                if (in_screen_line && in_screen_col) {
                    // We're in the actual screen area - fetch and display
                    // pixels
                    uint32_t screen_line = line_ - SCREEN_START_LINE;
                    uint32_t screen_col = current_column_ - SCREEN_START_COLUMN;

                    uint16_t addr =
                        calculateDisplayAddress(screen_line, screen_col);
                    uint8_t display_byte = memory_.read(addr);
                    uint8_t attr_byte = memory_.read(
                        calculateAttrAddress(screen_line, screen_col));

                    crt_.updatePixels(line_, current_column_, display_byte,
                                      attr_byte);
                } else {
                    // Border area
                    uint8_t border_attr = (border_color_ << 3);
                    crt_.updatePixels(line_, current_column_, 0x00,
                                      border_attr);
                }
            }
        }

        cpu_.transact();

        // Update position counters
        line_cycle_++;
        if (line_ == 0 && line_cycle_ == BORDER_T_STATES) {
            cpu_.interrupt();
        } else if (line_ == 0
                   && line_cycle_ == BORDER_T_STATES + INTERRUPT_DURATION) {
            cpu_.interrupt(false);
        }
        if (line_cycle_ >= T_STATES_PER_LINE) {
            line_cycle_ = 0;
            line_++;
            if (line_ >= FIELD_LINES) {
                line_ = 0;
                --flash_flipper_;
                if (flash_flipper_ == 0) {
                    flash_flipper_ = FLASH_RATE;
                    crt_.toggleFlash();
                }
            }
        }
    }

    void setBorderColor(uint8_t color) {
        border_color_ = color & 0x07;
    }

    uint8_t getBorderColor() const {
        return border_color_;
    }

 private:
    static constexpr uint32_t SCREEN_START_LINE = 64;
    static constexpr uint32_t SCREEN_START_COLUMN = 6;  // 48 pixels / 8
    static constexpr uint32_t SCREEN_WIDTH_BYTES = 32;
    static constexpr uint32_t SCREEN_HEIGHT = 192;
    static constexpr uint32_t BORDER_T_STATES = SCREEN_START_COLUMN * 4;
    static constexpr uint32_t SCREEN_WIDTH_T_STATES = SCREEN_WIDTH_BYTES * 4;
    static constexpr uint32_t FIELD_LINES =
        T_STATES_PER_FRAME / T_STATES_PER_LINE;
    static constexpr uint32_t FLASH_RATE = 16;
    static constexpr uint32_t INTERRUPT_DURATION = 32;

    uint16_t calculateDisplayAddress(uint32_t line, uint32_t col) {
        // Start of screen memory
        uint16_t addr = 0x4000;

        // Add Y portion
        addr |= ((line & 0xC0) << 5);  // Which third of the screen
        addr |= ((line & 0x07) << 8);  // Which character cell row
        addr |= ((line & 0x38) << 2);  // Remaining bits wherever

        // Add X portion
        addr |= col & 0b00011111;  // 5 bits of X go to bits 0-4

        return addr;
    }

    uint16_t calculateAttrAddress(uint32_t line, uint32_t col) {
        return 0x5800 + ((line >> 3) * 32) + col;
    }
};

class System {
 private:
    Memory memory_;
    IODeviceBus bus_;
    CRT crt_;
    CPU cpu_;
    ULA ula_;
    uint64_t current_t_state_;

    // Timing constants
    static constexpr uint64_t CHUNK_SIZE =
        13 * 8 * 224;  // Execute this many T-states at once

 public:
    System()
        : cpu_(memory_, bus_), ula_(memory_, cpu_, crt_), current_t_state_(0) {
        // Initialize subsystems
        bus_.addDevice(0x0001, &ula_);
    }

    void run() {
        bool quit = false;
        SDL_Event event;

        // Track both virtual and real time
        auto start_time = std::chrono::steady_clock::now();
        uint64_t next_refresh_t_state = current_t_state_;

        while (!quit) {
            // Handle SDL events
            while (SDL_PollEvent(&event)) {
                if (event.type == SDL_QUIT) {
                    quit = true;
                }
            }

            // Process a chunk of cycles
            uint64_t target_t_state = current_t_state_ + CHUNK_SIZE;
            while (current_t_state_ < target_t_state) {
                ula_.tick();
                current_t_state_++;
            }

            // Check if we need to refresh the display
            if (current_t_state_ >= next_refresh_t_state) {
                crt_.refresh();
                next_refresh_t_state += T_STATES_PER_FRAME;
            }

            // Sleep if we're ahead
            auto ahead_by =
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    start_time
                    + std::chrono::microseconds((current_t_state_ * 1000000)
                                                / CLOCK_RATE)
                    - std::chrono::steady_clock::now());

            if (ahead_by.count() > 1) {
                SDL_Delay(ahead_by.count() - 1);
            }
        }
    }

    void loadSNA(const std::string& filename);

    Memory& getMemory() {
        return memory_;
    }
};

// Load SNA snapshot
void System::loadSNA(const std::string& filename) {
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Could not open file: " + filename);
    }

    /*
        SNA Layout

        Offset   Size   Description
        -----------------------------------------------------------------------
        0        1      byte   I
        1        8      word   HL',DE',BC',AF'
        9        10     word   HL,DE,BC,IY,IX
        19       1      byte   Interrupt (bit 2 contains IFF2, 1=EI/0=DI)
        20       1      byte   R
        21       4      words  AF,SP
        25       1      byte   IntMode (0=IM0/1=IM1/2=IM2)
        26       1      byte   BorderColor (0..7, not used by Spectrum 1.7)
        27       49152  bytes  RAM dump 16384..65535
        -----------------------------------------------------------------------
        Total: 49179 bytes
    */

    auto readWord = [&file]() -> uint16_t {
        uint16_t value;
        file.read(reinterpret_cast<char*>(&value), sizeof(value));
        return value;
    };
    auto readByte = [&file]() -> uint8_t {
        uint8_t value;
        file.read(reinterpret_cast<char*>(&value), sizeof(value));
        return value;
    };

    cpu_.setPC(0x0072);
    cpu_.i = readByte();
    cpu_.hl2 = readWord();
    cpu_.de2 = readWord();
    cpu_.bc2 = readWord();
    cpu_.af2 = readWord();
    cpu_.hl = readWord();
    cpu_.de = readWord();
    cpu_.bc = readWord();
    cpu_.iy = readWord();
    cpu_.ix = readWord();
    cpu_.iff2 = (readByte() & 0x04) != 0;
    cpu_.r = readByte();
    cpu_.af = readWord();
    cpu_.sp = readWord();
    cpu_.im = readByte();
    ula_.setBorderColor(readByte());
    memory_.loadFromStream(file, 0x4000, 49152);
}

int main(int argc, char* argv[]) {
    try {
        // Parse command line arguments
        bool rom_loaded = false;

        System system;

        for (int i = 1; i < argc; i++) {
            std::string arg = argv[i];
            if (arg == "-h" || arg == "--help") {
                std::cout << "Usage: " << argv[0]
                          << " [options] [filename...]\n"
                          << "Options:\n"
                          << "  -h, --help           Show this help message\n"
                          << "If no filename is provided, boot into 48.rom.\n\n"
                          << "(.scr, .rom and .sna files are supported)\n";
                return 0;
            } else if (arg.ends_with(".rom")) {
                // Load the ROM file into memory
                system.getMemory().loadFromFile(arg, 0x0000, 16384);
                rom_loaded = true;
            } else if (arg.ends_with(".sna")) {
                // Load the SNA file into memory
                system.loadSNA(arg);
            } else if (arg.ends_with(".scr")) {
                // Load the SCR file into memory
                system.getMemory().loadFromFile(arg, 0x4000, 6912);
            } else {
                std::cerr << "Unknown file type: " << arg << std::endl;
                return 1;
            }
        }

        // Load the ROM (48.rom) into memory
        if (!rom_loaded) {
            system.getMemory().loadFromFile("48.rom", 0x0000, 16384);
        }

        system.run();
        return 0;
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
}