use crate::drives::human_size;
use std::fmt;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OperationKind {
    WipeAndFlash,
    FlashIso,
    FormatOnly,
}

impl OperationKind {
    pub const ALL: [OperationKind; 3] = [
        OperationKind::WipeAndFlash,
        OperationKind::FlashIso,
        OperationKind::FormatOnly,
    ];

    pub fn title(self) -> &'static str {
        match self {
            OperationKind::WipeAndFlash => "Wipe + Flash",
            OperationKind::FlashIso => "Flash ISO",
            OperationKind::FormatOnly => "Format Drive",
        }
    }

    pub fn description(self) -> &'static str {
        match self {
            OperationKind::WipeAndFlash => {
                "Clear old signatures, then write the Arch ISO byte-for-byte."
            }
            OperationKind::FlashIso => "Write the selected ISO directly to the USB device.",
            OperationKind::FormatOnly => {
                "Create one clean partition and filesystem for normal storage."
            }
        }
    }

    pub fn needs_iso(self) -> bool {
        matches!(self, OperationKind::WipeAndFlash | OperationKind::FlashIso)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileSystem {
    Fat32,
    Exfat,
    Ext4,
}

impl FileSystem {
    pub const ALL: [FileSystem; 3] = [FileSystem::Fat32, FileSystem::Exfat, FileSystem::Ext4];

    pub fn label(self) -> &'static str {
        match self {
            FileSystem::Fat32 => "FAT32",
            FileSystem::Exfat => "exFAT",
            FileSystem::Ext4 => "ext4",
        }
    }

    fn script_name(self) -> &'static str {
        match self {
            FileSystem::Fat32 => "fat32",
            FileSystem::Exfat => "exfat",
            FileSystem::Ext4 => "ext4",
        }
    }
}

#[derive(Debug, Clone)]
pub struct OperationRequest {
    pub kind: OperationKind,
    pub drive_path: String,
    pub iso_path: Option<PathBuf>,
    pub filesystem: FileSystem,
    pub label: String,
}

impl OperationRequest {
    pub fn expected_image_bytes(&self) -> Option<u64> {
        let path = self.iso_path.as_ref()?;
        std::fs::metadata(path).ok().map(|metadata| metadata.len())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommandSpec {
    pub program: String,
    pub args: Vec<String>,
    pub script: String,
}

impl fmt::Display for CommandSpec {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}", self.program)?;
        for arg in &self.args {
            write!(formatter, " {}", shell_quote(arg))?;
        }
        Ok(())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum OperationError {
    #[error("select a removable drive first")]
    MissingDrive,
    #[error("drive path must be an absolute /dev path")]
    InvalidDrivePath,
    #[error("select an ISO or disk image first")]
    MissingIso,
    #[error("image does not exist: {0}")]
    MissingIsoFile(String),
    #[error("image must be a regular file: {0}")]
    IsoNotFile(String),
    #[error("image should be an .iso or .img file: {0}")]
    InvalidImageExtension(String),
    #[error("filesystem label must contain letters, numbers, dashes, or underscores")]
    InvalidLabel,
}

pub fn build_command(request: &OperationRequest) -> Result<CommandSpec, OperationError> {
    validate_request(request)?;
    let script = build_script(request)?;

    Ok(privileged_command(script))
}

pub fn build_script(request: &OperationRequest) -> Result<String, OperationError> {
    validate_request(request)?;

    let drive = shell_quote(&request.drive_path);
    let script = match request.kind {
        OperationKind::WipeAndFlash => {
            let iso = request
                .iso_path
                .as_ref()
                .ok_or(OperationError::MissingIso)?;
            flash_script(&drive, &shell_quote_path(iso), true)
        }
        OperationKind::FlashIso => {
            let iso = request
                .iso_path
                .as_ref()
                .ok_or(OperationError::MissingIso)?;
            flash_script(&drive, &shell_quote_path(iso), false)
        }
        OperationKind::FormatOnly => format_script(
            &drive,
            &shell_quote(&sanitize_label(&request.label)?),
            request.filesystem,
        ),
    };

    Ok(script)
}

pub fn validate_request(request: &OperationRequest) -> Result<(), OperationError> {
    validate_drive_path(&request.drive_path)?;

    if request.kind.needs_iso() {
        let path = request
            .iso_path
            .as_ref()
            .ok_or(OperationError::MissingIso)?;
        validate_image_path(path)?;
    }

    if matches!(request.kind, OperationKind::FormatOnly) {
        sanitize_label(&request.label)?;
    }

    Ok(())
}

pub fn validate_drive_path(path: &str) -> Result<(), OperationError> {
    if path.trim().is_empty() {
        return Err(OperationError::MissingDrive);
    }

    let valid = path.starts_with("/dev/")
        && path.len() > "/dev/".len()
        && path
            .chars()
            .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '/' | '_' | '-'));

