use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;

use sal_os;
use sal_process;

#[derive(Debug)]
pub enum ImagePrepError {
    Io(String),
    InvalidInput(String),
    CommandFailed(String),
    NotImplemented(String),
}

impl std::fmt::Display for ImagePrepError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ImagePrepError::Io(e) => write!(f, "IO error: {}", e),
            ImagePrepError::InvalidInput(e) => write!(f, "Invalid input: {}", e),
            ImagePrepError::CommandFailed(e) => write!(f, "Command failed: {}", e),
            ImagePrepError::NotImplemented(e) => write!(f, "Not implemented: {}", e),
        }
    }
}

impl std::error::Error for ImagePrepError {}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Flavor {
    Ubuntu,
    Alpine,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetPlanOpts {
    #[serde(default = "default_dhcp4")]
    pub dhcp4: bool,
    #[serde(default)]
    pub dhcp6: bool,
    /// Static IPv6 address to assign in guest (temporary behavior)
    pub ipv6_addr: Option<String>, // e.g., "400::10/64"
    pub gw6: Option<String>,       // e.g., "400::1"
}

fn default_dhcp4() -> bool {
    true
}

impl Default for NetPlanOpts {
    fn default() -> Self {
        Self {
            dhcp4: true,
            dhcp6: false,
            ipv6_addr: Some("400::10/64".into()),
            gw6: Some("400::1".into()),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImagePrepOptions {
    pub flavor: Flavor,
    /// VM id (used for working directory layout and tap/mac derivations)
    pub id: String,
    /// Optional source path override, defaults to /images/<flavor default filename>
    pub source: Option<String>,
    /// Optional VM target directory, defaults to $HOME/hero/virt/vms/<id>
    pub target_dir: Option<String>,
    /// Netplan options
    #[serde(default)]
    pub net: NetPlanOpts,
    /// Disable cloud-init networking
    #[serde(default = "default_disable_cloud_init_net")]
    pub disable_cloud_init_net: bool,
}

fn default_disable_cloud_init_net() -> bool {
    true
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImagePrepResult {
    pub raw_disk: String,
    pub root_uuid: String,
    pub boot_uuid: String,
    pub work_qcow2: String,
}

fn hero_vm_root() -> String {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
    format!("{}/hero/virt/vms", home.trim_end_matches('/'))
}

fn default_source_for_flavor(flavor: &Flavor) -> (&'static str, bool) {
    match flavor {
        Flavor::Ubuntu => ("/images/noble-server-cloudimg-amd64.img", true),
        Flavor::Alpine => ("/images/alpine-virt-cloudimg-amd64.qcow2", true),
    }
}

fn fail(e: &str) -> ImagePrepError {
    ImagePrepError::CommandFailed(e.to_string())
}

fn run_script(script: &str) -> Result<sal_process::CommandResult, ImagePrepError> {
    match sal_process::run(script).silent(true).die(false).execute() {
        Ok(res) => {
            if res.success {
                Ok(res)
            } else {
                Err(ImagePrepError::CommandFailed(format!(
                    "{}{}",
                    res.stdout, res.stderr
                )))
            }
        }
        Err(e) => Err(ImagePrepError::CommandFailed(e.to_string())),
    }
}

/// Prepare a base cloud image for booting under Cloud Hypervisor:
/// - make a per-VM working copy
/// - attach via nbd, mount root/boot
/// - retag UUIDs, update fstab, write minimal grub.cfg
/// - generate netplan (DHCPv4, static IPv6 placeholder), disable cloud-init net
/// - convert to raw disk in VM dir
pub fn image_prepare(opts: &ImagePrepOptions) -> Result<ImagePrepResult, ImagePrepError> {
    // Resolve source image
    let (def_src, _must_exist) = default_source_for_flavor(&opts.flavor);
    let src = opts.source.clone().unwrap_or_else(|| def_src.to_string());
    if !Path::new(&src).exists() {
        return Err(ImagePrepError::InvalidInput(format!(
            "source image not found: {}",
            src
        )));
    }

    // Resolve VM dir
    let vm_dir = opts
        .target_dir
        .clone()
        .unwrap_or_else(|| format!("{}/{}", hero_vm_root(), opts.id));
    sal_os::mkdir(&vm_dir).map_err(|e| ImagePrepError::Io(e.to_string()))?;

    // Work qcow2 copy path and mount points
    let work_qcow2 = format!("{}/work.qcow2", vm_dir);
    let raw_path = format!("{}/disk.raw", vm_dir);
    let mnt_root = format!("/mnt/hero-img/{}/root", opts.id);
    let mnt_boot = format!("/mnt/hero-img/{}/boot", opts.id);

    // Only Ubuntu implemented for now
    match opts.flavor {
        Flavor::Ubuntu => {
            // Build bash script that performs all steps and echos "RAW|ROOT_UUID|BOOT_UUID" at end
            let net_ipv6 = opts.net.ipv6_addr.clone().unwrap_or_else(|| "400::10/64".into());
            let gw6 = opts.net.gw6.clone().unwrap_or_else(|| "400::1".into());
            let disable_ci_net = opts.disable_cloud_init_net;

            // Keep script small and robust; avoid brace-heavy awk to simplify escaping.
            let script = format!(
                "#!/bin/bash -e
set -euo pipefail

SRC={src}
VM_DIR={vm_dir}
WORK={work}
MNT_ROOT={mnt_root}
MNT_BOOT={mnt_boot}
RAW={raw}

mkdir -p \"$VM_DIR\"
mkdir -p \"$(dirname \"$MNT_ROOT\")\"
mkdir -p \"$MNT_ROOT\" \"$MNT_BOOT\"

# Make per-VM working copy (reflink if supported)
cp --reflink=auto -f \"$SRC\" \"$WORK\"

# Load NBD with sufficient partitions
modprobe nbd max_part=63

# Pick a free /dev/nbdX and connect the qcow2
NBD=\"\"
for i in $(seq 0 15); do
  DEV=\"/dev/nbd$i\"
  # Skip devices that have any mounted partitions (avoid reusing in-use NBDs)
  if findmnt -rn -S \"$DEV\" >/dev/null 2>&1 || \
     findmnt -rn -S \"${{DEV}}p1\" >/dev/null 2>&1 || \
     findmnt -rn -S \"${{DEV}}p14\" >/dev/null 2>&1 || \
     findmnt -rn -S \"${{DEV}}p15\" >/dev/null 2>&1 || \
     findmnt -rn -S \"${{DEV}}p16\" >/dev/null 2>&1; then
    continue
  fi
  # Ensure it's not connected (ignore errors if already disconnected)
  qemu-nbd --disconnect \"$DEV\" >/dev/null 2>&1 || true
  if qemu-nbd --format=qcow2 --connect=\"$DEV\" \"$WORK\"; then
    NBD=\"$DEV\"
    break
  fi
done
if [ -z \"$NBD\" ]; then
  echo \"No free /dev/nbdX device available\" >&2
  exit 1
fi

echo \"Selected NBD: $NBD\"

# Settle and probe partitions
udevadm settle >/dev/null 2>&1 || true
blockdev --rereadpt \"$NBD\" >/dev/null 2>&1 || true
partprobe \"$NBD\" >/dev/null 2>&1 || true
for t in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15; do
  if [ -b \"${{NBD}}p1\" ]; then
    sz=$(blockdev --getsize64 \"${{NBD}}p1\" 2>/dev/null || echo 0)
    if [ \"$sz\" -gt 0 ]; then
      break
    fi
  fi
  sleep 0.4
  udevadm settle >/dev/null 2>&1 || true
  blockdev --rereadpt \"$NBD\" >/dev/null 2>&1 || true
  partprobe \"$NBD\" >/dev/null 2>&1 || true
done

ROOT_DEV=\"${{NBD}}p1\"
# Prefer p16, else p15
if [ -b \"${{NBD}}p16\" ]; then
  BOOT_DEV=\"${{NBD}}p16\"
elif [ -b \"${{NBD}}p15\" ]; then
  BOOT_DEV=\"${{NBD}}p15\"
else
  echo \"Boot partition not found on $NBD (tried p16 and p15)\" >&2
  exit 33
fi

echo \"ROOT_DEV=$ROOT_DEV BOOT_DEV=$BOOT_DEV\"

if [ ! -b \"$ROOT_DEV\" ]; then
  echo \"Root partition not found: $ROOT_DEV\" >&2
  exit 32
fi

cleanup() {{
  set +e
  umount \"$MNT_BOOT\" 2>/dev/null || true
  umount \"$MNT_ROOT\" 2>/dev/null || true
  [ -n \"$NBD\" ] && qemu-nbd --disconnect \"$NBD\" 2>/dev/null || true
  rmmod nbd 2>/dev/null || true
}}
trap cleanup EXIT

# Ensure partitions are readable before mounting
for t in 1 2 3 4 5 6 7 8; do
  szr=$(blockdev --getsize64 \"$ROOT_DEV\" 2>/dev/null || echo 0)
  szb=$(blockdev --getsize64 \"$BOOT_DEV\" 2>/dev/null || echo 0)
  if [ \"$szr\" -gt 0 ] && [ \"$szb\" -gt 0 ] && blkid \"$ROOT_DEV\" >/dev/null 2>&1; then
    break
  fi
  sleep 0.4
  udevadm settle >/dev/null 2>&1 || true
  blockdev --rereadpt \"$NBD\" >/dev/null 2>&1 || true
  partprobe \"$NBD\" >/dev/null 2>&1 || true
done

# Mount and mutate (with retries to avoid races)
mounted_root=0
for t in 1 2 3 4 5 6 7 8 9 10; do
  if mount \"$ROOT_DEV\" \"$MNT_ROOT\"; then
    mounted_root=1
    break
  fi
  sleep 0.5
  udevadm settle >/dev/null 2>&1 || true
  partprobe \"$NBD\" >/dev/null 2>&1 || true
done
if [ \"$mounted_root\" -ne 1 ]; then
  echo \"Failed to mount root $ROOT_DEV\" >&2
  exit 32
fi

mounted_boot=0
for t in 1 2 3 4 5; do
  if mount \"$BOOT_DEV\" \"$MNT_BOOT\"; then
    mounted_boot=1
    break
  fi
  sleep 0.5
  udevadm settle >/dev/null 2>&1 || true
  partprobe \"$NBD\" >/dev/null 2>&1 || true
done
if [ \"$mounted_boot\" -ne 1 ]; then
  echo \"Failed to mount boot $BOOT_DEV\" >&2
  exit 33
fi

# Change UUIDs (best-effort)
tune2fs -U random \"$ROOT_DEV\" || true
tune2fs -U random \"$BOOT_DEV\" || true

ROOT_UUID=$(blkid -o value -s UUID \"$ROOT_DEV\")
BOOT_UUID=$(blkid -o value -s UUID \"$BOOT_DEV\")

# Update fstab
sed -i \"s/UUID=[a-f0-9-]* \\/ /UUID=$ROOT_UUID \\/ /\" \"$MNT_ROOT/etc/fstab\"
sed -i \"s/UUID=[a-f0-9-]* \\/boot /UUID=$BOOT_UUID \\/boot /\" \"$MNT_ROOT/etc/fstab\"

# Minimal grub.cfg (note: braces escaped for Rust format!)
mkdir -p \"$MNT_BOOT/grub\"
KERNEL=$(ls -1 \"$MNT_BOOT\"/vmlinuz-* | sort -V | tail -n1 | xargs -n1 basename)
INITRD=$(ls -1 \"$MNT_BOOT\"/initrd.img-* | sort -V | tail -n1 | xargs -n1 basename)
cat > \"$MNT_BOOT/grub/grub.cfg\" << EOF
set default=0
set timeout=3
menuentry 'Ubuntu Cloud' {{
    insmod part_gpt
    insmod ext2
    insmod gzio
    search --no-floppy --fs-uuid --set=root $BOOT_UUID
    linux /$KERNEL root=/dev/vda1 ro console=ttyS0
    initrd /$INITRD
}}
EOF

# Netplan config
rm -f \"$MNT_ROOT\"/etc/netplan/*.yaml
mkdir -p \"$MNT_ROOT\"/etc/netplan
cat > \"$MNT_ROOT/etc/netplan/01-netconfig.yaml\" << EOF
network:
  version: 2
  ethernets:
    eth0:
      dhcp4: {dhcp4}
      dhcp6: {dhcp6}
      addresses:
        - {ipv6}
      routes:
        - to: \"::/0\"
          via: {gw6}
      nameservers:
        addresses: [8.8.8.8, 1.1.1.1, 2001:4860:4860::8888]
EOF

# Disable cloud-init networking (optional but default)
if [ \"{disable_ci_net}\" = \"true\" ]; then
  mkdir -p \"$MNT_ROOT/etc/cloud/cloud.cfg.d\"
  echo \"network: {{config: disabled}}\" > \"$MNT_ROOT/etc/cloud/cloud.cfg.d/99-disable-network-config.cfg\"
fi

# Convert prepared image to raw (ensure source not locked)
umount \"$MNT_BOOT\" 2>/dev/null || true
umount \"$MNT_ROOT\" 2>/dev/null || true
if [ -n \"$NBD\" ]; then
  qemu-nbd --disconnect \"$NBD\" 2>/dev/null || true
  rmmod nbd 2>/dev/null || true
fi
rm -f \"$RAW\"
qemu-img convert -U -f qcow2 -O raw \"$WORK\" \"$RAW\"

# Output result triple
echo \"$RAW|$ROOT_UUID|$BOOT_UUID\"
",
                src = shell_escape(&src),
                vm_dir = shell_escape(&vm_dir),
                work = shell_escape(&work_qcow2),
                mnt_root = shell_escape(&mnt_root),
                mnt_boot = shell_escape(&mnt_boot),
                raw = shell_escape(&raw_path),
                dhcp4 = if opts.net.dhcp4 { "true" } else { "false" },
                dhcp6 = if opts.net.dhcp6 { "true" } else { "false" },
                ipv6 = shell_escape(&net_ipv6),
                gw6 = shell_escape(&gw6),
                disable_ci_net = if disable_ci_net { "true" } else { "false" },
            );

            let res = run_script(&script)?;
            let line = res.stdout.trim().lines().last().unwrap_or("").trim().to_string();
            let parts: Vec<_> = line.split('|').map(|s| s.to_string()).collect();
            if parts.len() != 3 {
                return Err(fail(&format!(
                    "unexpected output from image_prepare script, expected RAW|ROOT_UUID|BOOT_UUID, got: {}",
                    line
                )));
            }
            Ok(ImagePrepResult {
                raw_disk: parts[0].clone(),
                root_uuid: parts[1].clone(),
                boot_uuid: parts[2].clone(),
                work_qcow2,
            })
        }
        Flavor::Alpine => Err(ImagePrepError::NotImplemented(
            "Alpine image_prepare not implemented yet".into(),
        )),
    }
}

fn shell_escape(s: &str) -> String {
    if s.is_empty() {
        return "''".into();
    }
    if s.chars().all(|c| c.is_ascii_alphanumeric() || "-_./=:".contains(c)) {
        return s.into();
    }
    let mut out = String::from("'");
    for ch in s.chars() {
        if ch == '\'' {
            out.push_str("'\"'\"'");
        } else {
            out.push(ch);
        }
    }
    out.push('\'');
    out
}