use anyhow::{bail, Context, Result};
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

use crate::config::BuildConfig;

/// Where the kernel that ends up in the hekate `uImage` comes from.
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum KernelSource {
    /// The known-good switchroot L4T kernel (display/panel/regulator patches).
    /// This is the default because a vanilla kernel cannot drive the Switch
    /// DSI panel.
    Switchroot,
    /// The kernel built by this project, wrapped into a uImage with `mkimage`.
    /// Only useful once the kernel carries the L4T/switchroot patch set.
    Built,
}

/// SD-card path (relative to the card root) that hekate's L4T loader reads.
const BOOT_PREFIX: &str = "switchroot/alarm";
/// Block device holding the root filesystem, as seen by the L4T initramfs.
///
/// The Switch exposes the SD card as `mmcblk0` and the Arch root is partition 2
/// (`mmcblk0p1` is the FAT boot partition). Pinning this stops the switchroot
/// initramfs from probing every block device (and from stopping at the first
/// `switchroot`-labelled partition it finds).
const ROOTDEV: &str = "mmcblk0p2";
/// hekate boot entry that points at `BOOT_PREFIX`.
const INI_NAME: &str = "L4T-alarm.ini";
/// Reference L4T install shipped in this repository.
const SWITCHROOT_REF: &str = "switchroot/ubuntu-jammy";
/// Files that make up the L4T boot prefix, besides the kernel.
const L4T_SUPPORT_FILES: &[&str] = &["bl31.bin", "bl33.bin", "nx-plat.dtimg", "initramfs", "boot.scr"];
/// Load/entry address used by the switchroot L4T kernel image.
const UIMAGE_ADDR: &str = "0x80200000";

/// Build the hekate/L4T boot tree for the Switch.
///
/// Produces, under `output_dir`:
///
/// ```text
/// bootloader/ini/L4T-alarm.ini   # hekate "More configs..." entry
/// switchroot/alarm/uImage        # kernel
/// switchroot/alarm/bl31.bin      # ATF secure monitor
/// switchroot/alarm/bl33.bin      # L4T bootloader stage
/// switchroot/alarm/nx-plat.dtimg # platform device trees
/// switchroot/alarm/initramfs     # root-fs discovery / mount
/// ```
pub fn build_hekate_payload(
    config: &BuildConfig,
    build_dir: &Path,
    output_dir: &Path,
    kernel_source: KernelSource,
) -> Result<()> {
    let project_root = std::env::current_dir().context("failed to resolve project root")?;
    let reference = project_root.join(SWITCHROOT_REF);

    let boot_dir = output_dir.join(BOOT_PREFIX);
    let ini_dir = output_dir.join("bootloader/ini");
    fs::create_dir_all(&boot_dir)
        .with_context(|| format!("failed to create {}", boot_dir.display()))?;
    fs::create_dir_all(&ini_dir)
        .with_context(|| format!("failed to create {}", ini_dir.display()))?;

    // Drop the obsolete, non-functional boot files produced by earlier
    // revisions so they cannot shadow the correct entry on the SD card.
    for stale in [
        "l4t",
        "hekate_ipl.ini",
        "hekate.json",
        "l4t.ini",
        "boot.scr",
        "ALARM_BOOT_FILES.txt",
    ] {
        let path = output_dir.join(stale);
        if path.is_dir() {
            let _ = fs::remove_dir_all(&path);
        } else if path.exists() {
            let _ = fs::remove_file(&path);
        }
    }

    // 0. hekate's own runtime files: sys libs (lp0/minerva), Nyx GUI, res,
    //    payloads, update.bin. Without these hekate refuses to boot L4T with
    //    "missing lp0 (sleep) lib / missing minerva lib / update bootloader
    //    folder!" and Nyx is unavailable.
    let hekate_dir = ensure_hekate_bootloader(build_dir, &config.tegra.bootloader.hekate_version)?;
    let stale_bdk = output_dir.join("bootloader/bdk");
    if stale_bdk.exists() {
        let _ = fs::remove_dir_all(&stale_bdk);
    }
    copy_dir_recursive(&hekate_dir, &output_dir.join("bootloader"))?;

    // 1. hekate entry (written after the bundle so it is not overwritten).
    write_hekate_entry(config, &ini_dir)?;

    // 2. Kernel image.
    match kernel_source {
        KernelSource::Switchroot => {
            let src = reference.join("uImage");
            if !src.exists() {
                bail!(
                    "switchroot kernel not found at {}.\n\
                     Copy the L4T boot files from a switchroot release, or use \
                     `--kernel built` to wrap the locally built kernel.",
                    src.display()
                );
            }
            copy_file(&src, &boot_dir.join("uImage"))?;
            println!("Using switchroot L4T kernel (known-good display support)");
        }
        KernelSource::Built => {
            // Ship both: the raw Image is what the boot script actually uses
            // (`booti`, no decompression), and the uImage is kept for reference
            // and as a fallback for a `bootm`-based script.
            copy_file(&build_dir.join("kernel/Image"), &boot_dir.join("Image"))?;
            make_uimage(&build_dir.join("kernel/Image"), &boot_dir.join("uImage"))?;
            println!(
                "Copied the raw kernel Image and wrapped it into uImage \
                 (the boot script boots the Image with `booti`)"
            );
        }
    }

    // 3. Support files (ATF, bootloader stage, device trees, initramfs).
    for name in L4T_SUPPORT_FILES {
        let src = reference.join(name);
        if !src.exists() {
            bail!(
                "required L4T boot file '{}' missing at {}.\n\
                 These files come from a switchroot L4T release; place them in {}.",
                name,
                src.display(),
                reference.display()
            );
        }
        copy_file(&src, &boot_dir.join(name))?;
    }

    // 3b. Mainline boot set. The L4T `boot.scr` loads a Device Tree *image*
    //     (nx-plat.dtimg), picks an entry by SKU and then validates it against
    //     ${fdt_id_text}/${fdt_rev} — properties only L4T device trees carry.
    //     On a mainline DTB that check fails and the script reboots in a loop,
    //     so for a locally built kernel swap in our own script and a plain dtb.
    if kernel_source == KernelSource::Built {
        let scr = project_root.join("device-tree/boot-mainline.scr");
        if !scr.exists() {
            bail!(
                "mainline boot script not found at {}.\n\
                 Build it with:\n  \
                 mkimage -A arm64 -O linux -T script -C none -n \"ALARM mainline boot\" \
                 -d device-tree/boot-mainline.cmd device-tree/boot-mainline.scr",
                scr.display()
            );
        }
        copy_file(&scr, &boot_dir.join("boot.scr"))?;
        println!("Using the mainline boot script (boot.scr)");

        let dtb = build_dir.join("device-tree/tegra210-nintendo-switch.dtb");
        if !dtb.exists() {
            bail!(
                "Switch device tree not found at {}.\n\
                 Run `alarm-build kernel` first: install_switch_dt() injects the \
                 DTS into the kernel tree so `make dtbs` emits it.",
                dtb.display()
            );
        }
        copy_file(&dtb, &boot_dir.join("switch.dtb"))?;
        println!("Using the mainline Switch device tree (switch.dtb)");
    }

    // 4. Bootlogo / menu icon (optional).
    copy_optional(&reference.join("bootlogo_ubuntu.bmp"), &boot_dir.join("bootlogo.bmp"))?;
    copy_optional(&reference.join("icon_ubuntu_hue.bmp"), &boot_dir.join("icon.bmp"))?;

    println!("L4T boot tree created at: {}", output_dir.display());
    println!("  entry   : bootloader/ini/{}", INI_NAME);
    println!("  prefix  : /{}", BOOT_PREFIX);
    Ok(())
}

