/*
 *  Copyright (C) 2026 The DOSBox-X Team
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <vector>

#include "dosbox.h"
#include "bios_disk.h"
#include "imagedisk_teledisk.h"

static inline uint16_t td0_read_u16le(const uint8_t *p) {
    return (uint16_t)p[0] | ((uint16_t)p[1] << 8U);
}

static bool td0_decode_method_raw(const std::vector<uint8_t> &encoded, const uint16_t sector_size, std::vector<uint8_t> &decoded) {
    if (encoded.size() != sector_size)
        return false;
    decoded = encoded;
    return true;
}

static bool td0_decode_method_repeated_2byte_pattern(const std::vector<uint8_t> &encoded, const uint16_t sector_size, std::vector<uint8_t> &decoded) {
    decoded.clear();
    decoded.reserve(sector_size);

    size_t pos = 0;
    while (decoded.size() < sector_size) {
        if ((pos + 4U) > encoded.size())
            return false;

        const uint16_t count = td0_read_u16le(&encoded[pos]); pos += 2;
        const uint8_t b0 = encoded[pos++];
        const uint8_t b1 = encoded[pos++];

        for (uint16_t i = 0; i < count && decoded.size() < sector_size; i++) {
            decoded.push_back(b0);
            if (decoded.size() >= sector_size)
                break;
            decoded.push_back(b1);
        }
    }

    return decoded.size() == sector_size;
}

static bool td0_decode_method_rle(const std::vector<uint8_t> &encoded, const uint16_t sector_size, std::vector<uint8_t> &decoded) {
    decoded.clear();
    decoded.reserve(sector_size);

    size_t pos = 0;
    while (decoded.size() < sector_size) {
        if (pos >= encoded.size())
            return false;

        const uint8_t token = encoded[pos++];
        if (token == 0) {
            if (pos >= encoded.size())
                return false;

            const uint8_t n = encoded[pos++];
            if ((pos + n) > encoded.size())
                return false;

            for (uint8_t i = 0; i < n && decoded.size() < sector_size; i++)
                decoded.push_back(encoded[pos++]);
        }
        else {
            const size_t block_len = (size_t)token * 2U;
            if (pos >= encoded.size())
                return false;
            const uint8_t repeat_count = encoded[pos++];
            if ((pos + block_len) > encoded.size())
                return false;

            std::vector<uint8_t> block(block_len);
            memcpy(block.data(), &encoded[pos], block_len);
            pos += block_len;

            for (uint8_t r = 0; r < repeat_count && decoded.size() < sector_size; r++) {
                for (size_t i = 0; i < block_len && decoded.size() < sector_size; i++)
                    decoded.push_back(block[i]);
            }
        }
    }

    return decoded.size() == sector_size;
}

static bool td0_decode_sector_data(imageDiskTeledisk::sector_encoding_method method, const std::vector<uint8_t> &encoded, const uint16_t sector_size, std::vector<uint8_t> &decoded) {
    switch (method) {
    case imageDiskTeledisk::sector_encoding_method::raw:
        return td0_decode_method_raw(encoded, sector_size, decoded);
    case imageDiskTeledisk::sector_encoding_method::repeated_2byte_pattern:
        return td0_decode_method_repeated_2byte_pattern(encoded, sector_size, decoded);
    case imageDiskTeledisk::sector_encoding_method::rle:
        return td0_decode_method_rle(encoded, sector_size, decoded);
    default:
        return false;
    }
}

imageDiskTeledisk::td0entry *imageDiskTeledisk::findSector(uint8_t head,uint8_t track,uint8_t sector,unsigned int req_sector_size) {
    /* Delegate to the const overload; safe to cast away const because this object is non-const here. */
    const imageDiskTeledisk *self = this;
    return const_cast<td0entry *>(self->findSector(head, track, sector, req_sector_size));
}

const imageDiskTeledisk::td0entry *imageDiskTeledisk::findSector(uint8_t head,uint8_t track,uint8_t sector,unsigned int req_sector_size) const {
    const td0entry *best = nullptr;

    for (const auto &entry : dents) {
        if (entry.phys_head != head || entry.phys_track != track || entry.phys_sector != sector)
            continue;
        if (req_sector_size != ~0U && req_sector_size != 0 && entry.sector_size != req_sector_size)
            continue;
        if (!best || (!best->has_data && entry.has_data))
            best = &entry;
    }

    return best;
}