    valid.then_some(()).ok_or(OperationError::InvalidDrivePath)
}

pub fn validate_image_path(path: &Path) -> Result<(), OperationError> {
    let display = path.display().to_string();
    let metadata =
        std::fs::metadata(path).map_err(|_| OperationError::MissingIsoFile(display.clone()))?;

    if !metadata.is_file() {
        return Err(OperationError::IsoNotFile(display));
    }

    let extension = path
        .extension()
        .and_then(|extension| extension.to_str())
        .unwrap_or_default()
        .to_ascii_lowercase();

    if extension != "iso" && extension != "img" {
        return Err(OperationError::InvalidImageExtension(display));
    }

    Ok(())
}

pub fn sanitize_label(label: &str) -> Result<String, OperationError> {
    let label = label.trim();
    if label.is_empty() || label.len() > 11 {
        return Err(OperationError::InvalidLabel);
    }

    let valid = label
        .chars()
        .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-'));
    if !valid {
        return Err(OperationError::InvalidLabel);
    }

    Ok(label.to_ascii_uppercase())
}

pub fn shell_quote(value: &str) -> String {
    if value.is_empty() {
        return "''".to_owned();
    }

    format!("'{}'", value.replace('\'', "'\"'\"'"))
}

fn shell_quote_path(path: &Path) -> String {
    shell_quote(&path.display().to_string())
}

fn privileged_command(script: String) -> CommandSpec {
    let shell = "/bin/sh";
    let mut args = Vec::new();

    if std::env::var_os("MOTTE_NO_PKEXEC").is_some() || !Path::new("/usr/bin/pkexec").exists() {
        args.push("-c".to_owned());
        args.push(script.clone());
        return CommandSpec {
            program: shell.to_owned(),
            args,
            script,
        };
    }

    args.push(shell.to_owned());
    args.push("-c".to_owned());
    args.push(script.clone());
    CommandSpec {
        program: "/usr/bin/pkexec".to_owned(),
        args,
        script,
    }
}

fn flash_script(drive: &str, iso: &str, wipe_first: bool) -> String {
    let wipe = if wipe_first {
        format!(
            "echo \"Motte: clearing old filesystem signatures\"\n{}\nwipefs -a \"$dev\"\nblockdev --rereadpt \"$dev\" || true\npartprobe \"$dev\" || true\nudevadm settle || true\n",
            wipe_partition_signatures_snippet()
        )
    } else {
        String::new()
    };

    format!(
        r#"set -eu
dev={drive}
iso={iso}
echo "Motte: preparing $dev"
{unmount_partitions}
{wipe}echo "Motte: writing $iso to $dev"
echo "Motte: dd is writing and flushing the target device"
dd if="$iso" of="$dev" bs=4M iflag=fullblock conv=fsync status=progress
echo "Motte: refreshing target partition table"
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: target flush complete"
echo "Motte: flash complete"
"#,
        drive = drive,
        iso = iso,
        unmount_partitions = unmount_partitions_snippet(),
        wipe = wipe,
    )
}

fn format_script(drive: &str, label: &str, filesystem: FileSystem) -> String {
    format!(
        r#"set -eu
dev={drive}
label={label}
fs={filesystem}
{format_tools}
echo "Motte: checking formatter tools"
case "$fs" in
  fat32)
    ensure_any_command "mkfs.vfat mkfs.fat" "dosfstools" "FAT32 formatter not found"
    ;;
  exfat)
    ensure_any_command "mkfs.exfat" "exfatprogs" "exFAT formatter not found"
    ;;
  ext4)
    ensure_any_command "mkfs.ext4" "e2fsprogs" "ext4 formatter not found"
    ;;
esac
echo "Motte: preparing $dev"
{unmount_partitions}
echo "Motte: clearing old filesystem signatures"
{wipe_partition_signatures}
if ! wipefs -a "$dev"; then
  echo "Motte: retrying signature clear after device settle"
  udevadm settle || true
  sleep 1
  wipefs -a "$dev"
