use serde::Deserialize;
use serde_json::Value;
use std::process::Command;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DriveInfo {
    pub name: String,
    pub path: String,
    pub model: Option<String>,
    pub serial: Option<String>,
    pub size_bytes: u64,
    pub removable: bool,
    pub read_only: bool,
    pub transport: Option<String>,
    pub mountpoints: Vec<String>,
    pub partitions: Vec<PartitionInfo>,
}

impl DriveInfo {
    pub fn display_name(&self) -> String {
        match self
            .model
            .as_deref()
            .filter(|model| !model.trim().is_empty())
        {
            Some(model) => format!("{} ({})", model.trim(), self.path),
            None => self.path.clone(),
        }
    }

    pub fn human_size(&self) -> String {
        human_size(self.size_bytes)
    }

    pub fn safety_badge(&self) -> &'static str {
        if self.read_only {
            "read-only"
        } else if self.transport.as_deref() == Some("usb") {
            "usb"
        } else if self.removable {
            "removable"
        } else {
            "hotplug"
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PartitionInfo {
    pub name: String,
    pub path: String,
    pub size_bytes: u64,
    pub mountpoints: Vec<String>,
}

impl PartitionInfo {
    pub fn human_size(&self) -> String {
        human_size(self.size_bytes)
    }
}

pub trait DriveScanner {
    fn scan(&self) -> Result<Vec<DriveInfo>, DriveScanError>;
}

#[derive(Debug, Default, Clone, Copy)]
pub struct LsblkScanner;

impl DriveScanner for LsblkScanner {
    fn scan(&self) -> Result<Vec<DriveInfo>, DriveScanError> {
        let output = Command::new("lsblk")
            .args([
                "--json",
                "--bytes",
                "--output",
                "NAME,PATH,MODEL,SERIAL,SIZE,RM,RO,TYPE,MOUNTPOINTS,TRAN,HOTPLUG",
            ])
            .output()?;

        if !output.status.success() {
            return Err(DriveScanError::LsblkFailed(
                String::from_utf8_lossy(&output.stderr).trim().to_owned(),
            ));
        }

        parse_lsblk_json(&output.stdout)
    }
}

#[derive(Debug, thiserror::Error)]
pub enum DriveScanError {
    #[error("failed to run lsblk: {0}")]
    Command(#[from] std::io::Error),
    #[error("failed to parse lsblk JSON: {0}")]
    Json(#[from] serde_json::Error),
    #[error("lsblk failed: {0}")]
    LsblkFailed(String),
}

pub fn parse_lsblk_json(json: impl AsRef<[u8]>) -> Result<Vec<DriveInfo>, DriveScanError> {
    let lsblk: LsblkOutput = serde_json::from_slice(json.as_ref())?;
    let mut drives = lsblk
        .blockdevices
        .into_iter()
        .filter_map(device_to_drive)
        .collect::<Vec<_>>();

    drives.sort_by(|left, right| left.path.cmp(&right.path));
    Ok(drives)
}

fn device_to_drive(device: LsblkDevice) -> Option<DriveInfo> {
    if device.kind.as_deref() != Some("disk") {
        return None;
    }

    let transport = clean(device.tran);
    let is_usb = transport.as_deref() == Some("usb");
    let removable = device.rm || device.hotplug || is_usb;
    if !removable {
        return None;
    }

    let path = clean(device.path)?;
    let name = clean(device.name).unwrap_or_else(|| path.trim_start_matches("/dev/").to_owned());
    let mut mountpoints = compact_mountpoints(device.mountpoints);
    let partitions = device
        .children
        .unwrap_or_default()
        .into_iter()
        .filter_map(|child| {
            let child_path = clean(child.path)?;
            let child_name = clean(child.name).unwrap_or_else(|| child_path.clone());
            let child_mountpoints = compact_mountpoints(child.mountpoints);
            mountpoints.extend(child_mountpoints.iter().cloned());

            Some(PartitionInfo {
                name: child_name,
                path: child_path,
                size_bytes: child.size.unwrap_or(0),
                mountpoints: child_mountpoints,
            })
        })
        .collect::<Vec<_>>();

    mountpoints.sort();
    mountpoints.dedup();

    Some(DriveInfo {
        name,
        path,
        model: clean(device.model),
        serial: clean(device.serial),
        size_bytes: device.size.unwrap_or(0),
        removable,
        read_only: device.ro,
        transport,
        mountpoints,
        partitions,
    })
}

fn compact_mountpoints(mountpoints: Vec<Option<String>>) -> Vec<String> {
    mountpoints
        .into_iter()
        .flatten()
        .filter_map(|mountpoint| clean(Some(mountpoint)))
        .collect()
}

fn clean(value: Option<String>) -> Option<String> {
    let value = value?.trim().to_owned();
    (!value.is_empty()).then_some(value)
}

pub fn human_size(bytes: u64) -> String {
    const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
    let mut value = bytes as f64;
    let mut unit = 0;

    while value >= 1024.0 && unit < UNITS.len() - 1 {
        value /= 1024.0;
        unit += 1;
    }

    if unit == 0 {
        format!("{} {}", bytes, UNITS[unit])
    } else {
        format!("{value:.1} {}", UNITS[unit])
    }
}

#[derive(Debug, Deserialize)]
struct LsblkOutput {
    blockdevices: Vec<LsblkDevice>,
}

#[derive(Debug, Deserialize)]
struct LsblkDevice {
    name: Option<String>,
    path: Option<String>,
    model: Option<String>,
    serial: Option<String>,
    #[serde(default, deserialize_with = "deserialize_optional_u64")]
    size: Option<u64>,
    #[serde(default, deserialize_with = "deserialize_boolish")]
    rm: bool,
    #[serde(default, deserialize_with = "deserialize_boolish")]
    ro: bool,
    #[serde(rename = "type")]
    kind: Option<String>,
    #[serde(default, deserialize_with = "deserialize_mountpoints")]
    mountpoints: Vec<Option<String>>,
    tran: Option<String>,
    #[serde(default, deserialize_with = "deserialize_boolish")]
    hotplug: bool,
    #[serde(default)]
    children: Option<Vec<LsblkDevice>>,
}

fn deserialize_optional_u64<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<Value>::deserialize(deserializer)?;
    Ok(match value {
        Some(Value::Number(number)) => number.as_u64(),
        Some(Value::String(text)) => text.trim().parse().ok(),
        _ => None,
    })
}

fn deserialize_boolish<'de, D>(deserializer: D) -> Result<bool, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<Value>::deserialize(deserializer)?;
    Ok(match value {
        Some(Value::Bool(value)) => value,
        Some(Value::Number(number)) => number.as_u64().unwrap_or(0) != 0,
        Some(Value::String(text)) => matches!(
            text.trim().to_ascii_lowercase().as_str(),
            "1" | "true" | "yes" | "y"
        ),
        _ => false,
    })
}

fn deserialize_mountpoints<'de, D>(deserializer: D) -> Result<Vec<Option<String>>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<Value>::deserialize(deserializer)?;
    Ok(match value {
        Some(Value::Array(values)) => values
            .into_iter()
            .map(|value| match value {
                Value::String(text) => Some(text),
                _ => None,
            })
            .collect(),
        Some(Value::String(text)) => vec![Some(text)],
        _ => Vec::new(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_only_removable_drives() {
        let json = br#"
        {
          "blockdevices": [
            {"name":"nvme0n1","path":"/dev/nvme0n1","model":"Internal","size":512110190592,"rm":false,"ro":false,"type":"disk","mountpoints":[null],"tran":"nvme","hotplug":false},
            {"name":"sdb","path":"/dev/sdb","model":"USB Stick","serial":"ABC","size":"16008609792","rm":true,"ro":false,"type":"disk","mountpoints":[null],"tran":"usb","hotplug":true,
              "children":[{"name":"sdb1","path":"/dev/sdb1","size":16000000000,"rm":true,"ro":false,"type":"part","mountpoints":["/run/media/usb"]}]}
          ]
        }
        "#;

        let drives = parse_lsblk_json(json).expect("parse lsblk");

        assert_eq!(drives.len(), 1);
        assert_eq!(drives[0].path, "/dev/sdb");
        assert_eq!(drives[0].model.as_deref(), Some("USB Stick"));
        assert_eq!(drives[0].mountpoints, vec!["/run/media/usb"]);
        assert_eq!(drives[0].partitions[0].path, "/dev/sdb1");
    }

    #[test]
    fn accepts_numeric_bool_fields() {
        let json = br#"
        {
          "blockdevices": [
            {"name":"sdc","path":"/dev/sdc","model":"Reader","size":"1024","rm":0,"ro":0,"type":"disk","mountpoints":null,"tran":"usb","hotplug":1}
          ]
        }
        "#;

        let drives = parse_lsblk_json(json).expect("parse lsblk");

        assert_eq!(drives.len(), 1);
        assert_eq!(drives[0].size_bytes, 1024);
        assert_eq!(drives[0].safety_badge(), "usb");
    }

    #[test]
    fn formats_human_sizes() {
        assert_eq!(human_size(0), "0 B");
        assert_eq!(human_size(1024), "1.0 KiB");
        assert_eq!(human_size(4 * 1024 * 1024 * 1024), "4.0 GiB");
    }
}
