//! Host-side helper for build.rs: copy a directory tree into a freshly
//! formatted exFAT image. `mkfs.exfat` (exfatprogs) can format but not
//! populate, mtools doesn't speak exFAT, and mounting needs root or a loop
//! device. So this is a tiny exFAT *writer* of its own.
//!
//! It is deliberately separate from the kernel's driver (src/fs/exfat.rs)
//! and shares no code with it, so a misunderstanding of the spec in one is
//! not silently copied into the other. The host's `fsck.exfat` checks what
//! this writes (scripts/exfat-check.sh).
//!
//! It only needs to handle a *fresh* volume, which keeps it simple:
//! * Free clusters are handed out in order after the last used one.
//! * Files and directories get one contiguous run of clusters and the
//!   NoFatChain flag (their FAT entries stay unused), except files whose
//!   name starts with `frag-`: they get every *other* cluster, linked
//!   through the FAT, so the kernel has a fragmented file to read.
//! * Everything gets the same timestamp (1 Jan 2026), like the kernel,
//!   which has no clock yet.
//!
//! Layout reminders are in the kernel driver's comments.

use std::{fs, path::Path};

const ATTR_DIRECTORY: u16 = 0x10;
const ATTR_ARCHIVE: u16 = 0x20;
const FLAG_ALLOCATION_POSSIBLE: u8 = 0x01;
const FLAG_NO_FAT_CHAIN: u8 = 0x02;
/// 2026-01-01 00:00:00 in exFAT's DOS-style format (year since 1980 in the
/// top 7 bits, then month, day, hour, minute, seconds / 2).
const TIMESTAMP: u32 = (46 << 25) | (1 << 21) | (1 << 16);

fn get32(b: &[u8], at: usize) -> u32 {
    u32::from_le_bytes(b[at..at + 4].try_into().unwrap())
}
fn get64(b: &[u8], at: usize) -> u64 {
    u64::from_le_bytes(b[at..at + 8].try_into().unwrap())
}

struct Image {
    data: Vec<u8>,
    cluster: usize,  // bytes per cluster
    fat: usize,      // byte offset of the FAT
    heap: usize,     // byte offset of cluster 2
    count: u32,      // clusters in the heap
    bitmap: usize,   // byte offset of the allocation bitmap (contiguous on a fresh volume)
    upcase: Vec<u16>, // upcase[c] = upper-case of UTF-16 unit c (identity past the end)
    next: u32,       // next cluster to hand out
}

impl Image {
    fn pos(&self, cluster: u32) -> usize {
        self.heap + (cluster as usize - 2) * self.cluster
    }

    fn set_fat(&mut self, cluster: u32, value: u32) {
        let at = self.fat + 4 * cluster as usize;
        self.data[at..at + 4].copy_from_slice(&value.to_le_bytes());
    }

    fn used(&self, cluster: u32) -> bool {
        let i = (cluster - 2) as usize;
        self.data[self.bitmap + i / 8] & (1 << (i % 8)) != 0
    }

    /// Take `n` clusters: consecutive ones, or every other one (`spread`).
    fn alloc(&mut self, n: usize, spread: bool) -> Vec<u32> {
        let mut out = Vec::new();
        for _ in 0..n {
            let c = self.next;
            assert!(c - 2 < self.count, "exFAT image is full");
            let i = (c - 2) as usize;
            self.data[self.bitmap + i / 8] |= 1 << (i % 8);
            out.push(c);
            self.next += if spread { 2 } else { 1 };
        }
        out
    }

    /// Write `bytes` into `clusters` (zero-padding the last one).
    fn fill(&mut self, clusters: &[u32], bytes: &[u8]) {
        for (i, &c) in clusters.iter().enumerate() {
            let at = self.pos(c);
            let chunk = &bytes[(i * self.cluster).min(bytes.len())..((i + 1) * self.cluster).min(bytes.len())];
            self.data[at..at + self.cluster].fill(0);
            self.data[at..at + chunk.len()].copy_from_slice(chunk);
        }
    }

    /// The File + Stream Extension + File Name entries for one file.
    fn entry_set(&self, name: &str, attrs: u16, first: u32, size: u64, no_fat_chain: bool) -> Vec<u8> {
        let units: Vec<u16> = name.encode_utf16().collect();
        assert!(units.len() <= 255, "name too long: {name}");
        let names = units.len().div_ceil(15);
        let mut set = vec![0u8; 32 * (2 + names)];
        // File entry.
        set[0] = 0x85;
        set[1] = (1 + names) as u8; // SecondaryCount
        set[4..6].copy_from_slice(&attrs.to_le_bytes());
        for at in [8, 12, 16] {
            set[at..at + 4].copy_from_slice(&TIMESTAMP.to_le_bytes()); // created, modified, accessed
        }
        // Stream Extension entry.
        let s = &mut set[32..64];
        s[0] = 0xC0;
        s[1] = FLAG_ALLOCATION_POSSIBLE | if no_fat_chain && first != 0 { FLAG_NO_FAT_CHAIN } else { 0 };
        s[3] = units.len() as u8;
        let mut hash: u16 = 0;
        for &u in &units {
            let up = self.upcase.get(u as usize).copied().unwrap_or(u);
            for byte in up.to_le_bytes() {
                hash = hash.rotate_right(1).wrapping_add(byte as u16);
            }
        }
        s[4..6].copy_from_slice(&hash.to_le_bytes());
        s[8..16].copy_from_slice(&size.to_le_bytes()); // ValidDataLength
        s[20..24].copy_from_slice(&first.to_le_bytes());
        s[24..32].copy_from_slice(&size.to_le_bytes()); // DataLength
        // File Name entries, 15 UTF-16 units each.
        for (i, chunk) in units.chunks(15).enumerate() {
            let e = &mut set[64 + 32 * i..96 + 32 * i];
            e[0] = 0xC1;
            for (j, u) in chunk.iter().enumerate() {
                e[2 + 2 * j..4 + 2 * j].copy_from_slice(&u.to_le_bytes());
            }
        }
        // SetChecksum: every byte of the set except the checksum itself.
        let mut sum: u16 = 0;
        for (i, &b) in set.iter().enumerate() {
            if i != 2 && i != 3 {
                sum = sum.rotate_right(1).wrapping_add(b as u16);
            }
        }
        set[2..4].copy_from_slice(&sum.to_le_bytes());
        set
    }