fi
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: creating partition table"
parted -s "$dev" mklabel msdos
case "$fs" in
  fat32|exfat)
    parted -s "$dev" mkpart primary fat32 1MiB 100%
    parted -s "$dev" set 1 lba on || true
    ;;
  ext4)
    parted -s "$dev" mkpart primary ext4 1MiB 100%
    ;;
esac
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: waiting for the new partition"
attempt=0
part=""
while [ "$attempt" -lt 15 ]; do
  part=""
  part_count=0
  while read -r path type; do
    if [ "$type" = "part" ]; then
      part_count=$((part_count + 1))
      if [ -z "$part" ]; then
        part="$path"
      fi
    fi
  done <<MOTTE_PARTITION_SCAN
$(lsblk -rno PATH,TYPE "$dev" 2>/dev/null || true)
MOTTE_PARTITION_SCAN
  if [ "$part_count" -eq 1 ] && [ -n "$part" ] && [ -b "$part" ]; then
    break
  fi
  blockdev --rereadpt "$dev" || true
  partprobe "$dev" || true
  udevadm settle || true
  sleep 1
  attempt=$((attempt + 1))
done
if [ -z "$part" ] || [ ! -b "$part" ]; then
  echo "Motte: could not locate the new partition" >&2
  exit 1
fi
echo "Motte: clearing old signatures from $part"
wipefs -a "$part" || true
echo "Motte: formatting $part as $fs"
case "$fs" in
  fat32)
    fat_mkfs="$(first_available_command mkfs.vfat mkfs.fat)"
    "$fat_mkfs" -F 32 -n "$label" "$part"
    ;;
  exfat)
    mkfs.exfat -n "$label" "$part"
    ;;
  ext4)
    mkfs.ext4 -F -L "$label" "$part"
    ;;
esac
echo "Motte: flushing $dev"
blockdev --flushbufs "$dev" || true
echo "Motte: format complete"
"#,
        drive = drive,
        label = label,
        filesystem = shell_quote(filesystem.script_name()),
        format_tools = format_tool_install_snippet(),
        unmount_partitions = unmount_partitions_snippet(),
        wipe_partition_signatures = wipe_partition_signatures_snippet(),
    )
}

fn unmount_partitions_snippet() -> &'static str {
    r#"unmount_partition() {
  part="$1"
  if command -v udisksctl >/dev/null 2>&1; then
    udisksctl unmount -b "$part" >/dev/null 2>&1 || true
  fi
  findmnt -rn --source "$part" --output TARGET 2>/dev/null | while IFS= read -r mountpoint; do
    [ -n "$mountpoint" ] || continue
    echo "Motte: unmounting $part from $mountpoint"
    umount "$mountpoint" 2>/dev/null || umount -l "$mountpoint" 2>/dev/null || umount "$part" 2>/dev/null || umount -l "$part" 2>/dev/null || true
  done
}

wait_for_partition_unmounted() {
  part="$1"
  attempt=0
  while [ "$attempt" -lt 10 ]; do
    if ! findmnt -rn --source "$part" >/dev/null 2>&1; then
      return 0
    fi
    unmount_partition "$part"
    udevadm settle || true
    sleep 1
    attempt=$((attempt + 1))
  done
  echo "Motte: $part is still mounted after unmount attempts" >&2
  return 1
}

while IFS= read -r part; do
  [ "$part" = "$dev" ] && continue
  [ -n "$part" ] || continue
  if findmnt -rn --source "$part" >/dev/null 2>&1; then
    echo "Motte: releasing mounted $part"
    unmount_partition "$part"
    wait_for_partition_unmounted "$part"
  fi
done <<MOTTE_PARTS
$(lsblk -ln -o PATH "$dev")
MOTTE_PARTS"#
}

fn wipe_partition_signatures_snippet() -> &'static str {
    r#"while IFS= read -r part; do
  [ "$part" = "$dev" ] && continue
  if [ -n "$part" ]; then
    wipefs -a "$part" || true
  fi
done <<MOTTE_EXISTING_PARTS
$(lsblk -ln -o PATH "$dev")
MOTTE_EXISTING_PARTS"#
}

fn format_tool_install_snippet() -> &'static str {
    r#"first_available_command() {
  for command_name in "$@"; do
    if command -v "$command_name" >/dev/null 2>&1; then
      command -v "$command_name"
      return 0
    fi
  done
  return 1
}