/// Write `bootloader/ini/L4T-alarm.ini` in hekate's L4T schema.
///
/// See `switchroot/ubuntu-jammy/README_CONFIG.txt` for the full list of keys.
fn write_hekate_entry(config: &BuildConfig, ini_dir: &Path) -> Result<()> {
    let path = ini_dir.join(INI_NAME);
    fs::write(&path, hekate_entry(config))
        .with_context(|| format!("failed to write {}", path.display()))?;
    Ok(())
}

/// Render the hekate L4T entry. Split out so it can be unit tested without
/// touching the filesystem.
///
/// Keys:
/// * `l4t=1` / `boot_prefixes` / `id` — tell hekate's L4T loader where to look.
/// * `rootdev=mmcblk0p2` — pin the root filesystem instead of letting the
///   initramfs auto-detect it.
/// * `rootlabel_retries=1` — with `rootdev` given, a single label lookup is
///   enough; do not spend seconds retrying.
/// * `r2p_action=self` — reboot-to-payload returns to hekate.
/// * `fbconsole=0` — leave the framebuffer alone; the 4.9 L4T kernel drives the
///   DSI panel itself.
/// * `bootargs_extra` — appended to the kernel command line:
///   - `usb_logging` creates the USB serial gadget (`/dev/ttyGS0`); this is the
///     only console available without extra hardware, and the init script uses
///     it;
///   - `auto_rootdev_disable` stops the initramfs from overriding `rootdev`.
///
/// Keep this comment-free: hekate's ini parser treats a `#` line as a comment,
/// but a `;` line inside a section is parsed as a key/value pair and forwarded
/// into U-Boot's environment, which can break the boot script.
fn hekate_entry(config: &BuildConfig) -> String {
    format!(
        "[L4T ALARM]\n\
         l4t=1\n\
         boot_prefixes=/{prefix}/\n\
         id={label}\n\
         rootdev={rootdev}\n\
         rootlabel_retries=1\n\
         r2p_action=self\n\
         fbconsole=0\n\
         bootargs_extra=usb_logging auto_rootdev_disable\n",
        prefix = BOOT_PREFIX,
        label = config.image.label,
        rootdev = ROOTDEV,
    )
}

