#include "zx_disk.h"
#include "z80.h"
#include <stdio.h>
#include <stdlib.h>

#define ZX_FDC_MODE_IDLE        0
#define ZX_FDC_MODE_COMMAND     1
#define ZX_FDC_MODE_EXECUTION   2
#define ZX_FDC_MODE_RESULT      3

#define ZX_FDC_DATA_DIRECTION_INPUT     0x00
#define ZX_FDC_DATA_DIRECTION_OUTPUT    0x40

#define ZX_FDC_COMMAND_MASK                     0x1F
#define ZX_FDC_COMMAND_SPECIFY                  0x03
#define ZX_FDC_COMMAND_SENSE_DRIVE_STATUS       0x04
#define ZX_FDC_COMMAND_RECALIBRATE              0x07
#define ZX_FDC_COMMAND_SENSE_INTERRUPT_STATUS   0x08
#define ZX_FDC_COMMAND_SEEK                     0x0F
#define ZX_FDC_COMMAND_READ_ID                  0x0A
#define ZX_FDC_COMMAND_READ_DATA                0x06
#define ZX_FDC_COMMAND_READ_DELETED_DATA        0x0C

#define ZX_FDC_MAIN_DRIVE0_BUSY 1
#define ZX_FDC_MAIN_DRIVE1_BUSY 2
#define ZX_FDC_MAIN_DRIVE2_BUSY 4
#define ZX_FDC_MAIN_DRIVE3_BUSY 8
#define ZX_FDC_MAIN_BUSY        16
#define ZX_FDC_MAIN_EXEC_MODE   32
#define ZX_FDC_MAIN_DATA_DIR    64 // 0: CPU->FDC   1: FDC->CPU
#define ZX_FDC_MAIN_DATA_READY  128


namespace zx_disk
{
    struct sector_t
    {
        uint8_t id          { 0 };
        uint8_t size        { 0 };
        uint8_t st1         { 0 };
        uint8_t st2         { 0 };
        uint16_t data_length{ 0 };
        uint8_t *data       { nullptr };
    };

    struct track_t
    {
        uint8_t size        { 0 };
        uint8_t sector_size { 0 };
        uint8_t num_sectors { 0 };
        uint8_t filler_byte { 0x35 };
        sector_t *sectors   { nullptr };
    };

    struct disk_t
    {
        uint8_t num_tracks  { 0 };
        uint8_t num_sides   { 0 };
        track_t *tracks     { nullptr };
    };

    disk_t disk;

    uint8_t mode = ZX_FDC_MODE_IDLE;
    uint8_t call_count = 0;
    //uint8_t data_direction = 0x00;  // 0x40
    uint8_t fdd0busy = 0;
    bool command_success = false;
    uint8_t srt, hlt, hut, nd;
    bool seeking = false;

    uint8_t current_head = 0;
    uint8_t current_drive = 0;
    uint8_t current_track = 0;
    uint8_t current_sector = 0;
    uint16_t current_byte = 0;

    uint16_t bytes_to_read = 0;
    uint8_t eot = 0;
    uint8_t st1 = 0;
    uint8_t st2 = 0;

    int zx_fdc_main_status_port_in(int port);
    int zx_fdc_data_port_in(int port);
    void zx_fdc_data_port_out(int port, int val);

    void init()
    {
        z80::connect_port(0x2ffd, 0xf002, zx_fdc_main_status_port_in, nullptr);
        z80::connect_port(0x3ffd, 0xf002, zx_fdc_data_port_in, zx_fdc_data_port_out);
        load("goldenaxe1.dsk");
    }

    void reset()
    {

    }

