// USB mass-storage read/write validation check.
// Verifies that usbscsid-backed block devices support read and write I/O.

use std::process;

const PROGRAM: &str = "redbear-usb-storage-check";
const USAGE: &str = "Usage: redbear-usb-storage-check [--json]\n\n\
     USB storage read/write check. Discovers USB-backed block devices,\n\
     writes a test pattern to a safe sector, reads it back, and verifies.";

#[cfg(target_os = "redox")]
use std::fs;

const TEST_SECTOR: u64 = 2048;
const SECTOR_SIZE: usize = 512;
const TEST_PATTERN: &[u8; 24] = b"REDBEAR-USB-RW-PROOF\0\0\0\0";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CheckResult {
    Pass,
    Fail,
    Skip,
}

impl CheckResult {
    fn label(self) -> &'static str {
        match self {
            Self::Pass => "PASS",
            Self::Fail => "FAIL",
            Self::Skip => "SKIP",
        }
    }
}

struct Check {
    name: String,
    result: CheckResult,
    detail: String,
}

impl Check {
    fn pass(name: &str, detail: &str) -> Self {
        Check {
            name: name.to_string(),
            result: CheckResult::Pass,
            detail: detail.to_string(),
        }
    }
    fn fail(name: &str, detail: &str) -> Self {
        Check {
            name: name.to_string(),
            result: CheckResult::Fail,
            detail: detail.to_string(),
        }
    }
    fn skip(name: &str, detail: &str) -> Self {
        Check {
            name: name.to_string(),
            result: CheckResult::Skip,
            detail: detail.to_string(),
        }
    }
}

struct Report {
    checks: Vec<Check>,
    json_mode: bool,
}

impl Report {
    fn new(json_mode: bool) -> Self {
        Report {
            checks: Vec::new(),
            json_mode,
        }
    }
    fn add(&mut self, check: Check) {
        self.checks.push(check);
    }
    fn any_failed(&self) -> bool {
        self.checks.iter().any(|c| c.result == CheckResult::Fail)
    }
    fn print(&self) {
        if self.json_mode {
            self.print_json();
        } else {
            self.print_human();
        }
    }
    fn print_human(&self) {
        for check in &self.checks {
            let icon = match check.result {
                CheckResult::Pass => "[PASS]",
                CheckResult::Fail => "[FAIL]",
                CheckResult::Skip => "[SKIP]",
            };
            println!("{icon} {}: {}", check.name, check.detail);
        }
    }
    fn print_json(&self) {
        #[derive(serde::Serialize)]
        struct JsonCheck {
            name: String,
            result: String,
            detail: String,
        }
        #[derive(serde::Serialize)]
        struct JsonReport {
            device_path: String,
            sector_written: bool,
            sector_verified: bool,
            sector_restored: bool,
            checks: Vec<JsonCheck>,
        }
        let dev = self
            .checks
            .iter()
            .find(|c| c.name == "STORAGE_DISCOVERY")
            .map_or("none".to_string(), |c| c.detail.clone());
        let written = self
            .checks
            .iter()
            .any(|c| c.name == "STORAGE_WRITE" && c.result == CheckResult::Pass);
        let verified = self
            .checks
            .iter()
            .any(|c| c.name == "STORAGE_READBACK" && c.result == CheckResult::Pass);
        let restored = self
            .checks
            .iter()
            .any(|c| c.name == "STORAGE_RESTORE" && c.result == CheckResult::Pass);
        let checks: Vec<JsonCheck> = self
            .checks
            .iter()
            .map(|c| JsonCheck {
                name: c.name.clone(),
                result: c.result.label().to_string(),
                detail: c.detail.clone(),
            })
            .collect();
        if let Err(err) = serde_json::to_writer(
            std::io::stdout(),
            &JsonReport {
                device_path: dev,
                sector_written: written,
                sector_verified: verified,
                sector_restored: restored,
                checks,
            },
        ) {
            eprintln!("{PROGRAM}: failed to serialize JSON: {err}");
        }
    }
}

#[cfg(target_os = "redox")]
fn parse_args() -> Result<bool, String> {
    let mut json_mode = false;
    for arg in std::env::args().skip(1) {
        match arg.as_str() {
            "--json" => json_mode = true,
            "-h" | "--help" => {
                println!("{USAGE}");
                return Err(String::new());
            }
            _ => return Err(format!("unsupported argument: {arg}")),
        }
    }
    Ok(json_mode)
}