/// Wrap a raw arm64 `Image` into a gzip-compressed uImage for hekate.
fn make_uimage(kernel: &Path, dest: &Path) -> Result<()> {
    if !kernel.exists() {
        bail!(
            "kernel image not found at {}.\nRun `alarm-build kernel` first, or use \
             `--kernel switchroot`.",
            kernel.display()
        );
    }

    let gz = dest.with_extension("gz");
    // gzip -9 -n: deterministic output, maximum compression.
    let file = fs::File::create(&gz)
        .with_context(|| format!("failed to create {}", gz.display()))?;
    let status = Command::new("gzip")
        .args(["-9", "-n", "-c"])
        .arg(kernel)
        .stdout(file)
        .status()
        .context("failed to run gzip (is it installed?)")?;
    if !status.success() {
        bail!("gzip failed for {}", kernel.display());
    }

    let status = Command::new("mkimage")
        .args([
            "-A", "arm64",
            "-O", "linux",
            "-T", "kernel",
            "-C", "gzip",
            "-a", UIMAGE_ADDR,
            "-e", UIMAGE_ADDR,
            "-n", "ALARM L4T",
            "-d",
        ])
        .arg(&gz)
        .arg(dest)
        .status()
        .context("failed to run mkimage (install the 'u-boot-tools' package)")?;
    let _ = fs::remove_file(&gz);
    if !status.success() {
        bail!("mkimage failed");
    }
    Ok(())
}

fn copy_file(src: &Path, dst: &Path) -> Result<()> {
    fs::copy(src, dst)
        .with_context(|| format!("failed to copy {} -> {}", src.display(), dst.display()))?;
    Ok(())
}

fn copy_dir_recursive(src: &Path, dst: &Path) -> Result<()> {
    fs::create_dir_all(dst)?;
    for entry in fs::read_dir(src).with_context(|| format!("failed to read {}", src.display()))? {
        let entry = entry?;
        let src_path = entry.path();
        let dst_path = dst.join(entry.file_name());
        if src_path.is_dir() {
            copy_dir_recursive(&src_path, &dst_path)?;
        } else if src_path.is_file() {
            copy_file(&src_path, &dst_path)?;
        }
    }
    Ok(())
}

/// Official hekate release API.
const HEKATE_API: &str = "https://api.github.com/repos/CTCaer/hekate/releases";

/// Locate (or download) a complete hekate `bootloader/` folder.
///
/// hekate loads `bootloader/sys/libsys_lp0.bso` (sleep) and
/// `bootloader/sys/libsys_minerva.bso` (RAM config) from the SD card at boot,
/// and the Nyx menu lives in `bootloader/nyx/`.
fn ensure_hekate_bootloader(build_dir: &Path, version: &str) -> Result<PathBuf> {
    // 1. Explicit override: a directory containing `sys/` (the bootloader dir
    //    itself, or a parent that contains `bootloader/`).
    if let Some(dir) = std::env::var_os("HEKATE_BOOTLOADER_DIR") {
        let dir = PathBuf::from(dir);
        for candidate in [dir.clone(), dir.join("bootloader")] {
            if candidate.join("sys").is_dir() {
                return Ok(candidate);
            }
        }
        bail!(
            "HEKATE_BOOTLOADER_DIR does not contain a 'sys' directory: {}",
            dir.display()
        );
    }

    // 2. Bundled resources next to the project.
    let project_root = std::env::current_dir().context("failed to resolve project root")?;
    for candidate in [
        project_root.join("resources/hekate/bootloader"),
        project_root.join("resources/bootloader"),
        project_root.join("bootloader"),
    ] {
        if candidate.join("sys").is_dir() {
            return Ok(candidate);
        }
    }

    // 3. A previously downloaded copy.
    let cached = build_dir.join("hekate/bootloader");
    if cached.join("sys").is_dir() {
        return Ok(cached);
    }

    // 4. Download the official release bundle.
    download_hekate_bootloader(build_dir, version)?;
    if cached.join("sys").is_dir() {
        Ok(cached)
    } else {
        bail!(
            "hekate download did not contain bootloader/sys/*.bso (looked in {})",
            cached.display()
        )
    }
}

