mod config;
mod kernel;
mod rootfs;
mod packages;
mod payload;
mod utils;

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use colored::Colorize;
use std::path::PathBuf;

#[derive(Parser)]
#[command(name = "alarm-build")]
#[command(about = "ALARM build system for NVIDIA Tegra X1 (ODNX02-A2)")]
#[command(version)]
struct Cli {
    /// Build directory (default: ./build)
    #[arg(short, long, default_value = "./build")]
    build_dir: PathBuf,

    /// Output directory for images
    #[arg(short, long, default_value = "./output")]
    output_dir: PathBuf,

    /// Number of parallel jobs
    #[arg(short, long)]
    jobs: Option<u32>,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Initialize build environment and download sources
    Init {
        /// Skip downloading kernel source
        #[arg(long)]
        skip_kernel: bool,

        /// Skip downloading U-Boot source
        #[arg(long)]
        skip_uboot: bool,
    },

    /// Build the cross-compilation toolchain
    Toolchain,

    /// Build the Linux kernel
    Kernel {
        /// Use custom kernel config
        #[arg(long)]
        custom_config: Option<PathBuf>,
    },

    /// Build U-Boot bootloader
    Uboot,

    /// Build device tree overlays
    DeviceTree,

    /// Create root filesystem
    Rootfs {
        /// Include desktop environment
        #[arg(long)]
        desktop: bool,

        /// Extra packages to install (comma-separated)
        #[arg(long)]
        packages: Option<String>,
    },

    /// Build hekate L4T boot files (bootloader/ini + switchroot/alarm)
    Payload {
        /// Kernel to boot
        #[arg(long, value_enum, default_value_t = payload::KernelSource::Switchroot)]
        kernel: payload::KernelSource,
    },

    /// Build complete image (for SD card boot)
    Image {
        /// Image size in GB
        #[arg(short, long)]
        size: Option<u32>,

        /// Include desktop environment
        #[arg(long)]
        desktop: bool,

        /// Kernel to boot
        #[arg(long, value_enum, default_value_t = payload::KernelSource::Switchroot)]
        kernel: payload::KernelSource,
    },

    /// Flash image to SD card or USB
    Flash {
        /// Device path (e.g., /dev/sdX)
        device: PathBuf,
    },

    /// Clean build artifacts
    Clean {
        /// Remove all build directories
        #[arg(long)]
        all: bool,
    },

    /// Show build configuration
    Config,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    // Create directories
    std::fs::create_dir_all(&cli.build_dir)
        .context("Failed to create build directory")?;
    std::fs::create_dir_all(&cli.output_dir)
        .context("Failed to create output directory")?;

    match &cli.command {
        Commands::Init { skip_kernel, skip_uboot } => {
            println!("{}", "Initializing build environment...".green().bold());
            config::init_build_env(&cli.build_dir)?;
            
            if !skip_kernel {
                println!("{}", "Downloading Linux kernel...".cyan());
                kernel::download_source(&cli.build_dir)?;
            }
            
            if !skip_uboot {
                println!("{}", "Downloading U-Boot...".cyan());
                packages::download_uboot(&cli.build_dir)?;
            }
            
            println!("{}", "Initialization complete!".green().bold());
        }

        Commands::Toolchain => {
            println!("{}", "Building cross-compilation toolchain...".green().bold());
            utils::build_toolchain(&cli.build_dir)?;
        }

        Commands::Kernel { custom_config } => {
            println!("{}", "Building Linux kernel...".green().bold());
            let config = config::load_config(&cli.build_dir)?;
            kernel::build(&config, &cli.build_dir, custom_config.as_ref())?;
        }

        Commands::Payload { kernel } => {
            println!("{}", "Building hekate L4T boot files...".green().bold());
            let config = config::load_config(&cli.build_dir)?;
            payload::build_hekate_payload(&config, &cli.build_dir, &cli.output_dir, *kernel)?;
        }

        Commands::Uboot => {
            println!("{}", "Building U-Boot...".green().bold());
            let config = config::load_config(&cli.build_dir)?;
            packages::build_uboot(&config, &cli.build_dir)?;
        }

        Commands::DeviceTree => {
            println!("{}", "Building device tree...".green().bold());
            kernel::build_device_tree(&cli.build_dir)?;
        }

        Commands::Rootfs { desktop, packages: extra_pkgs } => {
            println!("{}", "Creating root filesystem...".green().bold());
            let config = config::load_config(&cli.build_dir)?;
            rootfs::create(&config, &cli.build_dir, *desktop, extra_pkgs.as_deref())?;
        }

        Commands::Image { size, desktop, kernel } => {
            println!("{}", "Building complete image...".green().bold());
            let config = config::load_config(&cli.build_dir)?;

            // Remove any previous image first: if the build fails part-way we
            // must not leave a stale image that `flash` would happily write.
            let image = cli.output_dir.join("alarm-tegra-x1.img");
            let _ = std::fs::remove_file(&image);

            // Build all components. The kernel and fallback device tree are
            // only needed when packaging the locally built kernel.
            if *kernel == payload::KernelSource::Built {
                kernel::build(&config, &cli.build_dir, None)?;
                kernel::build_device_tree(&cli.build_dir)?;
            }
            payload::build_hekate_payload(&config, &cli.build_dir, &cli.output_dir, *kernel)?;
            rootfs::create(&config, &cli.build_dir, *desktop, None)?;
            
            // Create final image for SD card boot
            let size_gb = size.unwrap_or(config.image.size_gb);
            rootfs::create_image(&config, &cli.build_dir, &cli.output_dir, size_gb)?;
            
            println!("{}", "Build complete!".green().bold());
            println!("Image:   {:?}", cli.output_dir.join("alarm-tegra-x1.img"));
            println!("\n{}", "To boot:".yellow().bold());
            println!("  1. Flash the image: sudo alarm-build flash /dev/sdX");
            println!("  2. Insert the SD card and power on into hekate");
            println!("  3. Select 'More configs...' -> 'L4T ALARM'");
        }

        Commands::Flash { device } => {
            println!("{}", "Flashing image...".green().bold());
            let image_path = cli.output_dir.join("alarm-tegra-x1.img");
            if !image_path.exists() {
                anyhow::bail!("Image not found at {:?}. Run 'image' command first.", image_path);
            }
            utils::flash_image(&image_path, device)?;
        }

        Commands::Clean { all } => {
            if *all {
                println!("{}", "Removing all build artifacts...".yellow());
                if cli.build_dir.exists() {
                    std::fs::remove_dir_all(&cli.build_dir)?;
                }
                if cli.output_dir.exists() {
                    std::fs::remove_dir_all(&cli.output_dir)?;
                }
            } else {
                println!("{}", "Cleaning build directory...".yellow());
                // Only remove specific build artifacts
                let _ = std::fs::remove_dir_all(cli.build_dir.join("kernel"));
                let _ = std::fs::remove_dir_all(cli.build_dir.join("uboot"));
            }
            println!("{}", "Clean complete!".green().bold());
        }

        Commands::Config => {
            config::show_config()?;
        }
    }

    // When we re-executed under sudo, give the generated files back to the
    // invoking user so subsequent non-root builds can overwrite them.
    utils::restore_ownership(std::slice::from_ref(&cli.output_dir));

    Ok(())
}