#[cfg(target_os = "redox")]
fn list_dir(path: &str) -> Vec<String> {
    match fs::read_dir(path) {
        Ok(entries) => entries
            .filter_map(|e| e.ok())
            .filter_map(|e| e.file_name().to_str().map(|s| s.to_string()))
            .collect(),
        Err(_) => Vec::new(),
    }
}

fn make_test_pattern() -> [u8; SECTOR_SIZE] {
    let mut buf = [0u8; SECTOR_SIZE];
    let mut offset = 0;
    while offset + TEST_PATTERN.len() <= SECTOR_SIZE {
        buf[offset..offset + TEST_PATTERN.len()].copy_from_slice(TEST_PATTERN);
        offset += TEST_PATTERN.len();
    }
    buf
}

/// On Redox, USB storage devices appear under /scheme/disk/ via usbscsid.
/// In QEMU tests, NVMe boot/extra disks also appear. We probe each disk
/// at the test sector offset to find one that is readable and writable.
#[cfg(target_os = "redox")]
fn find_usb_disk() -> Option<String> {
    let disks = list_dir("/scheme/disk");
    if disks.is_empty() {
        return None;
    }
    for disk_name in &disks {
        let path = format!("/scheme/disk/{}", disk_name);
        if let Ok(mut f) = fs::File::open(&path) {
            use std::io::{Read, Seek, SeekFrom};
            let offset = TEST_SECTOR * SECTOR_SIZE as u64;
            if f.seek(SeekFrom::Start(offset)).is_ok() {
                let mut buf = [0u8; 1];
                if f.read_exact(&mut buf).is_ok() {
                    drop(f);
                    return Some(path);
                }
            }
        }
    }
    None
}

/// Discover a writable block device and report USB storage class visibility.
#[cfg(target_os = "redox")]
fn check_storage_discovery() -> Check {
    match find_usb_disk() {
        Some(path) => {
            let usb_entries = list_dir("/scheme/usb");
            let mut storage_count = 0usize;
            for entry in &usb_entries {
                let port_path = format!("/scheme/usb/{}", entry);
                for port in list_dir(&port_path) {
                    let desc_path = format!("{}/{}/descriptors", port_path, port);
                    if let Ok(data) = fs::read_to_string(&desc_path) {
                        if let Ok(desc) = serde_json::from_str::<serde_json::Value>(&data) {
                            let class = desc.get("class").and_then(|v| v.as_u64()).unwrap_or(0);
                            if class == 8 {
                                storage_count += 1;
                            }
                        }
                    }
                }
            }
            Check::pass(
                "STORAGE_DISCOVERY",
                &format!(
                    "{} ({} USB storage class device(s) visible)",
                    path, storage_count
                ),
            )
        }
        None => Check::fail(
            "STORAGE_DISCOVERY",
            "no writable block device found under /scheme/disk/",
        ),
    }
}

/// Write a test pattern to the test sector.
#[cfg(target_os = "redox")]
fn check_storage_write(disk_path: &str, original_out: &mut [u8; SECTOR_SIZE]) -> Check {
    use std::io::{Read, Seek, SeekFrom, Write};
    let offset = TEST_SECTOR * SECTOR_SIZE as u64;

    let mut f = match fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(disk_path)
    {
        Ok(f) => f,
        Err(e) => {
            return Check::fail(
                "STORAGE_WRITE",
                &format!("failed to open {} for read/write: {e}", disk_path),
            );
        }
    };

    // Save original content for later restore
    if let Err(e) = f.seek(SeekFrom::Start(offset)) {
        return Check::fail(
            "STORAGE_WRITE",
            &format!("failed to seek to sector {TEST_SECTOR}: {e}"),
        );
    }
    if let Err(e) = f.read_exact(original_out) {
        return Check::fail(
            "STORAGE_WRITE",
            &format!("failed to read original sector {TEST_SECTOR}: {e}"),
        );
    }

    let pattern = make_test_pattern();
    if let Err(e) = f.seek(SeekFrom::Start(offset)) {
        return Check::fail("STORAGE_WRITE", &format!("failed to seek for write: {e}"));
    }
    if let Err(e) = f.write_all(&pattern) {
        return Check::fail(
            "STORAGE_WRITE",
            &format!("failed to write test pattern to sector {TEST_SECTOR}: {e}"),
        );
    }
    if let Err(e) = f.flush() {
        return Check::fail("STORAGE_WRITE", &format!("failed to flush write: {e}"));
    }
    Check::pass(
        "STORAGE_WRITE",
        &format!(
            "wrote test pattern to sector {TEST_SECTOR} of {}",
            disk_path
        ),
    )
}