/// Download and extract the hekate release zip into `build_dir/hekate`.
fn download_hekate_bootloader(build_dir: &Path, version: &str) -> Result<()> {
    let dest = build_dir.join("hekate");
    fs::create_dir_all(&dest)?;

    let api = if version.trim().is_empty() {
        format!("{HEKATE_API}/latest")
    } else {
        format!("{HEKATE_API}/tags/v{}", version.trim_start_matches('v'))
    };
    println!("Fetching hekate release info from {}", api);
    let json = http_get(&api)?;
    let zip_url = select_zip_asset(&json)
        .with_context(|| format!("no hekate_ctcaer_*.zip asset found for {api}"))?;

    let zip = dest.join("hekate.zip");
    println!("Downloading hekate bootloader bundle...");
    download_file(&zip_url, &zip)?;
    extract_zip(&zip, &dest)
}

/// Find the `hekate_ctcaer_*.zip` download URL in a release JSON body.
fn select_zip_asset(json: &str) -> Option<String> {
    let value: serde_json::Value = serde_json::from_str(json).ok()?;
    value
        .get("assets")?
        .as_array()?
        .iter()
        .filter_map(|asset| {
            let name = asset.get("name")?.as_str()?;
            if name.starts_with("hekate_ctcaer") && name.ends_with(".zip") {
                asset
                    .get("browser_download_url")?
                    .as_str()
                    .map(str::to_string)
            } else {
                None
            }
        })
        .next()
}

fn http_get(url: &str) -> Result<String> {
    let output = if which::which("curl").is_ok() {
        Command::new("curl")
            .args(["-fsSL", "-H", "User-Agent: alarm-build"])
            .arg(url)
            .output()
            .context("failed to run curl")?
    } else if which::which("wget").is_ok() {
        Command::new("wget")
            .args(["-qO-"])
            .arg(url)
            .output()
            .context("failed to run wget")?
    } else {
        bail!("neither curl nor wget is installed; cannot fetch {}", url);
    };
    if !output.status.success() {
        bail!("failed to fetch {} (is there network access?)", url);
    }
    Ok(String::from_utf8_lossy(&output.stdout).into_owned())
}

fn download_file(url: &str, dest: &Path) -> Result<()> {
    let status = if which::which("curl").is_ok() {
        Command::new("curl")
            .args(["-fL", "--retry", "3", "-o"])
            .arg(dest)
            .arg(url)
            .status()
    } else {
        Command::new("wget")
            .arg("-O")
            .arg(dest)
            .arg(url)
            .status()
    }
    .context("failed to run downloader")?;
    if !status.success() {
        bail!("failed to download {}", url);
    }
    Ok(())
}

fn extract_zip(zip: &Path, dest: &Path) -> Result<()> {
    for tool in ["bsdtar", "unzip", "7z"] {
        if which::which(tool).is_err() {
            continue;
        }
        let status = match tool {
            "bsdtar" => Command::new(tool)
                .arg("-xf")
                .arg(zip)
                .arg("-C")
                .arg(dest)
                .status(),
            "unzip" => Command::new(tool)
                .args(["-oq"])
                .arg(zip)
                .arg("-d")
                .arg(dest)
                .status(),
            _ => Command::new(tool)
                .args(["x", "-y"])
                .arg(format!("-o{}", dest.display()))
                .arg(zip)
                .status(),
        }
        .with_context(|| format!("failed to run {}", tool))?;
        if status.success() {
            return Ok(());
        }
    }
    bail!("no working zip extractor found (install unzip, bsdtar, or 7z)")
}

fn copy_optional(src: &Path, dst: &Path) -> Result<()> {
    if src.exists() {
        copy_file(src, dst)?;
    }
    Ok(())
}

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

    #[test]
    fn entry_uses_hekate_l4t_schema() {
        let config = BuildConfig::default_for_odnx02_a2();
        let entry = hekate_entry(&config);

        // These keys are what make hekate's L4T loader find and boot the
        // kernel; without them the screen stays black.
        assert!(entry.contains("l4t=1"), "missing l4t=1:\n{entry}");
        assert!(
            entry.contains("boot_prefixes=/switchroot/alarm/"),
            "missing boot_prefixes:\n{entry}"
        );
        assert!(entry.contains("id=ALARM"), "missing id:\n{entry}");
        assert!(
            entry.contains("[L4T ALARM]"),
            "missing entry header:\n{entry}"
        );
        // The root device must be pinned, and the USB serial gadget requested:
        // without rootdev the initramfs probes every block device, without
        // usb_logging there is no serial console.
        assert!(entry.contains("rootdev=mmcblk0p2"), "missing rootdev:\n{entry}");
        assert!(
            entry.contains("usb_logging"),
            "missing usb_logging in bootargs_extra:\n{entry}"
        );
        assert!(
            entry.contains("auto_rootdev_disable"),
            "missing auto_rootdev_disable:\n{entry}"
        );
        // Kernel command line keys must not leak into the ini.
        assert!(!entry.contains("console=tty"), "leaked kernel args:\n{entry}");
        assert!(!entry.contains("wakeetrae2"), "leaked junk args:\n{entry}");
    }
}