    /// Write everything in `dir` to the image; returns its directory entries.
    fn write_tree(&mut self, dir: &Path) -> Vec<u8> {
        let mut children: Vec<_> = fs::read_dir(dir).unwrap().map(|e| e.unwrap().path()).collect();
        children.sort();
        let mut entries = Vec::new();
        for path in children {
            let name = path.file_name().unwrap().to_str().unwrap().to_string();
            if path.is_dir() {
                let sub = self.write_tree(&path);
                // At least one cluster; a directory's size is always whole clusters.
                let n = sub.len().div_ceil(self.cluster).max(1);
                let clusters = self.alloc(n, false);
                self.fill(&clusters, &sub);
                entries.extend(self.entry_set(&name, ATTR_DIRECTORY, clusters[0], (n * self.cluster) as u64, true));
            } else {
                let bytes = fs::read(&path).unwrap();
                let n = bytes.len().div_ceil(self.cluster);
                let spread = name.starts_with("frag-");
                let clusters = self.alloc(n, spread);
                self.fill(&clusters, &bytes);
                if spread {
                    for w in clusters.windows(2) {
                        self.set_fat(w[0], w[1]);
                    }
                    if let Some(&last) = clusters.last() {
                        self.set_fat(last, 0xFFFF_FFFF); // end of chain
                    }
                }
                let first = clusters.first().copied().unwrap_or(0);
                entries.extend(self.entry_set(&name, ATTR_ARCHIVE, first, bytes.len() as u64, !spread));
            }
        }
        entries
    }
}

/// Copy `tree` into the exFAT file system in `image` (made by mkfs.exfat).
pub fn populate(image: &Path, tree: &Path) {
    let data = fs::read(image).unwrap();
    assert_eq!(&data[3..11], b"EXFAT   ", "not an exFAT image");
    let sector = 1usize << data[108];
    let cluster = sector << data[109];
    let fat = get32(&data, 80) as usize * sector;
    let heap = get32(&data, 88) as usize * sector;
    let count = get32(&data, 92);
    let root = get32(&data, 96);
    let mut img = Image { data, cluster, fat, heap, count, bitmap: 0, upcase: Vec::new(), next: 0 };

    // The root directory (one cluster on a fresh volume) holds the
    // allocation bitmap (0x81) and up-case table (0x82) entries.
    let root_at = img.pos(root);
    let mut free_slot = None;
    for i in 0..cluster / 32 {
        let e = &img.data[root_at + 32 * i..root_at + 32 * i + 32];
        match e[0] {
            0x81 => img.bitmap = img.pos(get32(e, 20)),
            0x82 => {
                let (at, len) = (img.pos(get32(e, 20)), get64(e, 24) as usize);
                // Compressed table: 0xFFFF, n = "the next n characters map
                // to themselves".
                let raw: Vec<u16> = img.data[at..at + len].chunks(2).map(|c| u16::from_le_bytes([c[0], c[1]])).collect();
                let mut i = 0;
                while i < raw.len() {
                    if raw[i] == 0xFFFF && i + 1 < raw.len() {
                        for _ in 0..raw[i + 1] {
                            let c = img.upcase.len() as u16;
                            img.upcase.push(c);
                        }
                        i += 2;
                    } else {
                        img.upcase.push(raw[i]);
                        i += 1;
                    }
                }
            }
            0x00 => {
                free_slot = Some(i);
                break;
            }
            _ => {}
        }
    }
    assert!(img.bitmap != 0 && !img.upcase.is_empty(), "mkfs.exfat left no bitmap or up-case table?");
    // Hand out clusters after the last one mkfs used.
    img.next = (2..count + 2).rev().find(|&c| img.used(c)).unwrap() + 1;

    let entries = img.write_tree(tree);
    let slot = free_slot.expect("root directory full");
    assert!(32 * slot + entries.len() <= cluster, "too many files for one root cluster");
    img.data[root_at + 32 * slot..root_at + 32 * slot + entries.len()].copy_from_slice(&entries);

    // PercentInUse (boot sector byte 112). It isn't covered by the boot
    // checksum, so changing it doesn't invalidate the boot region.
    let used = (2..count + 2).filter(|&c| img.used(c)).count();
    img.data[112] = (used * 100 / count as usize) as u8;
    fs::write(image, &img.data).unwrap();
}