//! Complete FDS disk images, including the internal FAT/EROFS/ext4 disk.
//! Cartridge validation remains shared with the cartridge writer. Internal
//! installation is deliberately separate from its reserved-partition policy.
use fds_burn::image::{
    self, LINUX_TYPE, Partition, TABLE_BYTES, crc32, put32, put64, u32le, u64le,
};
use fds_common::{Error, Result};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::{
    fs::File,
    os::{
        fd::AsRawFd,
        unix::fs::{FileExt, FileTypeExt},
    },
};

const EFI_TYPE: [u8; 16] = [
    0x28, 0x73, 0x2a, 0xc1, 0x1f, 0xf8, 0xd2, 0x11, 0xba, 0x4b, 0, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b,
];
fn bad(message: &str) -> Error {
    Error(format!("Invalid FDS disk image: {message}"))
}

#[derive(Debug, PartialEq, Eq, Serialize)]
pub struct Geometry {
    pub bytes: u64,
    pub kind: String,
    pub disk_uuid: String,
    pub partitions: Vec<Partition>,
    #[serde(skip)]
    head: Vec<u8>,
    #[serde(skip)]
    tail: Vec<u8>,
}

pub fn inspect(file: &File, bytes: u64) -> Result<Geometry> {
    if bytes % 512 != 0 || bytes < (2048 + 2048 + 33) * 512 {
        return Err(bad(
            "use a complete, sector-aligned GPT .img, not a compressed archive or partition payload",
        ));
    }
    let mut head = vec![0; 1024 + TABLE_BYTES];
    let mut tail = vec![0; 512 + TABLE_BYTES];
    file.read_exact_at(&mut head, 0)?;
    let tail_offset = bytes - tail.len() as u64;
    file.read_exact_at(&mut tail, tail_offset)?;
    if let Ok(cartridge) = image::inspect(file, bytes) {
        return Ok(Geometry {
            bytes,
            kind: format!("{:?}", cartridge.class).to_lowercase(),
            disk_uuid: cartridge.disk_uuid,
            partitions: cartridge.partitions,
            head,
            tail,
        });
    }
    // The only additional accepted layout is the complete internal disk.
    // Both GPT copies, CRCs, bounds, types, names and filesystem signatures
    // must agree before a destination can be opened for writing.
    let sectors = bytes / 512;
    if head[510..512] != [0x55, 0xaa]
        || head[446] != 0
        || head[450] != 0xee
        || u32le(&head, 454) != 1
        || u32le(&head, 458) != (sectors - 1).min(u32::MAX as u64) as u32
        || head[462..510].iter().any(|b| *b != 0)
    {
        return Err(bad(
            "missing protective MBR; expected a complete FDS GPT image",
        ));
    }
    let primary = &head[512..1024];
    let backup = &tail[TABLE_BYTES..];
    for (header, current, alternate, table) in [
        (primary, 1, sectors - 1, 2),
        (backup, sectors - 1, 1, sectors - 33),
    ] {
        let mut checked = header.to_vec();
        put32(&mut checked, 16, 0);
        if &header[..8] != b"EFI PART"
            || u32le(header, 8) != 0x10000
            || u32le(header, 12) != 92
            || u32le(header, 20) != 0
            || header[92..].iter().any(|b| *b != 0)
            || crc32(&checked[..92]) != u32le(header, 16)
            || u64le(header, 24) != current
            || u64le(header, 32) != alternate
            || u64le(header, 40) != 34
            || u64le(header, 48) != sectors - 34
            || u64le(header, 72) != table
            || u32le(header, 80) != 128
            || u32le(header, 84) != 128
            || header[56..72].iter().all(|b| *b == 0)
        {
            return Err(bad("GPT geometry or header checksum is invalid"));
        }
    }
    let table = &head[1024..];
    if primary[40..72] != backup[40..72]
        || primary[80..92] != backup[80..92]
        || table != &tail[..TABLE_BYTES]
        || crc32(table) != u32le(primary, 88)
        || table[3 * 128..].iter().any(|b| *b != 0)
    {
        return Err(bad(
            "internal disk GPT copies disagree or do not contain exactly three partitions",
        ));
    }
    let mut partitions = Vec::new();
    let mut ids = std::collections::BTreeSet::new();
    let mut next = 2048;
    for (index, name) in ["FDS_BOOT", "FDS_RECOVERY", "FDS_INTERNAL"]
        .iter()
        .enumerate()
    {
        let entry = &table[index * 128..(index + 1) * 128];
        let first = u64le(entry, 32);
        let last = u64le(entry, 40);
        let mut encoded = [0; 72];
        for (n, unit) in name.encode_utf16().enumerate() {
            encoded[n * 2..n * 2 + 2].copy_from_slice(&unit.to_le_bytes());
        }
        if entry[..16] != if index == 0 { EFI_TYPE } else { LINUX_TYPE }
            || entry[16..32].iter().all(|b| *b == 0)
            || !ids.insert(entry[16..32].to_vec())
            || u64le(entry, 48) != 0
            || entry[56..] != encoded
            || first < next
            || first % 2048 != 0
            || (index == 0 && first != 2048)
            || first > sectors - 34
            || last < first
            || last > sectors - 34
            || (last - first + 1) % 2048 != 0
        {
            return Err(bad(
                "unsupported internal partition type, name, UUID or bounds",
            ));
        }
        let start = first * 512;
        let mut signature = [0; 2048];
        file.read_exact_at(&mut signature, start)?;
        let valid = match index {
            0 => {
                signature[510..512] == [0x55, 0xaa]
                    && &signature[82..90] == b"FAT32   "
                    && signature[11..13] == [0, 2]
            }
            1 => signature[1024..1028] == [0xe2, 0xe1, 0xf5, 0xe0],
            _ => signature[1080..1082] == [0x53, 0xef],
        };
        if !valid {
            return Err(bad("internal disk requires FAT32, EROFS and ext4 in order"));
        }
        partitions.push(Partition {
            number: (index + 1) as u8,
            name: (*name).into(),
            start,
            bytes: (last - first + 1) * 512,
        });
        next = last + 1;
    }
    Ok(Geometry {
        bytes,
        kind: "internal".into(),
        disk_uuid: image::hex(&primary[56..72]),
        partitions,
        head,
        tail,
    })
}