    void load(const char *filename)
    {
        FILE *f = fopen(filename, "rb");
        if (!f) return;
        fseek(f, 0, SEEK_END);
        const int file_size = ftell(f);
        fseek(f, 0, SEEK_SET);
        uint8_t *buffer = (uint8_t*)malloc(file_size);
        fread(buffer, file_size, 1, f);
        fclose(f);

        if (disk.tracks) {
            for (int i=0; i<disk.num_tracks; ++i) {
                for (int j=0; j<disk.tracks[i].num_sectors; ++j) {
                    free(disk.tracks[i].sectors[j].data);
                }
                free(disk.tracks[i].sectors);
            }
            free(disk.tracks);
        }

        // Hardcoded to read extended disks. Also, no security checks at all. Fuck the police.
        disk.num_tracks = buffer[0x30];
        disk.num_sides = buffer[0x31];
        disk.tracks = (track_t*)malloc(disk.num_tracks*sizeof(track_t));
        for (int i=0; i<disk.num_tracks; ++i) disk.tracks[i].size = buffer[0x34+i];
        uint32_t pos = 0x100;
        for (int i=0; i<disk.num_tracks; ++i)
        {
            track_t &track = disk.tracks[i];
            track.sector_size = buffer[pos+0x14];
            track.num_sectors = buffer[pos+0x15];
            track.filler_byte = buffer[pos+0x17];
            track.sectors = (sector_t*)malloc(track.num_sectors*sizeof(sector_t));
            pos += 0x18;
            for (int j=0; j<track.num_sectors; ++j)
            {
                sector_t &sector = track.sectors[j];
                sector.id = buffer[pos+0x02];
                sector.size = buffer[pos+0x03];
                sector.st1 = buffer[pos+0x04];
                sector.st2 = buffer[pos+0x05];
                sector.data_length = buffer[pos+0x06] + uint16_t((buffer[pos+0x07])<<8);
                if (sector.data_length==0) sector.data_length = uint16_t(sector.size)<<8;
                pos += 8;
            }
            if (pos&0xff) pos = (pos & 0xffffff00) + 0x100;
            for (int j=0; j<track.num_sectors; ++j)
            {
                sector_t &sector = track.sectors[j];
                const uint16_t size = uint16_t(sector.size)<<8;
                sector.data = (uint8_t*)malloc(size);
                for (int k=0; k<size; ++k) sector.data[k] = buffer[pos++];
            }
        }
    }

    uint8_t find_sector_id(uint8_t sector_id)
    {
        track_t &track = disk.tracks[current_track];
        for (int j=0; j<track.num_sectors; ++j)
        {
            if (track.sectors[j].id == sector_id) return j;
        }
        return 0xff;
    }

    uint8_t main_status_register()
    {
        const uint8_t val = (fdd0busy) |
                ((mode==ZX_FDC_MODE_IDLE) ? 0x00 : 0x10) |
                ((mode==ZX_FDC_MODE_EXECUTION) ? 0x20 : 0x00) |
                ((mode==ZX_FDC_MODE_RESULT) ? 0x40 : 0x00) |
                (0x80);
        return val;

    }

    int zx_fdc_main_status_port_in(int port)
    {
        //if (mode != ZX_FDC_MODE_EXECUTION) printf("FDC port 0x2ffd IN\n");
        return main_status_register();
    }

    uint8_t ST0()
    {
        return current_drive | (current_head<<2) | (seeking ? 0x20 : 0x00) | (fdd0busy ? 0x00 : 0x80);
    }

    uint8_t ST3()
    {
        const bool ready = (current_head==0) && (current_drive==0);
        return current_drive | (current_head<<2) | ((current_track==0)?0x10:0) |
            (ready?0x20:0) | 0x00; // TS
    }

    void start_command(uint8_t command);

    uint8_t (*process_current_command)(uint8_t) = nullptr;
    uint8_t process_command_unknown(uint8_t command);
    uint8_t process_command_specify(uint8_t command);
    uint8_t process_command_sense_drive_status(uint8_t command);
    uint8_t process_command_recalibrate(uint8_t command);
    uint8_t process_command_sense_interrupt_status(uint8_t command);
    uint8_t process_command_seek(uint8_t command);
    uint8_t process_command_read_id(uint8_t command);
    uint8_t process_command_read_data(uint8_t command);

    int zx_fdc_data_port_in(int port)
    {
        //if (mode != ZX_FDC_MODE_EXECUTION) printf("FDC port 0x3ffd IN\n");
        if (mode == ZX_FDC_MODE_COMMAND) {
            printf("IGNORED!\n");
            return 0;
        }
        if (process_current_command) return process_current_command(0);
        return 0;
    }

    void zx_fdc_data_port_out(int port, int val)
    {
        //printf("FDC port 0x3ffd OUT: 0x%02x\n", val);
        if (mode == ZX_FDC_MODE_RESULT) {
            printf("IGNORED!\n");
            return;
        }
        if (process_current_command)
            process_current_command(val);
        else
            start_command(val);
    }

    void start_command(uint8_t command)
    {
        if ((fdd0busy) && ((command & ZX_FDC_COMMAND_MASK) != ZX_FDC_COMMAND_SENSE_INTERRUPT_STATUS)) return;

        mode = ZX_FDC_MODE_COMMAND;
        call_count = 1;
        debug::composeCommand(command);
        switch(command & ZX_FDC_COMMAND_MASK)
        {
            case ZX_FDC_COMMAND_SPECIFY:
                process_current_command = process_command_specify;
                break;
            case ZX_FDC_COMMAND_SENSE_DRIVE_STATUS:
                process_current_command = process_command_sense_drive_status;
                break;
            case ZX_FDC_COMMAND_RECALIBRATE:
                process_current_command = process_command_recalibrate;
                fdd0busy = 1;
                break;
            case ZX_FDC_COMMAND_SENSE_INTERRUPT_STATUS:
                process_current_command = process_command_sense_interrupt_status;
                mode = ZX_FDC_MODE_RESULT;
                call_count = 0;
                break;
            case ZX_FDC_COMMAND_SEEK:
                process_current_command = process_command_seek;
                fdd0busy = 1;
                seeking = true;
                break;
            case ZX_FDC_COMMAND_READ_ID:
                process_current_command = process_command_read_id;
                break;
            case ZX_FDC_COMMAND_READ_DATA:
                process_current_command = process_command_read_data;
                break;
            case ZX_FDC_COMMAND_READ_DELETED_DATA:
                process_current_command = process_command_read_data;
                break;
            default:
            {
                debug::printCommand();
                process_current_command = process_command_unknown;
                mode = ZX_FDC_MODE_RESULT;
                call_count = 0;
            }
        }
    }