Int13Status imageDiskTeledisk::Read_Sector(uint32_t head,uint32_t cylinder,uint32_t sector,void * data,unsigned int req_sector_size) {
    const auto* ent = findSector((uint8_t)head, (uint8_t)cylinder, (uint8_t)sector, req_sector_size);

    if (!ent || !ent->has_data)
        return Int13Status::SectorNotFound;

    if (req_sector_size == 0) req_sector_size = ent->getSectorSize();

    if (ent->getSectorSize() != req_sector_size)
        return Int13Status::SectorNotFound;

    /* Copy the data regardless, mirroring real hardware which still transfers the
     * (possibly corrupt) sector before reporting the flagged condition. */
    memcpy(data, ent->data.data(), req_sector_size);

    /* Surface sector flags recorded in the .td0 image as INT 13h status. */
    if (ent->flags & td0_sector_flags::crc_error)
        return Int13Status::DataError;             /* recorded with a CRC error */
    if (ent->flags & td0_sector_flags::deleted_mark)
        return Int13Status::AddressMarkNotFound;   /* deleted-data address mark */

    return Int13Status::NoError;
}

Int13Status imageDiskTeledisk::Write_Sector(uint32_t head,uint32_t cylinder,uint32_t sector,const void * data,unsigned int req_sector_size) {
    (void)head; (void)cylinder; (void)sector; (void)data; (void)req_sector_size;
    /* TeleDisk images are read-only: writes cannot be re-encoded back into the .td0 archive,
     * so accepting a write in memory only would silently discard the guest's data on remount.
     * Report the disk as write-protected instead. */
    return Int13Status::WriteProtected;
}

Int13Status imageDiskTeledisk::Read_AbsoluteSector(uint32_t sectnum, void * data) {
    unsigned int c,h,s;

    if (sectors == 0 || heads == 0)
        return Int13Status::SectorNotFound;

    s = (sectnum % sectors) + 1;
    h = (sectnum / sectors) % heads;
    c = (sectnum / sectors / heads);
    return Read_Sector(h,c,s,data);
}

Int13Status imageDiskTeledisk::Write_AbsoluteSector(uint32_t sectnum,const void *data) {
    unsigned int c,h,s;

    if (sectors == 0 || heads == 0)
        return Int13Status::SectorNotFound;

    s = (sectnum % sectors) + 1;
    h = (sectnum / sectors) % heads;
    c = (sectnum / sectors / heads);
    return Write_Sector(h,c,s,data);
}