install_package() {
  package_name="$1"
  if command -v pacman >/dev/null 2>&1; then
    echo "Motte: installing $package_name with pacman"
    pacman -Sy --noconfirm --needed "$package_name"
  elif command -v apt-get >/dev/null 2>&1; then
    echo "Motte: installing $package_name with apt-get"
    export DEBIAN_FRONTEND=noninteractive
    apt-get update
    apt-get install -y "$package_name"
  elif command -v dnf >/dev/null 2>&1; then
    echo "Motte: installing $package_name with dnf"
    dnf install -y "$package_name"
  elif command -v zypper >/dev/null 2>&1; then
    echo "Motte: installing $package_name with zypper"
    zypper --non-interactive install "$package_name"
  else
    echo "Motte: no supported package manager found for $package_name" >&2
    return 127
  fi
}

ensure_any_command() {
  command_names="$1"
  package_name="$2"
  missing_message="$3"
  for command_name in $command_names; do
    if command -v "$command_name" >/dev/null 2>&1; then
      return 0
    fi
  done
  echo "Motte: $missing_message; attempting to install $package_name"
  install_package "$package_name"
  for command_name in $command_names; do
    if command -v "$command_name" >/dev/null 2>&1; then
      return 0
    fi
  done
  echo "Motte: $missing_message; install $package_name manually" >&2
  return 127
}
"#
}

#[derive(Debug, Clone, PartialEq)]
pub enum OperationEvent {
    Started(String),
    Log(String),
    Progress { copied_bytes: u64, total_bytes: u64 },
    Finished(Result<(), String>),
}

pub fn start_operation(
    command: CommandSpec,
    expected_bytes: Option<u64>,
) -> Receiver<OperationEvent> {
    let (tx, rx) = mpsc::channel();
    thread::spawn(move || run_operation(command, expected_bytes, tx));
    rx
}

fn run_operation(command: CommandSpec, expected_bytes: Option<u64>, tx: Sender<OperationEvent>) {
    let _ = tx.send(OperationEvent::Started(command.to_string()));

    let mut child = match Command::new(&command.program)
        .args(&command.args)
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
    {
        Ok(child) => child,
        Err(error) => {
            let _ = tx.send(OperationEvent::Finished(Err(format!(
                "failed to start operation: {error}"
            ))));
            return;
        }
    };

    let stdout = child.stdout.take();
    let stderr = child.stderr.take();
    let stdout_thread = stdout.map(|stream| forward_stream(stream, tx.clone(), expected_bytes));
    let stderr_thread = stderr.map(|stream| forward_stream(stream, tx.clone(), expected_bytes));

    let result = match child.wait() {
        Ok(status) if status.success() => Ok(()),
        Ok(status) => Err(format!("operation exited with status {status}")),
        Err(error) => Err(format!("failed to wait for operation: {error}")),
    };

    if let Some(thread) = stdout_thread {
        let _ = thread.join();
    }
    if let Some(thread) = stderr_thread {
        let _ = thread.join();
    }

    let _ = tx.send(OperationEvent::Finished(result));
}

fn forward_stream<R>(
    mut stream: R,
    tx: Sender<OperationEvent>,
    expected_bytes: Option<u64>,
) -> thread::JoinHandle<()>
where
    R: Read + Send + 'static,
{
    thread::spawn(move || {
        let mut buffer = [0; 1024];
        let mut pending = String::new();

        loop {
            let bytes_read = match stream.read(&mut buffer) {
                Ok(0) => break,
                Ok(bytes_read) => bytes_read,
                Err(error) => {
                    let _ = tx.send(OperationEvent::Log(format!("stream read failed: {error}")));
                    break;
                }
            };

            let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
            for character in chunk.chars() {
                if character == '\n' || character == '\r' {
                    flush_pending(&tx, &mut pending, expected_bytes);
                } else {
                    pending.push(character);
                }
            }
        }

        flush_pending(&tx, &mut pending, expected_bytes);
    })
}

fn flush_pending(tx: &Sender<OperationEvent>, pending: &mut String, expected_bytes: Option<u64>) {
    let line = pending.trim().to_owned();
    pending.clear();

    if line.is_empty() {
        return;
    }

    if let (Some(copied_bytes), Some(total_bytes)) =
        (parse_dd_progress_bytes(&line), expected_bytes)
    {
        let _ = tx.send(OperationEvent::Progress {
            copied_bytes: copied_bytes.min(total_bytes),
            total_bytes,
        });
    }

    let _ = tx.send(OperationEvent::Log(line));
}

pub fn parse_dd_progress_bytes(line: &str) -> Option<u64> {
    let line = line.trim();
    let (bytes, rest) = line.split_once(" bytes")?;
    if !rest.contains("copied") {
        return None;
    }

    bytes.trim().replace(',', "").parse().ok()
}