fn overlay(buffer: &mut [u8], offset: u64, patch: &[u8], position: u64) {
    let begin = offset.max(position);
    let end = (offset + buffer.len() as u64).min(position + patch.len() as u64);
    if begin < end {
        buffer[(begin - offset) as usize..(end - offset) as usize]
            .copy_from_slice(&patch[(begin - position) as usize..(end - position) as usize]);
    }
}

/// Verify all written bytes, including the relocated backup GPT. Unallocated
/// space outside the image is not erased, and no filesystem is expanded.
pub fn transfer(
    source: &File,
    target: &File,
    geometry: &Geometry,
    sha256: &str,
    target_bytes: u64,
    mut progress: impl FnMut(&str, u64) -> Result<()>,
) -> Result<()> {
    if target_bytes < geometry.bytes || target_bytes % 512 != 0 {
        return Err(bad("target is too small or not sector aligned"));
    }
    if source.metadata()?.len() != geometry.bytes
        || image::digest(source, geometry.bytes, |n| progress("checking", n))? != sha256
        || inspect(source, geometry.bytes)? != *geometry
    {
        return Err(Error(
            "Image changed after preview; target untouched".into(),
        ));
    }
    let mut head = geometry.head.clone();
    let mut tail = geometry.tail.clone();
    let sectors = target_bytes / 512;
    put32(&mut head, 458, (sectors - 1).min(u32::MAX as u64) as u32);
    for (header, lba, other, entries) in [
        (&mut head[512..1024], 1, sectors - 1, 2),
        (&mut tail[TABLE_BYTES..], sectors - 1, 1, sectors - 33),
    ] {
        put64(header, 24, lba);
        put64(header, 32, other);
        put64(header, 48, sectors - 34);
        put64(header, 72, entries);
        put32(header, 16, 0);
        let crc = crc32(&header[..92]);
        put32(header, 16, crc);
    }
    let old_tail = vec![0; tail.len()];
    let patch = |buffer: &mut [u8], offset| {
        overlay(buffer, offset, &head, 0);
        if target_bytes != geometry.bytes {
            overlay(
                buffer,
                offset,
                &old_tail,
                geometry.bytes - old_tail.len() as u64,
            );
        }
        overlay(buffer, offset, &tail, target_bytes - tail.len() as u64);
    };
    let mut original = vec![0; 1024 * 1024];
    let mut written = vec![0; original.len()];
    for phase in ["writing", "verifying"] {
        progress(phase, 0)?;
        let mut hash = Sha256::new();
        let mut offset = 0;
        while offset < geometry.bytes {
            let n = original.len().min((geometry.bytes - offset) as usize);
            source.read_exact_at(&mut original[..n], offset)?;
            hash.update(&original[..n]);
            patch(&mut original[..n], offset);
            if phase == "writing" {
                target.write_all_at(&original[..n], offset)?;
            } else {
                target.read_exact_at(&mut written[..n], offset)?;
                if written[..n] != original[..n] {
                    return Err(Error("Disk readback mismatch; flash failed".into()));
                }
            }
            offset += n as u64;
            if offset % (64 * 1024 * 1024) == 0 || offset == geometry.bytes {
                progress(phase, offset)?;
            }
        }
        if image::hex(&hash.finalize()) != sha256 || source.metadata()?.len() != geometry.bytes {
            return Err(Error(
                "Source changed during transfer; target is incomplete".into(),
            ));
        }
        if phase == "writing" {
            target.write_all_at(&tail, target_bytes - tail.len() as u64)?;
            target.sync_all()?;
            if target.metadata()?.file_type().is_block_device() {
                if unsafe { libc::ioctl(target.as_raw_fd(), 0x1261 as libc::Ioctl) } < 0 {
                    return Err(std::io::Error::last_os_error().into());
                }
            } else {
                let rc = unsafe {
                    libc::posix_fadvise(target.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED)
                };
                if rc != 0 {
                    return Err(std::io::Error::from_raw_os_error(rc).into());
                }
            }
        }
    }
    let mut actual_tail = vec![0; tail.len()];
    target.read_exact_at(&mut actual_tail, target_bytes - tail.len() as u64)?;
    let observed = inspect(target, target_bytes)?;
    if actual_tail != tail
        || observed.partitions != geometry.partitions
        || observed.disk_uuid != geometry.disk_uuid
        || observed.kind != geometry.kind
    {
        return Err(Error(
            "Written GPT or backup metadata failed verification".into(),
        ));
    }
    target.sync_all()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use fds_common::manifest::Class;
    use std::os::fd::FromRawFd;

    fn memory() -> File {
        let fd = unsafe { libc::memfd_create(c"fds-flash-test".as_ptr(), libc::MFD_CLOEXEC) };
        assert!(fd >= 0);
        unsafe { File::from_raw_fd(fd) }
    }
    fn fixture() -> (File, Geometry, String) {
        let source = memory();
        let layout =
            image::Layout::new(Class::System, 1024 * 1024, None, [1; 16], [2; 16]).unwrap();
        layout.write(&source).unwrap();
        source
            .write_all_at(&[0xe2, 0xe1, 0xf5, 0xe0], 1024 * 1024 + 1024)
            .unwrap();
        let geometry = inspect(&source, layout.bytes).unwrap();
        let hash = image::digest(&source, layout.bytes, |_| Ok(())).unwrap();
        (source, geometry, hash)
    }
    #[test]
    fn exact_larger_and_overlapping_backup_locations() {
        let (source, geometry, hash) = fixture();
        for extra in [0, 512, 4 * 1024 * 1024] {
            let target = memory();
            let size = geometry.bytes + extra;
            target.set_len(size).unwrap();
            transfer(&source, &target, &geometry, &hash, size, |_, _| Ok(())).unwrap();
            assert_eq!(
                inspect(&target, size).unwrap().partitions,
                geometry.partitions
            );
        }
        assert_eq!(
            image::digest(&source, geometry.bytes, |_| Ok(())).unwrap(),
            hash
        );
    }
    #[test]
    fn changed_source_is_rejected_before_any_write() {
        let (source, geometry, hash) = fixture();
        let target = memory();
        target.set_len(geometry.bytes).unwrap();
        let before = image::digest(&target, geometry.bytes, |_| Ok(())).unwrap();
        source.write_all_at(b"changed", 1024 * 1024 + 8192).unwrap();
        assert!(
            transfer(
                &source,
                &target,
                &geometry,
                &hash,
                geometry.bytes,
                |_, _| Ok(())
            )
            .is_err()
        );
        assert_eq!(
            image::digest(&target, geometry.bytes, |_| Ok(())).unwrap(),
            before
        );
    }
    #[test]
    fn readback_corruption_and_io_failure_never_succeed() {
        let (source, geometry, hash) = fixture();
        let target = memory();
        target.set_len(geometry.bytes).unwrap();
        let failure = transfer(
            &source,
            &target,
            &geometry,
            &hash,
            geometry.bytes,
            |phase, n| {
                if phase == "verifying" && n == 0 {
                    target.write_all_at(b"corrupt", 1024 * 1024 + 8192)?;
                }
                Ok(())
            },
        )
        .unwrap_err();
        assert!(failure.to_string().contains("readback mismatch"));
        let readonly = File::open(format!("/proc/self/fd/{}", target.as_raw_fd())).unwrap();
        assert!(
            transfer(
                &source,
                &readonly,
                &geometry,
                &hash,
                geometry.bytes,
                |_, _| Ok(())
            )
            .is_err()
        );
        let interrupted = transfer(
            &source,
            &target,
            &geometry,
            &hash,
            geometry.bytes,
            |phase, n| {
                if phase == "writing" && n > 0 {
                    return Err(Error("Target removed".into()));
                }
                Ok(())
            },
        )
        .unwrap_err();
        assert!(interrupted.to_string().contains("removed"));
    }
}