imageDiskTeledisk::imageDiskTeledisk(FILE *imgFile, const char *imgName, uint32_t imgSizeK, bool isHardDisk) : imageDisk(ID_TELEDISK) {
    (void)isHardDisk;//UNUSED
    (void)imgSizeK;//UNUSED
    td0_image_header header;

    heads = 1;
    cylinders = 0;
    image_base = 0;
    sectors = 0;
    active = false;
    sector_size = 0;
    reserved_cylinders = 0;
    diskSizeK = 0;
    diskimg = imgFile;

    if (imgName)
        diskname = imgName;

    if (!diskimg)
        return;

    fseek(diskimg,0,SEEK_SET);
    if (fread(&header,1,sizeof(header),diskimg) != sizeof(header))
        return;

    /* Only support "TD" (normal compression), not "td" advanced compression */
    if (header.signature[0] != 'T' || header.signature[1] != 'D')
        return;

    /* Optional comment block, flagged by the high bit of the stepping field */
    if (header.stepping & 0x80U) {
        td0_comment_header comment;
        if (fread(&comment,1,sizeof(comment),diskimg) != sizeof(comment))
            return;

        const uint16_t comment_len = td0_read_u16le(comment.length);
        if (fseek(diskimg, (long)comment_len, SEEK_CUR) != 0)
            return;
    }

    bool parse_ok = false;
    for (;;) {
        td0_track_header track_header;
        if (fread(&track_header,1,sizeof(track_header),diskimg) != sizeof(track_header))
            return;

        const uint8_t num_sectors = track_header.num_sectors;
        if (num_sectors == 0xFFU) {
            parse_ok = true;
            break;
        }

        const uint8_t phys_track = track_header.cylinder;
        const uint8_t phys_head = track_header.head & 1U;

        for (uint8_t i = 0; i < num_sectors; i++) {
            td0_sector_header sector_header;
            if (fread(&sector_header,1,sizeof(sector_header),diskimg) != sizeof(sector_header))
                return;

            td0entry ent;
            const uint8_t size_code = (sector_header.size_code > 6U) ? 0U : sector_header.size_code;

            ent.phys_track = phys_track;
            ent.phys_head = phys_head;
            ent.phys_sector = sector_header.sector;
            ent.logical_track = sector_header.cylinder;
            ent.logical_head = sector_header.head;
            ent.sector_size = (uint16_t)(128U << size_code);
            ent.flags = sector_header.flags;

//            LOG_MSG("TD1: C/H/S=%u/%u/%u size=%u flags=0x%02X",
//                (unsigned int)ent.phys_track,
//                (unsigned int)ent.phys_head,
//                (unsigned int)ent.phys_sector,
//                (unsigned int)ent.sector_size,
//                (unsigned int)ent.flags
//            );
            if ((ent.flags & td0_sector_flags::id_no_data) || (sector_header.size_code & 0xF8U)) {
                ent.has_data = false;
                ent.data.assign(ent.sector_size, 0);
            }
            else if (ent.flags & td0_sector_flags::dos_skipped) {
                ent.has_data = true;
                ent.data.assign(ent.sector_size, 0);
            }
            else {
                td0_data_header data_header;
                const long data_header_offset = ftell(diskimg);
                if (fread(&data_header,1,sizeof(data_header),diskimg) != sizeof(data_header))
                    return;

                const uint16_t data_block_size = td0_read_u16le(data_header.block_size);
                if (data_block_size < 1U)
                    return;

                const auto encoded_size = (uint16_t)(data_block_size - 1U);
                std::vector<uint8_t> encoded(encoded_size);
                if (encoded_size != 0 && fread(encoded.data(),1,encoded_size,diskimg) != encoded_size)
                    return;

                const auto method =
                    static_cast<imageDiskTeledisk::sector_encoding_method>(data_header.method);

                ent.has_data = td0_decode_sector_data(method, encoded, ent.sector_size, ent.data);
                if (!ent.has_data) {
                    LOG_MSG(
                        "TD0: failed to decode sector C/H/S=%u/%u/%u method=%u file_ofs=%ld",
                        (unsigned int)ent.phys_track,
                        (unsigned int)ent.phys_head,
                        (unsigned int)ent.phys_sector,
                        (unsigned int)method,
                        data_header_offset
                    );
                    ent.data.assign(ent.sector_size, 0);
                }
            }

            dents.push_back(ent);
            diskSizeK += ent.sector_size;
        }
    }

    diskSizeK /= 1024;

    if (parse_ok && !dents.empty()) {
        uint16_t max_cylinder = 0;
        uint16_t max_head = 0;

        for (const auto &entry : dents) {
            max_cylinder = std::max<int>(entry.phys_track, max_cylinder);

            if (entry.phys_track != 0)
                continue;

            max_head = std::max<int>(entry.phys_head, max_head);
            if (entry.phys_head == 0) {
                sectors++;
                if (sector_size == 0)
                    sector_size = entry.getSectorSize();
            }
        }

        cylinders = max_cylinder + 1;
        heads = max_head + 1;
        const bool founddisk = sector_size != 0 && sectors != 0 &&
                               cylinders != 0 && heads != 0;

        LOG_MSG(
            "TD0 geometry detection: C/H/S %u/%u/%u %u bytes/sector entries=%u",
            cylinders, heads, sectors, sector_size, (unsigned int)dents.size()
        );

        /* active stays false (as initialized) unless we resolved a usable geometry */
        if (founddisk) {
            active = true;
            incrementFDD();
            updateFloppyDPT();
        }
    }
}

imageDiskTeledisk::~imageDiskTeledisk() {
    if (diskimg) {
        fclose(diskimg);
        diskimg = nullptr;
    }
}