pub fn progress_label(copied_bytes: u64, total_bytes: u64) -> String {
    let percent = if total_bytes == 0 {
        0.0
    } else {
        copied_bytes as f32 / total_bytes as f32 * 100.0
    };

    format!(
        "{} / {} ({percent:.1}%)",
        human_size(copied_bytes),
        human_size(total_bytes)
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn quotes_shell_values_safely() {
        assert_eq!(shell_quote(""), "''");
        assert_eq!(shell_quote("/dev/sdb"), "'/dev/sdb'");
        assert_eq!(shell_quote("arch user's.iso"), "'arch user'\"'\"'s.iso'");
    }

    #[test]
    fn validates_iso_paths() {
        let mut file = NamedTempFile::with_suffix(".iso").expect("temp iso");
        writeln!(file, "fake iso").expect("write temp iso");

        validate_image_path(file.path()).expect("valid iso path");
    }

    #[test]
    fn rejects_non_images() {
        let file = NamedTempFile::new().expect("temp file");
        let err = validate_image_path(file.path()).expect_err("reject extension");

        assert!(matches!(err, OperationError::InvalidImageExtension(_)));
    }

    #[test]
    fn builds_wipe_and_flash_script() {
        let mut file = NamedTempFile::with_suffix(".iso").expect("temp iso");
        writeln!(file, "fake iso").expect("write temp iso");
        let request = OperationRequest {
            kind: OperationKind::WipeAndFlash,
            drive_path: "/dev/sdb".to_owned(),
            iso_path: Some(file.path().to_path_buf()),
            filesystem: FileSystem::Fat32,
            label: "MOTTE".to_owned(),
        };

        let script = build_script(&request).expect("build script");

        assert!(script.contains("wipefs -a \"$dev\""));
        assert!(script.contains("MOTTE_EXISTING_PARTS"));
        assert!(script.contains("dd if=\"$iso\" of=\"$dev\""));
        assert!(script.contains("iflag=fullblock"));
        assert!(script.contains("conv=fsync"));
        assert!(script.contains("Motte: refreshing target partition table"));
        assert!(script.contains("blockdev --rereadpt \"$dev\""));
        assert!(script.contains("partprobe \"$dev\""));
        assert!(script.contains("Motte: target flush complete"));
        assert!(!script.contains("\nsync\n"));
        assert!(script.contains("udisksctl unmount -b \"$part\""));
        assert!(script.contains("findmnt -rn --source \"$part\""));
        assert!(script.contains("umount \"$mountpoint\""));
        assert!(script.contains("wait_for_partition_unmounted \"$part\""));
    }

    #[test]
    fn builds_format_script_for_fat32() {
        let request = OperationRequest {
            kind: OperationKind::FormatOnly,
            drive_path: "/dev/sdc".to_owned(),
            iso_path: None,
            filesystem: FileSystem::Fat32,
            label: "arch_usb".to_owned(),
        };

        let script = build_script(&request).expect("build script");

        assert!(script.contains("parted -s \"$dev\" mklabel msdos"));
        assert!(script.contains("Motte: checking formatter tools"));
        assert!(script.contains("ensure_any_command \"mkfs.vfat mkfs.fat\" \"dosfstools\""));
        assert!(script.contains("parted -s \"$dev\" mkpart primary fat32"));
        assert!(script.contains("parted -s \"$dev\" set 1 lba on"));
        assert!(script.contains("blockdev --rereadpt \"$dev\""));
        assert!(script.contains("Motte: waiting for the new partition"));
        assert!(script.contains("while [ \"$attempt\" -lt 15 ]"));
        assert!(script.contains("lsblk -rno PATH,TYPE \"$dev\""));
        assert!(script.contains("first_available_command mkfs.vfat mkfs.fat"));
        assert!(script.contains("pacman -Sy --noconfirm --needed \"$package_name\""));
        assert!(!script.contains("sed -n '2p'"));
        assert!(script.contains("label='ARCH_USB'"));
    }

    #[test]
    fn parses_dd_progress_lines() {
        assert_eq!(
            parse_dd_progress_bytes("1,048,576 bytes (1.0 MB, 1.0 MiB) copied, 1 s, 1 MB/s"),
            Some(1_048_576)
        );
        assert_eq!(parse_dd_progress_bytes("512 records in"), None);
    }

    #[test]
    fn labels_progress() {
        assert_eq!(progress_label(512, 1024), "512 B / 1.0 KiB (50.0%)");
    }
}