    // ===================================================================
    // FDC COMMAND: UNKNOWN
    // ===================================================================
    uint8_t process_command_unknown(uint8_t command)
    {
        process_current_command = nullptr;
        mode = ZX_FDC_MODE_IDLE;
        const uint8_t val = ST0();
        //printf("--> (returning 0x%02x)\n", val);
        debug::composeResult(val);
        debug::printResult();
        return val;
    }
    
    // ===================================================================
    // FDC COMMAND: SPECIFY (0x03)
    // ===================================================================
    uint8_t process_command_specify(uint8_t command)
    {
        debug::composeCommand(command);
        switch (call_count) {
            case 1:
                srt = (command & 0xf0) >> 4;
                hut = command & 0x0f;
                call_count++;
                break;
            case 2:
                hlt = (command & 0xfe) >> 1;
                nd = command & 0x01;
                call_count=0;
                process_current_command = nullptr;
                mode = ZX_FDC_MODE_IDLE;
                debug::printCommand();
                break;
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: SENSE DRIVE STATUS (0x04)
    // ===================================================================
    uint8_t process_command_sense_drive_status(uint8_t command)
    {
        switch (mode)
        {
            case ZX_FDC_MODE_COMMAND:
                current_head = (command & 0x4)>>2;
                current_drive = command & 0x3;
                call_count = 0;
                mode = ZX_FDC_MODE_RESULT;
                debug::composeCommand(command);
                debug::printCommand();
                break;
            case ZX_FDC_MODE_RESULT:
            {
                process_current_command = nullptr;
                mode = ZX_FDC_MODE_IDLE;
                const uint8_t val = ST3();
                //printf("--> (returning 0x%02x)\n", val);
                debug::composeResult(val);
                debug::printResult();
                return val;
            }
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: RECALIBRATE (0x07)
    // ===================================================================
    uint8_t process_command_recalibrate(uint8_t command)
    {
        switch (mode)
        {
            case ZX_FDC_MODE_COMMAND:
                current_drive = command & 0x3;
                call_count = 0;
                current_track = 0;
                seeking = true;
                process_current_command = nullptr;
                mode = ZX_FDC_MODE_IDLE;
                debug::composeCommand(command);
                debug::printCommand();
                break;
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: SENSE INTERRUPT STATUS (0x08)
    // ===================================================================
    uint8_t process_command_sense_interrupt_status(uint8_t command)
    {
        switch (call_count)
        {
            case 0:
            {
                debug::printCommand();
                call_count++;
                const uint8_t val = ST0();
                //printf("--> (returning 0x%02x)\n", val);
                debug::composeResult(val);
                return val;
            }
            case 1:
                process_current_command = nullptr;
                fdd0busy = 0;
                seeking = false;
                call_count = 0;
                mode = ZX_FDC_MODE_IDLE;
                //printf("--> (returning 0x%02x)\n", current_track);
                debug::composeResult(current_track);
                debug::printResult();
                return current_track;
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: SEEK (0x0F)
    // ===================================================================
    uint8_t process_command_seek(uint8_t command)
    {
        debug::composeCommand(command);
        switch (call_count)
        {
            case 1:
            {
                current_head = (command & 0x4)>>2;
                current_drive = command & 0x3;
                call_count++;
                break;
            }
            case 2:
                current_track = command;
                process_current_command = nullptr;
                call_count = 0;
                mode = ZX_FDC_MODE_IDLE;
                debug::printCommand();
                break;
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: READ ID (0x0A)
    // ===================================================================
    uint8_t process_command_read_id(uint8_t command)
    {
        switch (mode)
        {
            case ZX_FDC_MODE_COMMAND:
                current_head = (command & 0x4)>>2;
                current_drive = command & 0x3;
                call_count = 0;
                current_sector = 0;
                mode = ZX_FDC_MODE_RESULT;
                debug::composeCommand(command);
                debug::printCommand();
                break;
            case ZX_FDC_MODE_RESULT:
            {
                switch (call_count)
                {
                    case 0:
                    {
                        call_count++;
                        fdd0busy = 1;
                        const uint8_t val = ST0();
                        fdd0busy = 0;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 1:
                    {
                        call_count++;
                        const uint8_t val = 0x00; // ST1
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 2:
                    {
                        call_count++;
                        const uint8_t val = 0x00; // ST2
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 3:
                    {
                        call_count++;
                        const uint8_t val = current_track;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 4:
                    {
                        call_count++;
                        const uint8_t val = current_head;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 5:
                    {
                        call_count++;
                        const uint8_t val = disk.tracks[current_track].sectors[current_sector].id;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 6:
                    {
                        call_count = 0;
                        process_current_command = nullptr;
                        mode = ZX_FDC_MODE_IDLE;
                        const uint8_t val = disk.tracks[current_track].sectors[current_sector].size;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        debug::printResult();
                        return val;
                    }
                }
            }
        }
        return 0;
    }

    // ===================================================================
    // FDC COMMAND: READ DATA (0x06)
    // ===================================================================
    uint8_t process_command_read_data(uint8_t command)
    {
        switch (mode)
        {
            case ZX_FDC_MODE_COMMAND:
                debug::composeCommand(command);
                switch (call_count)
                {
                    case 1:
                        call_count++;
                        current_head = (command & 0x4)>>2;
                        current_drive = command & 0x3;
                        break;
                    case 2:
                        call_count++;
                        current_track = command;
                        break;
                    case 3:
                        call_count++;
                        current_head = command;
                        break;
                    case 4:
                        call_count++;
                        current_sector = find_sector_id(command);
                        break;
                    case 5:
                        call_count++;
                        bytes_to_read = command;
                        break;
                    case 6:
                        call_count++;
                        eot = command;
                        break;
                    case 7:
                        call_count++;
                        break;
                    case 8:
                        if ( (bytes_to_read==0) && (command != 0xff) ) bytes_to_read = command;
                        bytes_to_read = bytes_to_read << 8;
                        current_byte = 0;
                        call_count = 0;
                        mode = ZX_FDC_MODE_EXECUTION;
                        debug::printCommand();
                        break;
                }
                break;

            case ZX_FDC_MODE_EXECUTION:
            {
                const uint8_t val = disk.tracks[current_track].sectors[current_sector].data[current_byte];

                current_byte++;
                if (current_byte == disk.tracks[current_track].sectors[current_sector].data_length) {
                    current_byte = 0;
                    current_sector = disk.tracks[current_track].sectors[current_sector].id;
                    current_sector++;
                    current_sector = find_sector_id(current_sector);
                    if (current_sector == 0xff) {
                        current_track++;
                        current_sector = 0;
                    }
                }

                bytes_to_read--;
                if (bytes_to_read==0) {
                    st1 = disk.tracks[current_track].sectors[current_sector].st1;
                    st2 = disk.tracks[current_track].sectors[current_sector].st2;
                    mode = ZX_FDC_MODE_RESULT;
                }

                return val;
                break;
            }

            case ZX_FDC_MODE_RESULT:
            {
                switch (call_count)
                {
                    case 0:
                    {
                        call_count++;
                        fdd0busy = 1;
                        const uint8_t val = ST0();
                        fdd0busy = 0;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 1:
                    {
                        call_count++;
                        const uint8_t val = st1;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 2:
                    {
                        call_count++;
                        const uint8_t val = st2;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 3:
                    {
                        call_count++;
                        const uint8_t val = current_track;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 4:
                    {
                        call_count++;
                        const uint8_t val = current_head;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 5:
                    {
                        call_count++;
                        const uint8_t val = disk.tracks[current_track].sectors[current_sector].id;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        return val;
                    }
                    case 6:
                    {
                        call_count = 0;
                        process_current_command = nullptr;
                        mode = ZX_FDC_MODE_IDLE;
                        const uint8_t val = disk.tracks[current_track].sectors[current_sector].size;
                        //printf("--> (returning 0x%02x)\n", val);
                        debug::composeResult(val);
                        debug::printResult();
                        return val;
                    }
                }
                break;
            }
        }
        return 0;
    }


    namespace debug
    {
        uint8_t values[9];
        uint8_t num_values = 0;

        void composeCommand(const uint8_t value)
        {
            values[num_values++] = value;
        }

        void composeResult(const uint8_t value)
        {
            values[num_values++] = value;
        }

        void printCommand()
        {
            printf("FDC COMMAND %02x: ( ", values[0] & 0x1f);
            for (int i=0; i<num_values; ++i) printf("%02x ", values[i]);
            printf(")\n");
            num_values = 0;
        }

        void printResult()
        {
            printf("FDC RESULT: ( ");
            for (int i=0; i<num_values; ++i) printf("%02x ", values[i]);
            printf(")\n");
            num_values = 0;
        }

    }

}