/// Read back the test sector and verify the pattern matches.
#[cfg(target_os = "redox")]
fn check_storage_readback(disk_path: &str) -> Check {
    use std::io::{Read, Seek, SeekFrom};
    let offset = TEST_SECTOR * SECTOR_SIZE as u64;

    let mut f = match fs::File::open(disk_path) {
        Ok(f) => f,
        Err(e) => {
            return Check::fail(
                "STORAGE_READBACK",
                &format!("failed to reopen {}: {e}", disk_path),
            );
        }
    };

    if let Err(e) = f.seek(SeekFrom::Start(offset)) {
        return Check::fail(
            "STORAGE_READBACK",
            &format!("failed to seek to sector {TEST_SECTOR}: {e}"),
        );
    }

    let mut buf = [0u8; SECTOR_SIZE];
    if let Err(e) = f.read_exact(&mut buf) {
        return Check::fail(
            "STORAGE_READBACK",
            &format!("failed to read sector {TEST_SECTOR}: {e}"),
        );
    }

    let pattern = make_test_pattern();
    if buf == pattern {
        Check::pass(
            "STORAGE_READBACK",
            &format!("sector {TEST_SECTOR} readback matches test pattern"),
        )
    } else {
        let first_mismatch = buf
            .iter()
            .zip(pattern.iter())
            .enumerate()
            .find(|(_, (a, b))| a != b)
            .map(|(i, _)| i)
            .unwrap_or(SECTOR_SIZE);
        Check::fail(
            "STORAGE_READBACK",
            &format!("sector {TEST_SECTOR} readback mismatch at byte offset {first_mismatch}"),
        )
    }
}

/// Restore the original sector content.
#[cfg(target_os = "redox")]
fn check_storage_restore(disk_path: &str, original: &[u8; SECTOR_SIZE]) -> Check {
    use std::io::{Seek, SeekFrom, Write};
    let offset = TEST_SECTOR * SECTOR_SIZE as u64;

    let mut f = match fs::OpenOptions::new().write(true).open(disk_path) {
        Ok(f) => f,
        Err(e) => {
            return Check::fail(
                "STORAGE_RESTORE",
                &format!("failed to open {} for restore: {e}", disk_path),
            );
        }
    };

    if let Err(e) = f.seek(SeekFrom::Start(offset)) {
        return Check::fail(
            "STORAGE_RESTORE",
            &format!("failed to seek for restore: {e}"),
        );
    }
    if let Err(e) = f.write_all(original) {
        return Check::fail(
            "STORAGE_RESTORE",
            &format!("failed to restore original sector: {e}"),
        );
    }
    if let Err(e) = f.flush() {
        return Check::fail("STORAGE_RESTORE", &format!("failed to flush restore: {e}"));
    }
    Check::pass(
        "STORAGE_RESTORE",
        &format!("restored original content of sector {TEST_SECTOR}"),
    )
}

fn run() -> Result<(), String> {
    #[cfg(not(target_os = "redox"))]
    {
        if std::env::args().any(|a| a == "-h" || a == "--help") {
            println!("{USAGE}");
            return Err(String::new());
        }
        println!("{PROGRAM}: USB storage check requires Redox runtime");
        return Ok(());
    }
    #[cfg(target_os = "redox")]
    {
        let json_mode = parse_args()?;
        let mut report = Report::new(json_mode);

        let discovery = check_storage_discovery();
        let disk_path = discovery
            .detail
            .split_whitespace()
            .next()
            .map(|s| s.to_string());
        report.add(discovery);

        let disk_path = match disk_path {
            Some(p) => p,
            None => {
                report.print();
                return Err("no USB storage device discovered".to_string());
            }
        };

        let mut original = [0u8; SECTOR_SIZE];
        let write_check = check_storage_write(&disk_path, &mut original);
        let write_ok = write_check.result == CheckResult::Pass;
        report.add(write_check);

        if write_ok {
            report.add(check_storage_readback(&disk_path));
        } else {
            report.add(Check::skip(
                "STORAGE_READBACK",
                "skipped because write check failed",
            ));
        }

        if write_ok {
            report.add(check_storage_restore(&disk_path, &original));
        } else {
            report.add(Check::skip(
                "STORAGE_RESTORE",
                "skipped because write check failed",
            ));
        }

        report.print();
        if report.any_failed() {
            return Err("one or more USB storage checks failed".to_string());
        }
        Ok(())
    }
}

fn main() {
    if let Err(err) = run() {
        if err.is_empty() {
            process::exit(0);
        }
        eprintln!("{PROGRAM}: {err}");
        process::exit(1);
    }
}