// Unlicense — cochranblock.org
// Contributors: GotEmCoach, KOVA, Claude Opus 4.6
//! pixel-forge — Pixel art game asset generator.
//! Stardew Valley / Starbound quality. Local-first, Rust-native.

use pixel_forge::*;
use clap::{Parser, Subcommand};

// Federal compliance docs — baked into the binary via include_str!
const GOVDOC_SBOM: &str = include_str!("../govdocs/SBOM.md");
const GOVDOC_SECURITY: &str = include_str!("../govdocs/SECURITY.md");
const GOVDOC_SSDF: &str = include_str!("../govdocs/SSDF.md");
const GOVDOC_FIPS: &str = include_str!("../govdocs/FIPS.md");
const GOVDOC_CMMC: &str = include_str!("../govdocs/CMMC.md");
const GOVDOC_PRIVACY: &str = include_str!("../govdocs/PRIVACY.md");
const GOVDOC_SUPPLY_CHAIN: &str = include_str!("../govdocs/SUPPLY_CHAIN.md");
const GOVDOC_ACCESSIBILITY: &str = include_str!("../govdocs/ACCESSIBILITY.md");
const GOVDOC_FEDRAMP: &str = include_str!("../govdocs/FedRAMP_NOTES.md");
const GOVDOC_ITAR_EAR: &str = include_str!("../govdocs/ITAR_EAR.md");
const GOVDOC_FEDERAL_USE: &str = include_str!("../govdocs/FEDERAL_USE_CASES.md");

const BASE_CLASSES: &[&str] = &[
    "character", "weapon", "potion", "terrain", "enemy",
    "tree", "building", "animal", "effect", "food",
    "armor", "tool", "vehicle", "ui", "misc",
];

/// Discover class dirs: known classes + any extra dirs found on disk (skip _ and . prefixed).
fn discover_classes(data_path: &std::path::Path) -> Vec<String> {
    let mut names: Vec<String> = BASE_CLASSES.iter().map(|s| s.to_string()).collect();
    if let Ok(entries) = std::fs::read_dir(data_path) {
        for entry in entries.flatten() {
            let name = entry.file_name().to_string_lossy().to_string();
            if name.starts_with('_') || name.starts_with('.') { continue; }
            if entry.path().is_dir() && !names.contains(&name) {
                names.push(name);
            }
        }
    }
    names
}


#[derive(Parser)]
#[command(name = "pixel-forge", about = "Pixel art game asset generator")]
struct Cli {
    #[command(subcommand)]
    cmd: Option<Cmd>,
    /// Dump machine-readable SBOM (SPDX format) and exit.
    #[arg(long)]
    sbom: bool,
}

#[derive(Subcommand)]
enum Cmd {
    /// Generate a sprite from a text description.
    Sprite {
        /// What to generate: "knight idle", "potion red", "tree oak"
        prompt: String,
        /// Output size in pixels (square). Default 32.
        #[arg(short, long, default_value_t = 32)]
        size: u32,
        /// Number of animation frames. 1 = static sprite.
        #[arg(short, long, default_value_t = 1)]
        frames: u32,
        /// Palette: stardew, starbound, snes, nes, gameboy, or custom .pal path.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "sprite.png")]
        output: String,
    },
    /// Generate a seamless tileset.
    Tileset {
        /// Theme: "dungeon", "forest", "cave", "village", "space"
        theme: String,
        /// Tile size in pixels (square). Default 16.
        #[arg(short, long, default_value_t = 32)]
        size: u32,
        /// Number of tile variations. Default 16.
        #[arg(short, long, default_value_t = 16)]
        count: u32,
        /// Palette name or path.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "tileset.png")]
        output: String,
    },
    /// List available palettes.
    Palettes,
    /// Download/cache the diffusion model.
    Setup,
    /// Train a tiny pixel art diffusion model from a dataset of PNGs.
    /// Dataset: directory of PNGs, optionally organized by class subdirectory.
    Train {
        /// Path to training data directory (PNGs, optionally in class subdirs).
        #[arg(short, long, default_value = "data")]
        data: String,
        /// Output model file path (.safetensors).
        #[arg(short, long, default_value = "pixel-forge-cinder.safetensors")]
        output: String,
        /// Number of training epochs.
        #[arg(short, long, default_value_t = 100)]
        epochs: usize,
        /// Batch size.
        #[arg(short, long, default_value_t = 64)]
        batch_size: usize,
        /// Learning rate.
        #[arg(short, long, default_value_t = 1e-3)]
        lr: f64,
        /// Image size (square). Must match training data.
        #[arg(long, default_value_t = 32)]
        img_size: u32,
        /// Use MediumUNet (~5.8M params) instead of TinyUNet (~1.1M params).
        #[arg(long)]
        medium: bool,
        /// Use AnvilUNet (~16M params) — the XL model.
        #[arg(long)]
        anvil: bool,
        /// Disable EMA (saves ~50% RAM during training).
        #[arg(long)]
        no_ema: bool,
        /// Train with v-prediction (model predicts velocity instead of clean image).
        #[arg(long)]
        v_pred: bool,
        /// Minimum learning rate for cosine LR decay. 0 = flat LR.
        #[arg(long, default_value_t = 1e-5)]
        lr_min: f64,
        /// Warm-up epochs: linear ramp from lr_min to lr.
        #[arg(long, default_value_t = 5)]
        warmup: usize,
        /// Conditioning data dir for stage-aware training (silhouettes, colorblocks).
        #[arg(long)]
        condition: Option<String>,
        /// Mixed precision: f16 forward pass, f32 optimizer. ~2x faster on CUDA.
        #[arg(long)]
        fp16: bool,
        /// Save checkpoint every N epochs. 1 = every epoch (overwrites same file).
        #[arg(long, default_value_t = 25)]
        checkpoint_every: usize,
        /// Resume from existing safetensors checkpoint.
        #[arg(long)]
        resume: Option<String>,
        /// Sponge Mesh: auto-retry on NaN/plateau (max 3 retries).
        #[arg(long)]
        sponge: bool,
        /// GPU backend: auto, metal, cuda, vulkan, cpu.
        #[arg(long, default_value = "auto")]
        device: String,
    },
    /// Ingest Gemini-generated sprite sheets into training data.
    /// Slices grids, removes backgrounds, downscales to 32x32, sorts by filename.
    IngestGemini {
        /// Path to directory of Gemini output PNGs (named CLASS_batchN.png).
        #[arg(short, long, default_value = "data/raw/gemini")]
        input: String,
        /// Output training data directory.
        #[arg(short, long, default_value = "data_v2_32")]
        output: String,
    },
    /// Curate raw downloaded datasets into class-sorted training directories.
    Curate {
        /// Path to raw downloads directory.
        #[arg(short, long, default_value = "data/raw")]
        raw: String,
        /// Path to output class directories.
        #[arg(short, long, default_value = "data")]
        output: String,
        /// Target tile size (square).
        #[arg(long, default_value_t = 32)]
        size: u32,
        /// Augment classes with fewer than this many tiles by adding h-flip variants.
        #[arg(long, default_value_t = 0)]
        augment_below: usize,
    },
    /// Precompute per-class skeleton images from training data.
    /// Skeletons give the sampler a head start — 70% structure, 30% noise.
    Skeletons {
        /// Path to training data directory.
        #[arg(short, long, default_value = "data")]
        data: String,
        /// Image size (must match training data).
        #[arg(long, default_value_t = 32)]
        img_size: u32,
    },
    /// Extract structural layers for stage-aware cascade.
    /// Generates 5-channel conditioning data: SDF, normal X/Y, outline, luminance.
    PrepStages {
        /// Path to training data directory with class subdirs.
        #[arg(short, long, default_value = "data_v2_32")]
        data: String,
        /// Number of palette colors for color block-in stage (unused, kept for compat).
        #[arg(long, default_value_t = 8)]
        palette_colors: usize,
    },
    /// Stage cascade: Cinder-sil → structure map → Quench-detail → sprite.
    StageCascade {
        /// Class to generate.
        class: String,
        /// Cinder-sil model (silhouette generator).
        #[arg(long, default_value = "pixel-forge-cinder-sil.safetensors")]
        sil_model: String,
        /// Quench-detail model (6ch conditioned detail filler).
        #[arg(long, default_value = "pixel-forge-quench-detail.safetensors")]
        detail_model: String,
        /// Number of sprites.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Silhouette generation steps.
        #[arg(long, default_value_t = 20)]
        sil_steps: usize,
        /// Detail generation steps.
        #[arg(long, default_value_t = 30)]
        detail_steps: usize,
        /// Palette.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file.
        #[arg(short, long, default_value = "stage-cascade.png")]
        output: String,
        /// Generate 4-directional sprite sheet via relighting.
        #[arg(long)]
        four_dir: bool,
        /// Skip silhouette generation, use real structure maps from this dir.
        #[arg(long)]
        real_struct: Option<String>,
    },
    /// Relight a sprite into 4-directional views (front/left/right/back).
    /// Uses SDF + normal maps — no model needed, pure math.
    Relight {
        /// Input sprite PNG.
        image: String,
        /// Output sprite sheet (4 views side by side).
        #[arg(short, long, default_value = "relit.png")]
        output: String,
        /// Also save individual direction PNGs (front.png, left.png, etc).
        #[arg(long)]
        split: bool,
    },
    /// Generate pixel art using a trained model (no SD required).
    Generate {
        /// Class to generate: character, weapon, potion, terrain, enemy, etc.
        class: String,
        /// Path to trained model (.safetensors).
        #[arg(short, long, default_value = "pixel-forge-cinder.safetensors")]
        model: String,
        /// Output size in pixels (must match training size).
        #[arg(short, long, default_value_t = 32)]
        size: u32,
        /// Number of images to generate.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Denoising steps (more = higher quality, slower).
        #[arg(long, default_value_t = 40)]
        steps: usize,
        /// Palette for color quantization.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "generated.png")]
        output: String,
        /// Use MediumUNet (Quench) model architecture.
        #[arg(long)]
        medium: bool,
        /// Seed for deterministic generation. Same seed = same sprite.
        #[arg(long)]
        seed: Option<u64>,
        /// Generate 4-directional sprite sheet (front/left/right/back) via relighting.
        #[arg(long)]
        four_dir: bool,
        /// CFG guidance scale. Higher = stronger class adherence.
        #[arg(long, default_value_t = 3.0)]
        cfg: f64,
    },
    /// Record a swipe (good/bad) on a generated sprite for Judge training.
    Swipe {
        /// Path to the sprite image (32×32 PNG).
        image: String,
        /// "good" or "bad" (swipe right or left).
        verdict: String,
        /// Class of the sprite.
        #[arg(short, long, default_value = "misc")]
        class: String,
        /// Path to swipe store file.
        #[arg(long, default_value = "swipes.bin")]
        store: String,
    },
    /// Train the Judge model from recorded swipes.
    TrainJudge {
        /// Path to swipe store file.
        #[arg(short, long, default_value = "swipes.bin")]
        store: String,
        /// Output Judge model path.
        #[arg(short, long, default_value = "judge.safetensors")]
        output: String,
    },
    /// Score sprites through a trained Judge — filter good from bad.
    Judge {
        /// Path to sprite images (or directory of PNGs).
        input: String,
        /// Path to trained Judge model.
        #[arg(short, long, default_value = "judge.safetensors")]
        model: String,
    },
    /// Generate a scene: 8×8 grid of sprites (128×128 composite).
    Scene {
        /// Generation mode: "bootstrap" (rule-seeded), "model" (trained Combiner), "seed" (user-seeded).
        #[arg(short, long, default_value = "bootstrap")]
        mode: String,
        /// Path to Combiner model (for mode=model).
        #[arg(short = 'M', long, default_value = "combiner.safetensors")]
        combiner_model: String,
        /// Biome for bootstrap: forest, dungeon, village, cave, plains.
        #[arg(short, long, default_value = "forest")]
        biome: String,
        /// Sampling temperature (lower = more deterministic).
        #[arg(short, long, default_value_t = 0.8)]
        temperature: f32,
        /// Number of scenes to generate.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Output file path.
        #[arg(short, long, default_value = "scene.png")]
        output: String,
    },
    /// Train the Combiner model on rule-seeded or user-accepted scenes.
    TrainCombiner {
        /// Number of bootstrap scenes to generate for training.
        #[arg(short, long, default_value_t = 500)]
        num_scenes: usize,
        /// Training epochs.
        #[arg(short, long, default_value_t = 20)]
        epochs: usize,
        /// Output model path.
        #[arg(short, long, default_value = "combiner.safetensors")]
        output: String,
    },
    /// Update Generator LoRA from Judge feedback.
    TrainLora {
        /// Path to base Generator model.
        #[arg(short, long, default_value = "pixel-forge-cinder.safetensors")]
        model: String,
        /// Path to trained Judge model.
        #[arg(short, long, default_value = "judge.safetensors")]
        judge: String,
        /// Output LoRA weights path.
        #[arg(short, long, default_value = "generator-lora.safetensors")]
        output: String,
        /// Class to generate during LoRA training.
        #[arg(short, long, default_value = "character")]
        class: String,
    },
    /// Show swipe store stats.
    SwipeStats {
        /// Path to swipe store file.
        #[arg(short, long, default_value = "swipes.bin")]
        store: String,
    },
    /// Full pipeline: generate → judge → combine → render scene.
    Pipeline {
        /// Path to Generator model.
        #[arg(short, long, default_value = "pixel-forge-cinder.safetensors")]
        model: String,
        /// Path to Judge model (skip judging if missing).
        #[arg(short, long, default_value = "judge.safetensors")]
        judge: String,
        /// Path to Combiner model (use bootstrap if missing).
        #[arg(short = 'M', long, default_value = "combiner.safetensors")]
        combiner: String,
        /// Path to LoRA weights (skip if missing).
        #[arg(short, long, default_value = "generator-lora.safetensors")]
        lora: String,
        /// Palette for color quantization.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output scene image path.
        #[arg(short, long, default_value = "scene.png")]
        output: String,
    },
    /// Train expert heads on frozen Quench base.
    TrainExperts {
        /// Path to frozen Quench model.
        #[arg(long, default_value = "pixel-forge-quench.safetensors")]
        quench: String,
        /// Training data directory.
        #[arg(short, long, default_value = "data")]
        data: String,
        /// Output experts file.
        #[arg(short, long, default_value = "experts.safetensors")]
        output: String,
        /// Epochs.
        #[arg(short, long, default_value_t = 20)]
        epochs: usize,
        /// Batch size.
        #[arg(short, long, default_value_t = 32)]
        batch_size: usize,
    },
    /// Auto-generate: detect device, pick best model, generate.
    Auto {
        /// Class to generate: character, weapon, potion, terrain, enemy, etc.
        class: String,
        /// Output size in pixels (must match training size).
        #[arg(short, long, default_value_t = 32)]
        size: u32,
        /// Number of images to generate.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Denoising steps.
        #[arg(long, default_value_t = 40)]
        steps: usize,
        /// Palette for color quantization.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "generated.png")]
        output: String,
    },
    /// Re-probe device capabilities (clears cached profile).
    Probe {
        /// Output as JSON (for cluster probing).
        #[arg(long)]
        json: bool,
    },
    /// Probe all nodes in the forge cluster.
    ClusterProbe,
    /// Deploy pixel-forge binary to all nodes (sync + build).
    ClusterDeploy,
    /// Sync trained models to all nodes.
    ClusterSync,
    /// Generate across the full cluster.
    ClusterGenerate {
        /// Class to generate.
        class: String,
        /// Total sprites to generate (distributed across nodes).
        #[arg(short, long, default_value_t = 16)]
        count: u32,
        /// Denoising steps.
        #[arg(long, default_value_t = 40)]
        steps: usize,
        /// Palette.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "cluster-output.png")]
        output: String,
    },
    /// GPU scheduling — delegates to kova c2 gpu. Pass args after --.
    Gpu {
        /// Arguments forwarded to `kova c2 gpu`.
        args: Vec<String>,
    },
    /// Forge: generate → discriminator gate → PoA sign → Ghost Fabric packet.
    Forge {
        /// Class to generate.
        class: String,
        /// Number of signed sprites to produce.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Denoising steps.
        #[arg(long, default_value_t = 40)]
        steps: usize,
        /// Discriminator quality threshold (0.94 ≈ loss < 0.06).
        #[arg(long, default_value_t = 0.94)]
        threshold: f32,
        /// Max generation attempts before giving up.
        #[arg(long, default_value_t = 50)]
        max_attempts: u32,
        /// Path to discriminator model.
        #[arg(long, default_value = "discriminator.safetensors")]
        disc: String,
        /// Palette for color quantization.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file path.
        #[arg(short, long, default_value = "forged.png")]
        output: String,
        /// GPS latitude (default: Dundalk, MD).
        #[arg(long, default_value_t = 39.2504)]
        lat: f32,
        /// GPS longitude (default: Dundalk, MD).
        #[arg(long, default_value_t = -76.5205)]
        lon: f32,
        /// Use MoE cascade (Quench → Cinder + Experts) instead of auto-detect.
        #[arg(long)]
        cascade: bool,
    },
    /// Sign an existing sprite with PoA — produce a Ghost Fabric packet.
    Sign {
        /// Path to 32x32 PNG sprite.
        image: String,
        /// Class of the sprite.
        #[arg(short, long, default_value = "misc")]
        class: String,
        /// Quality score override (0.0-1.0). If omitted, runs discriminator.
        #[arg(long)]
        score: Option<f32>,
        /// Path to discriminator model (used if --score not set).
        #[arg(long, default_value = "discriminator.safetensors")]
        disc: String,
        /// GPS latitude.
        #[arg(long, default_value_t = 39.2504)]
        lat: f32,
        /// GPS longitude.
        #[arg(long, default_value_t = -76.5205)]
        lon: f32,
    },
    /// Show node public key (for PoA verification).
    NodeKey,
    /// Plugin mode — JSON request/response for kova integration.
    Plugin {
        /// Keep process alive, read JSON lines continuously.
        #[arg(long)]
        r#loop: bool,
    },
    /// Quantize a model from f32 to f16 (halves file size + faster on ARM/Metal).
    Quantize {
        /// Input model file (f32 safetensors).
        input: String,
        /// Output file. Defaults to input-f16.safetensors.
        #[arg(short, long)]
        output: Option<String>,
    },
    /// MoE cascade: Quench foundation → Cinder detail + Experts.
    Cascade {
        /// Class to generate.
        class: String,
        /// Quench model path.
        #[arg(long, default_value = "pixel-forge-quench.safetensors")]
        quench: String,
        /// Cinder model path.
        #[arg(long, default_value = "pixel-forge-cinder.safetensors")]
        cinder: String,
        /// Expert weights path (optional).
        #[arg(long, default_value = "experts.safetensors")]
        experts: String,
        /// Number of images.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Quench foundation steps.
        #[arg(long, default_value_t = 25)]
        quench_steps: usize,
        /// Cinder detail steps.
        #[arg(long, default_value_t = 15)]
        cinder_steps: usize,
        /// Palette.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file.
        #[arg(short, long, default_value = "cascade.png")]
        output: String,
        /// CFG guidance scale. Higher = stronger class adherence.
        #[arg(long, default_value_t = 3.0)]
        cfg: f64,
        /// Seed for deterministic generation. Same seed = same sprite.
        #[arg(long)]
        seed: Option<u64>,
    },
    /// Upscale 16x16 PNGs from data/ to 32x32 in data_v2_32/ (nearest-neighbor).
    Upscale {
        /// Source directory with 16x16 class subdirs.
        #[arg(short, long, default_value = "data")]
        input: String,
        /// Destination directory for 32x32 output.
        #[arg(short, long, default_value = "data_v2_32")]
        output: String,
    },
    /// Federal compliance docs — baked into the binary.
    Govdocs {
        /// Document to display: sbom, security, ssdf, fips, cmmc, privacy, supply-chain, accessibility, fedramp, itar-ear, federal-use-cases. Omit for index.
        doc: Option<String>,
    },
    /// Distributed training: each GPU trains a different model tier.
    /// lf (3070 8GB) → Anvil, bt (5700 XT 8GB) → Quench, gd (3050 Ti 4GB) → Cinder.
    TrainFleet {
        /// Path to training data directory.
        #[arg(short, long, default_value = "data_v3_32")]
        data: String,
        /// Number of training epochs.
        #[arg(short, long, default_value_t = 100)]
        epochs: usize,
        /// Learning rate.
        #[arg(short, long, default_value_t = 2e-4)]
        lr: f64,
        /// Only run on these nodes (comma-separated: lf,bt,gd). Default: all reachable.
        #[arg(long)]
        nodes: Option<String>,
        /// Dry run: show what would be dispatched, don't start.
        #[arg(long)]
        dry_run: bool,
    },
    /// Train a per-class siloed MicroUNet model.
    /// One model per class group — models are leaf nodes, code is trunk.
    TrainSilo {
        /// Silo name (matches [class.NAME] in class_config.toml).
        #[arg(long)]
        class: String,
        /// Path to training data directory.
        #[arg(short, long, default_value = "data_v3_32")]
        data: String,
        /// Output model file (.safetensors).
        #[arg(short, long)]
        output: Option<String>,
        /// Number of training epochs.
        #[arg(short, long, default_value_t = 200)]
        epochs: usize,
        /// Batch size.
        #[arg(short, long, default_value_t = 32)]
        batch_size: usize,
        /// Learning rate.
        #[arg(short, long, default_value_t = 2e-4)]
        lr: f64,
        /// Channel config: "16,32" or "16,32,32".
        #[arg(long)]
        channels: Option<String>,
        /// Config file for class→silo routing.
        #[arg(long, default_value = "class_config.toml")]
        config: String,
        /// Checkpoint every N epochs.
        #[arg(long, default_value_t = 25)]
        checkpoint_every: usize,
        /// Disable EMA.
        #[arg(long)]
        no_ema: bool,
        /// Resume from checkpoint.
        #[arg(long)]
        resume: Option<String>,
        /// Use any-gpu Vulkan backend instead of candle. Required for AMD RDNA1/2
        /// (e.g. RX 5700 XT on bt) where ROCm is unavailable.
        #[cfg(feature = "vulkan")]
        #[arg(long)]
        vulkan: bool,
    },
    /// Generate from a per-class siloed model.
    Silo {
        /// Class to generate (routed via class_config.toml).
        class: String,
        /// Config file.
        #[arg(long, default_value = "class_config.toml")]
        config: String,
        /// Number of images.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Palette.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file.
        #[arg(short, long, default_value = "silo.png")]
        output: String,
        /// Seed for deterministic generation.
        #[arg(long)]
        seed: Option<u64>,
        /// Override diffusion steps.
        #[arg(long)]
        steps: Option<usize>,
    },
    /// Desktop generation: Anvil single-stage (needs 3+ GB VRAM).
    Anvil {
        /// Class to generate.
        class: String,
        /// Anvil model path.
        #[arg(long, default_value = "pixel-forge-anvil.safetensors")]
        anvil: String,
        /// Number of images.
        #[arg(short, long, default_value_t = 1)]
        count: u32,
        /// Diffusion steps.
        #[arg(short, long, default_value_t = 40)]
        steps: usize,
        /// Palette.
        #[arg(short, long, default_value = "stardew")]
        palette: String,
        /// Output file.
        #[arg(short, long, default_value = "anvil.png")]
        output: String,
        /// CFG guidance scale. Higher = stronger class adherence.
        #[arg(long, default_value_t = 3.0)]
        cfg: f64,
    },
    /// Tiered pipeline: silo shape → Cinder detail → palette quantize.
    /// Produces sharper, more class-faithful sprites by routing through
    /// a per-class MicroUNet before Cinder refinement.
    Tiered {
        /// Class to generate: warrior, dragon, building, tree, etc.
        class: String,
        /// Directory containing per-class silo models and class_config.toml.
        #[arg(long, default_value = "models")]
        silo_dir: String,
        /// Cinder detail model (Stage 3).
        #[arg(long, default_value = "pixel-forge-cinder.safetensors")]
        detail_model: String,
        /// Number of sprites to generate.
        #[arg(short, long, default_value_t = 4)]
        count: u32,
        /// Palette for output quantization.
        #[arg(short, long, default_value = "endesga32")]
        palette: String,
        /// Silo DDIM steps (Stage 1).
        #[arg(long, default_value_t = 20)]
        silo_steps: usize,
        /// Cinder DDIM steps (Stage 3).
        #[arg(long, default_value_t = 40)]
        detail_steps: usize,
        /// Noise blend level: 0.0 = all silo, 1.0 = all noise. Default 0.45.
        #[arg(long, default_value_t = 0.45)]
        refine_t: f32,
        /// CFG guidance scale.
        #[arg(long, default_value_t = 3.0)]
        cfg: f64,
        /// Seed for deterministic generation.
        #[arg(long)]
        seed: Option<u64>,
        /// PaletteNet model path (Phase 2). When set, predicts class-specific colors.
        #[arg(long)]
        palette_model: Option<String>,
        /// Output file path.
        #[arg(short, long, default_value = "tiered.png")]
        output: String,
    },
    /// Train Cinder (TinyUNet) on AMD/Vulkan via any-gpu. Required for bt's
    /// 5700 XT (no ROCm). Matches candle Cinder architecture; supports 6ch input
    /// when --condition is supplied, and resume from a 3ch candle checkpoint.
    #[cfg(feature = "vulkan")]
    TrainCinderVk {
        /// Path to training data directory.
        #[arg(short, long, default_value = "data_v3_32")]
        data: String,
        /// Conditioning dir (for 6ch Cinder-detail). Same layout as --data.
        #[arg(long)]
        condition: Option<String>,
        /// Output model file (.safetensors).
        #[arg(short, long, default_value = "models/cinder-detail-vk.safetensors")]
        output: String,
        /// Resume from existing 3ch or 6ch candle checkpoint.
        #[arg(long)]
        resume: Option<String>,
        /// Number of training epochs.
        #[arg(short, long, default_value_t = 100)]
        epochs: usize,
        /// Batch size (large = more GPU efficiency, more VRAM).
        #[arg(short, long, default_value_t = 16)]
        batch_size: u32,
        /// Learning rate.
        #[arg(long, default_value_t = 1e-4)]
        lr: f64,
    },
    /// Generate conditioning data for 6ch Cinder-detail fine-tuning.
    /// Downscales each training image to 16x16 then back to 32x32 (nearest-neighbor),
    /// producing a coarse palette-quantized hint analogous to silo outputs.
    /// Output dir must have the same class structure as the source data dir.
    PrepSiloCond {
        /// Source training data directory (e.g. data_v3_32).
        #[arg(short, long, default_value = "data_v3_32")]
        data: String,
        /// Output conditioning directory (e.g. data_silo_cond_32).
        #[arg(short, long, default_value = "data_silo_cond_32")]
        output: String,
    },
    /// A/B test whether a 6ch conditioned Cinder is actually using the conditioning
    /// channel. Runs forward pass twice with the same noisy input but different
    /// conditioning hints (real coarsened image vs random noise) and reports
    /// per-pixel mean absolute difference. Large diff → model uses conditioning.
    /// Tiny diff → model has learned to ignore the new channel.
    InspectCond {
        /// Path to a 6ch Cinder model (must be 6ch — checked via conv_in.weight shape).
        #[arg(long)]
        model: String,
        /// Class to use for conditioning (must have a silo for real-cond mode).
        #[arg(long, default_value = "building")]
        class: String,
        /// Number of test samples to average over.
        #[arg(long, default_value_t = 4)]
        count: u32,
        /// Seed for reproducibility.
        #[arg(long, default_value_t = 42)]
        seed: u64,
    },
    /// Train the PaletteNet — tiny MLP that predicts class-appropriate colors.
    /// Supervised from K-means dominant colors per class in training data.
    TrainPaletteNet {
        /// Path to training data directory.
        #[arg(short, long, default_value = "data_v3_32")]
        data: String,
        /// Output model file.
        #[arg(short, long, default_value = "models/palette.safetensors")]
        output: String,
        /// Training epochs.
        #[arg(short, long, default_value_t = 50)]
        epochs: usize,
        /// Batch size.
        #[arg(short, long, default_value_t = 64)]
        batch_size: usize,
        /// Learning rate.
        #[arg(long, default_value_t = 1e-3)]
        lr: f64,
        /// Number of palette colors to predict per class.
        #[arg(long, default_value_t = 8)]
        palette_colors: usize,
    },
}

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

    // --sbom: dump machine-readable SPDX and exit
    if cli.sbom {
        let v = env!("CARGO_PKG_VERSION");
        print!("\
SPDXVersion: SPDX-2.3
DataLicense: CC0-1.0
SPDXID: SPDXRef-DOCUMENT
DocumentName: pixel-forge
DocumentNamespace: https://github.com/cochranblock/pixel-forge
Creator: Tool: pixel-forge-{v}

PackageName: pixel-forge
PackageVersion: {v}
PackageLicenseDeclared: Unlicense
PackageDownloadLocation: https://github.com/cochranblock/pixel-forge

## Direct Dependencies

PackageName: clap
PackageVersion: 4
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: serde
PackageVersion: 1
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: serde_json
PackageVersion: 1
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: anyhow
PackageVersion: 1
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: image
PackageVersion: 0.25
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: candle-core
PackageVersion: 0.8
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: candle-nn
PackageVersion: 0.8
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: candle-transformers
PackageVersion: 0.8
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: hf-hub
PackageVersion: 0.4
PackageLicenseDeclared: Apache-2.0

PackageName: tokenizers
PackageVersion: 0.21
PackageLicenseDeclared: Apache-2.0

PackageName: rand
PackageVersion: 0.8
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: bincode
PackageVersion: 1
PackageLicenseDeclared: MIT

PackageName: zstd
PackageVersion: 0.13
PackageLicenseDeclared: MIT

PackageName: eframe
PackageVersion: 0.33
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: egui
PackageVersion: 0.33
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: base64
PackageVersion: 0.22
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: ed25519-dalek
PackageVersion: 2
PackageLicenseDeclared: BSD-3-Clause

PackageName: sha2
PackageVersion: 0.10
PackageLicenseDeclared: MIT OR Apache-2.0

PackageName: half
PackageVersion: 2
PackageLicenseDeclared: MIT OR Apache-2.0
");
        return Ok(());
    }

    let cmd = match cli.cmd {
        Some(cmd) => cmd,
        None => return app::run(),
    };

    match cmd {
        Cmd::Sprite {
            prompt,
            size,
            frames,
            palette,
            output,
        } => {
            let pal = palette::load_palette(&palette)?;
            println!("generating {frames}x {size}x{size} sprite: \"{prompt}\"");
            println!("palette: {} ({} colors)", palette, pal.len());

            let raw_images = pipeline::generate(&prompt, size, frames)?;

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let resized = grid::snap_to_grid(&img, size);
                    palette::quantize(&resized, &pal)
                })
                .collect();

            if frames == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet = sheet::pack_horizontal(&processed);
                sheet.save(&output)?;
            }
            println!("saved: {output}");
        }
        Cmd::Tileset {
            theme,
            size,
            count,
            palette,
            output,
        } => {
            let pal = palette::load_palette(&palette)?;
            println!("generating {count}x {size}x{size} tiles: \"{theme}\"");

            let mut tiles = Vec::new();
            for i in 0..count {
                let tile_prompt = format!("{theme} tile variation {}", i + 1);
                let raw = pipeline::generate(&tile_prompt, size, 1)?;
                let snapped = grid::snap_to_grid(&raw[0], size);
                let quantized = palette::quantize(&snapped, &pal);
                tiles.push(quantized);
            }

            let sheet = sheet::pack_grid(&tiles, 4);
            sheet.save(&output)?;
            println!("saved: {output}");
        }
        Cmd::Palettes => {
            palette::list_palettes();
        }
        Cmd::Setup => {
            pipeline::download_model()?;
            println!("model cached and ready");
        }
        Cmd::Curate { raw, output, size, augment_below } => {
            curate::curate(&raw, &output, size, augment_below)?;
        }
        Cmd::StageCascade { class, sil_model, detail_model, count, sil_steps, detail_steps, palette: palette_name, output, four_dir, real_struct } => {
            let cond = class_cond::lookup(&class.to_lowercase());

            let pal = palette::load_palette(&palette_name)?;
            let raw_images = if let Some(ref struct_dir) = real_struct {
                // Bypass silhouette generation — use real structure maps
                moe::detail_only_sample(
                    &detail_model, struct_dir, &class, &cond, 32, count, detail_steps,
                )?
            } else {
                moe::stage_cascade_sample(
                    &sil_model, &detail_model, &cond, 32, count, sil_steps, detail_steps,
                )?
            };

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, 32);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if four_dir {
                let mut all_sheets = Vec::new();
                for sprite in &processed {
                    let (sheet, _) = relight::four_dir_sheet(sprite);
                    all_sheets.push(sheet);
                }
                if all_sheets.len() == 1 {
                    all_sheets[0].save(&output)?;
                } else {
                    let sw = all_sheets[0].width();
                    let sh = all_sheets[0].height();
                    let mut combined = image::RgbaImage::new(sw, sh * all_sheets.len() as u32);
                    for (i, s) in all_sheets.iter().enumerate() {
                        image::imageops::overlay(&mut combined, s, 0, (i as u32 * sh) as i64);
                    }
                    combined.save(&output)?;
                }
                println!("saved: {output} ({count} × 4 directions)");
            } else {
                if count == 1 {
                    processed[0].save(&output)?;
                } else {
                    let sheet_img = sheet::pack_grid(&processed, 8);
                    sheet_img.save(&output)?;
                }
                println!("saved: {output}");
            }
        }
        Cmd::IngestGemini { input, output } => {
            let input_path = std::path::Path::new(&input);
            let output_path = std::path::Path::new(&output);

            let mut total = 0usize;
            let mut class_counts: std::collections::HashMap<String, usize> = std::collections::HashMap::new();

            for entry in std::fs::read_dir(input_path)? {
                let entry = entry?;
                let path = entry.path();
                if path.extension().and_then(|e| e.to_str()) != Some("png") { continue; }

                let fname = entry.file_name().to_string_lossy().to_string();
                // Parse class from filename: CLASS_batchN.png or CLASS.png
                let class_name = fname.split('_').next()
                    .unwrap_or("misc")
                    .to_lowercase()
                    .replace("-", "_");

                let img = image::open(&path)?.to_rgba8();
                let w = img.width();
                let h = img.height();

                // Auto-detect grid: try common cell sizes
                let cell_size = if w >= 6 * 128 && h >= 5 * 128 { 128 }
                    else if w >= 6 * 64 && h >= 5 * 64 { 64 }
                    else if w >= 6 * 32 && h >= 5 * 32 { 32 }
                    else { (w / 6).max(32) as u32 };

                let cols = w / cell_size;
                let rows = h / cell_size;

                let class_dir = output_path.join(&class_name);
                std::fs::create_dir_all(&class_dir)?;

                let mut file_count = 0usize;
                for row in 0..rows {
                    for col in 0..cols {
                        let x0 = col * cell_size;
                        let y0 = row * cell_size;
                        let cell = image::imageops::crop_imm(&img, x0, y0, cell_size, cell_size).to_image();

                        // Skip empty/background cells (>90% same color)
                        let pixels: Vec<_> = cell.pixels().collect();
                        let first = pixels[0];
                        let same_count = pixels.iter().filter(|p| {
                            (p[0] as i16 - first[0] as i16).abs() < 20 &&
                            (p[1] as i16 - first[1] as i16).abs() < 20 &&
                            (p[2] as i16 - first[2] as i16).abs() < 20
                        }).count();
                        if same_count as f32 / pixels.len() as f32 > 0.90 { continue; }

                        // Remove black/green background → transparency
                        let mut cleaned = image::RgbaImage::new(cell_size, cell_size);
                        for (x, y, px) in cell.enumerate_pixels() {
                            let is_black = px[0] < 15 && px[1] < 15 && px[2] < 15;
                            let is_green = px[1] > 200 && px[0] < 50 && px[2] < 50; // chroma key
                            if is_black || is_green {
                                cleaned.put_pixel(x, y, image::Rgba([0, 0, 0, 0]));
                            } else {
                                cleaned.put_pixel(x, y, *px);
                            }
                        }

                        // Downscale to 32x32 if needed
                        let final_sprite = if cell_size != 32 {
                            image::imageops::resize(&cleaned, 32, 32, image::imageops::FilterType::Nearest)
                        } else {
                            cleaned
                        };

                        // Skip if too few opaque pixels after cleanup
                        let opaque = final_sprite.pixels().filter(|p| p[3] > 10).count();
                        if opaque < 50 { continue; }

                        let existing = *class_counts.get(&class_name).unwrap_or(&0);
                        let out_path = class_dir.join(format!("gemini_{class_name}_{:05}.png", existing));
                        final_sprite.save(&out_path)?;
                        *class_counts.entry(class_name.clone()).or_insert(0) += 1;
                        file_count += 1;
                    }
                }
                if file_count > 0 {
                    println!("{fname}: {file_count} sprites → {class_name}");
                }
                total += file_count;
            }

            println!("\ntotal: {total} sprites ingested");
            for (cls, cnt) in class_counts.iter() {
                println!("  {cls}: {cnt}");
            }
        }
        Cmd::Relight { image, output, split } => {
            let img = image::open(&image)?.to_rgba8();
            let (sheet, sprites) = relight::four_dir_sheet(&img);
            sheet.save(&output)?;
            println!("saved: {output} ({}x{}, 4 views)", sheet.width(), sheet.height());

            if split {
                let stem = std::path::Path::new(&output).file_stem()
                    .and_then(|s| s.to_str()).unwrap_or("relit");
                let dir = std::path::Path::new(&output).parent().unwrap_or(std::path::Path::new("."));
                for (name, sprite) in &sprites {
                    let path = dir.join(format!("{stem}_{name}.png"));
                    sprite.save(&path)?;
                    println!("  {}", path.display());
                }
            }
        }
        Cmd::Skeletons { data, img_size } => {
            train::compute_skeletons_v2(&data, img_size)?;
            // Copy to CWD so generation commands find it
            let src = format!("{data}/skeletons_v2.safetensors");
            let dst = "skeletons_v2.safetensors";
            if std::path::Path::new(&src).exists() && src != dst {
                std::fs::copy(&src, dst)?;
                println!("copied: {dst}");
            }
        }
        Cmd::PrepStages { data, palette_colors: _ } => {
            let data_path = std::path::Path::new(&data);
            let class_names = discover_classes(data_path);
            let struct_base = data_path.join("_structure");  // RGBA: SDF, nx, ny, outline
            let lum_base = data_path.join("_luminance");     // Grayscale: light/shadow
            let sil_base = data_path.join("_silhouettes");   // Binary mask
            std::fs::create_dir_all(&struct_base)?;
            std::fs::create_dir_all(&lum_base)?;
            std::fs::create_dir_all(&sil_base)?;

            let mut total = 0usize;
            for class_name in &class_names {
                let class_dir = data_path.join(class_name);
                if !class_dir.is_dir() { continue; }

                let struct_dir = struct_base.join(class_name);
                let lum_dir = lum_base.join(class_name);
                let sil_dir = sil_base.join(class_name);
                std::fs::create_dir_all(&struct_dir)?;
                std::fs::create_dir_all(&lum_dir)?;
                std::fs::create_dir_all(&sil_dir)?;

                let mut count = 0usize;
                for entry in std::fs::read_dir(&class_dir)? {
                    let entry = entry?;
                    let path = entry.path();
                    if path.extension().and_then(|e| e.to_str()) != Some("png") { continue; }

                    let img = image::open(&path)?.to_rgba8();
                    let w = img.width() as usize;
                    let h = img.height() as usize;
                    let fname = entry.file_name();

                    // Binary mask: opaque pixels
                    let mut mask = vec![false; w * h];
                    for (x, y, px) in img.enumerate_pixels() {
                        if px[3] > 10 {
                            mask[y as usize * w + x as usize] = true;
                        }
                    }

                    // SDF via distance transform (Chamfer approximation)
                    // Forward + backward pass with 3-4-3 weights
                    let mut dist = vec![f32::MAX; w * h];
                    // Set border pixels to 0, interior to MAX
                    for y in 0..h {
                        for x in 0..w {
                            let i = y * w + x;
                            if !mask[i] {
                                dist[i] = 0.0;
                            } else {
                                // Check if any neighbor is outside — edge pixel
                                let is_edge = (x == 0 || !mask[i - 1])
                                    || (x == w - 1 || !mask[i + 1])
                                    || (y == 0 || !mask[i - w])
                                    || (y == h - 1 || !mask[i + w]);
                                dist[i] = if is_edge { 0.5 } else { f32::MAX };
                            }
                        }
                    }
                    // Forward pass (top-left to bottom-right)
                    for y in 1..h {
                        for x in 1..w {
                            let i = y * w + x;
                            if !mask[i] { continue; }
                            let up = dist[(y - 1) * w + x] + 1.0;
                            let left = dist[y * w + (x - 1)] + 1.0;
                            let diag = dist[(y - 1) * w + (x - 1)] + 1.414;
                            dist[i] = dist[i].min(up).min(left).min(diag);
                        }
                    }
                    // Backward pass (bottom-right to top-left)
                    for y in (0..h - 1).rev() {
                        for x in (0..w - 1).rev() {
                            let i = y * w + x;
                            if !mask[i] { continue; }
                            let down = dist[(y + 1) * w + x] + 1.0;
                            let right = dist[y * w + (x + 1)] + 1.0;
                            let diag = dist[(y + 1) * w + (x + 1)] + 1.414;
                            dist[i] = dist[i].min(down).min(right).min(diag);
                        }
                    }

                    // Normalize SDF to [0, 1]
                    let max_dist = dist.iter().cloned().fold(0.0f32, f32::max).max(1.0);
                    for d in &mut dist {
                        *d /= max_dist;
                    }

                    // Normals from SDF gradient (Sobel-like)
                    let mut nx = vec![0.0f32; w * h];
                    let mut ny = vec![0.0f32; w * h];
                    for y in 1..h - 1 {
                        for x in 1..w - 1 {
                            let i = y * w + x;
                            if !mask[i] { continue; }
                            // Horizontal gradient
                            let gx = dist[i + 1] - dist[i - 1];
                            // Vertical gradient
                            let gy = dist[i + w] - dist[i - w];
                            let len = (gx * gx + gy * gy).sqrt().max(1e-6);
                            nx[i] = gx / len;
                            ny[i] = gy / len;
                        }
                    }

                    // Outline: pixels where SDF ≈ 0.5 (on the edge)
                    // More precisely: mask=true but at least one neighbor is mask=false
                    let mut outline = vec![0u8; w * h];
                    for y in 0..h {
                        for x in 0..w {
                            let i = y * w + x;
                            if !mask[i] { continue; }
                            let has_bg = (x > 0 && !mask[i - 1])
                                || (x < w - 1 && !mask[i + 1])
                                || (y > 0 && !mask[i - w])
                                || (y < h - 1 && !mask[i + w]);
                            if has_bg { outline[i] = 255; }
                        }
                    }

                    // Luminance from original sprite (0.299R + 0.587G + 0.114B)
                    let mut lum = vec![0u8; w * h];
                    for (x, y, px) in img.enumerate_pixels() {
                        let i = y as usize * w + x as usize;
                        if mask[i] {
                            let l = 0.299 * px[0] as f32 + 0.587 * px[1] as f32 + 0.114 * px[2] as f32;
                            lum[i] = l.clamp(0.0, 255.0) as u8;
                        }
                    }

                    // Structure: RGBA = SDF, normal_x, normal_y, outline
                    let mut struct_img = image::RgbaImage::new(w as u32, h as u32);
                    // Luminance: grayscale RGB
                    let mut lum_img = image::RgbImage::new(w as u32, h as u32);
                    // Silhouette: binary
                    let mut sil_img = image::RgbImage::new(w as u32, h as u32);

                    for y in 0..h {
                        for x in 0..w {
                            let i = y * w + x;
                            if mask[i] {
                                let r = (dist[i] * 255.0) as u8;
                                let g = ((0.5 + 0.5 * nx[i]) * 255.0) as u8;
                                let b = ((0.5 + 0.5 * ny[i]) * 255.0) as u8;
                                struct_img.put_pixel(x as u32, y as u32, image::Rgba([r, g, b, outline[i]]));
                                lum_img.put_pixel(x as u32, y as u32, image::Rgb([lum[i], lum[i], lum[i]]));
                                sil_img.put_pixel(x as u32, y as u32, image::Rgb([255, 255, 255]));
                            } else {
                                struct_img.put_pixel(x as u32, y as u32, image::Rgba([0, 128, 128, 0]));
                                lum_img.put_pixel(x as u32, y as u32, image::Rgb([0, 0, 0]));
                                sil_img.put_pixel(x as u32, y as u32, image::Rgb([0, 0, 0]));
                            }
                        }
                    }
                    struct_img.save(struct_dir.join(&fname))?;
                    lum_img.save(lum_dir.join(&fname))?;
                    sil_img.save(sil_dir.join(&fname))?;

                    count += 1;
                }
                if count > 0 {
                    println!("{class_name}: {count} sprites → 5-layer structure");
                }
                total += count;
            }
            println!("total: {total} structure maps");
            println!("  _structure/ RGBA: SDF(R), normal_x(G), normal_y(B), outline(A)");
            println!("  _luminance/ grayscale: light/shadow from original");
            println!("  _silhouettes/ binary: shape mask");
            println!("  → 5 conditioning channels for detail model (8ch input: 5 cond + 3 noisy)");
        }
        Cmd::Train {
            data,
            output,
            epochs,
            batch_size,
            lr,
            img_size,
            medium,
            anvil,
            no_ema,
            v_pred,
            lr_min,
            warmup,
            condition,
            fp16,
            checkpoint_every,
            resume,
            sponge,
            device,
        } => {
            // Vulkan backend: any-gpu now powers the siloed MicroUNet path only.
            // The legacy Cinder/TinyUNet Vulkan path was retired — use `train-silo --vulkan`
            // for per-class training on AMD (RX 5700 XT etc.).
            #[cfg(feature = "vulkan")]
            if device == "vulkan" {
                anyhow::bail!(
                    "vulkan training is now routed through `pixel-forge train-silo --vulkan --class <name>` \
                     (MicroUNet path). Use that for AMD GPUs where candle's CUDA/Metal backends can't reach."
                );
            }
            #[cfg(not(feature = "vulkan"))]
            if device == "vulkan" {
                anyhow::bail!("vulkan backend requires --features vulkan. Rebuild with: cargo build --release --features vulkan --no-default-features");
            }

            // Fix: default output to correct model filename based on tier
            let output = if output == "pixel-forge-cinder.safetensors" {
                if anvil {
                    "pixel-forge-anvil.safetensors".to_string()
                } else if medium {
                    "pixel-forge-quench.safetensors".to_string()
                } else {
                    output
                }
            } else {
                output
            };
            let config = train::TrainConfig {
                data_dir: data,
                output,
                epochs,
                batch_size,
                lr,
                img_size,
                medium,
                anvil,
                ema: !no_ema,
                v_prediction: v_pred,
                lr_min,
                warmup_epochs: warmup,
                condition_dir: condition,
                mixed_precision: fp16,
                checkpoint_every,
                resume,
                ..Default::default()
            };
            if sponge {
                sponge::sponge_train(&config)?;
            } else {
                let stop = train::train(&config)?;
                match stop {
                    train::TrainStop::Complete { final_loss } => {
                        println!("training complete, final loss: {final_loss:.6}");
                    }
                    train::TrainStop::NanLoss(epoch) => {
                        eprintln!("training aborted: NaN loss at epoch {}", epoch + 1);
                        eprintln!("try: --sponge for auto-retry, or lower --lr");
                        std::process::exit(1);
                    }
                    train::TrainStop::Plateau { epoch, loss } => {
                        eprintln!("training aborted: loss plateau at epoch {} (loss={loss:.8})", epoch + 1);
                        eprintln!("try: --sponge for auto-retry, or different data");
                        std::process::exit(1);
                    }
                }
            }
        }
        Cmd::Generate {
            class,
            model,
            size,
            count,
            steps,
            palette: palette_name,
            output,
            medium,
            seed,
            four_dir,
            cfg,
        } => {
            if class.is_empty() {
                anyhow::bail!("class name cannot be empty. Try: character, weapon, dragon, potion, tree");
            }
            if count == 0 {
                anyhow::bail!("count must be at least 1");
            }
            if !std::path::Path::new(&model).exists() {
                anyhow::bail!("model not found: {model}\n\nTo train a model:\n  pixel-forge train --data data_v2_32 --epochs 100 --img-size 32\n\nOr download a pre-trained model and place it at:\n  {model}");
            }

            let cond = class_cond::lookup(&class.to_lowercase());

            println!("generating {count}x {size}x{size} class={class} (super={})", cond.super_id);

            let pal = palette::load_palette(&palette_name)?;
            let raw_images = if medium {
                train::sample_medium_seeded_cfg(&model, &cond, size, count, steps, seed, cfg)?
            } else {
                train::sample_seeded_cfg(&model, &cond, size, count, steps, seed, cfg)?
            };

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, size);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if four_dir {
                // Relight each sprite into 4-direction sheets
                let mut all_sheets = Vec::new();
                for sprite in &processed {
                    let (sheet, _) = relight::four_dir_sheet(sprite);
                    all_sheets.push(sheet);
                }
                if all_sheets.len() == 1 {
                    all_sheets[0].save(&output)?;
                } else {
                    // Stack vertically: each row is one character's 4 views
                    let sw = all_sheets[0].width();
                    let sh = all_sheets[0].height();
                    let mut combined = image::RgbaImage::new(sw, sh * all_sheets.len() as u32);
                    for (i, s) in all_sheets.iter().enumerate() {
                        image::imageops::overlay(&mut combined, s, 0, (i as u32 * sh) as i64);
                    }
                    combined.save(&output)?;
                }
                println!("saved: {output} ({count} sprites × 4 directions)");
            } else {
                if count == 1 {
                    processed[0].save(&output)?;
                } else {
                    let sheet_img = sheet::pack_grid(&processed, 8);
                    sheet_img.save(&output)?;
                }
                println!("saved: {output}");
            }
        }
        Cmd::Swipe { image, verdict, class, store } => {
            let good = match verdict.to_lowercase().as_str() {
                "good" | "right" | "yes" | "y" | "1" => true,
                "bad" | "left" | "no" | "n" | "0" => false,
                _ => anyhow::bail!("verdict must be 'good' or 'bad', got '{verdict}'"),
            };

            let class_names = discover_classes(std::path::Path::new("data_v2_32"));
            let class_id = class_names.iter()
                .position(|n| n == &class.to_lowercase())
                .unwrap_or(14) as u32;

            // Load sprite pixels (channel-first f32)
            let img = image::open(&image)?.to_rgb8();
            let img = image::imageops::resize(&img, 16, 16, image::imageops::FilterType::Nearest);
            let n = 256usize; // 16*16
            let mut pixels = vec![0.0f32; 768];
            for y in 0..16u32 {
                for x in 0..16u32 {
                    let p = img.get_pixel(x, y);
                    let idx = (y * 16 + x) as usize;
                    pixels[idx] = p[0] as f32 / 255.0;
                    pixels[n + idx] = p[1] as f32 / 255.0;
                    pixels[2 * n + idx] = p[2] as f32 / 255.0;
                }
            }

            let mut swipe_store = swipe_store::SwipeStore::load_or_default(&store, 200);
            let len = swipe_store.record(pixels, good, class_id);
            swipe_store.save(&store)?;

            let verdict_str = if good { "good" } else { "bad" };
            println!("recorded swipe: {} ({}) → {} swipes in buffer", verdict_str, class, len);
            if swipe_store.judge_ready() {
                println!("judge ready — run `pixel-forge train-judge` to train");
            }
            if swipe_store.should_retrain() {
                println!("retrain triggered (every 5 swipes)");
            }
        }
        Cmd::TrainJudge { store, output } => {
            let device = pipeline::best_device();
            let swipe_store = swipe_store::SwipeStore::load_or_default(&store, 200);

            let (varmap, _model, result) = judge::train_judge(&swipe_store, &device)?;
            judge::save_judge(&varmap, &output)?;

            println!("{result}");
            let file_size = std::fs::metadata(&output)?.len();
            println!("saved: {} ({:.1} KB)", output, file_size as f64 / 1024.0);
        }
        Cmd::Judge { input, model } => {
            let device = pipeline::best_device();
            let (_varmap, judge_model) = judge::load_judge(&model, &device)?;

            let path = std::path::Path::new(&input);
            let files: Vec<std::path::PathBuf> = if path.is_dir() {
                pixel_forge::walk_pngs(&input)
            } else {
                vec![path.to_path_buf()]
            };

            println!("judging {} sprites...", files.len());
            for file in &files {
                let img = image::open(file)?.to_rgb8();
                let img = image::imageops::resize(&img, 16, 16, image::imageops::FilterType::Nearest);
                let n = 256usize;
                let mut pixels = vec![0.0f32; 768];
                for y in 0..16u32 {
                    for x in 0..16u32 {
                        let p = img.get_pixel(x, y);
                        let idx = (y * 16 + x) as usize;
                        pixels[idx] = p[0] as f32 / 255.0;
                        pixels[n + idx] = p[1] as f32 / 255.0;
                        pixels[2 * n + idx] = p[2] as f32 / 255.0;
                    }
                }
                let score = judge_model.score_one(&pixels, &device)?;
                let verdict = if score > 0.5 { "GOOD" } else { "BAD " };
                println!("  {} {:.3} {}", verdict, score, file.display());
            }
        }
        Cmd::Scene { mode, combiner_model, biome, temperature, count, output } => {
            match mode.as_str() {
                "bootstrap" => {
                    let biome_enum = match biome.to_lowercase().as_str() {
                        "forest" => scene::Biome::Forest,
                        "dungeon" => scene::Biome::Dungeon,
                        "village" => scene::Biome::Village,
                        "cave" => scene::Biome::Cave,
                        "plains" => scene::Biome::Plains,
                        _ => {
                            println!("unknown biome '{biome}', using forest");
                            scene::Biome::Forest
                        }
                    };

                    for i in 0..count {
                        let grid = scene::generate_seeded(biome_enum);
                        let img = grid.render_placeholder();
                        let out = if count == 1 {
                            output.clone()
                        } else {
                            format!("{}-{}.png", output.trim_end_matches(".png"), i + 1)
                        };
                        img.save(&out)?;
                        println!("scene {}/{}: {} sprites, saved {}", i + 1, count, grid.sprite_count(), out);
                    }
                }
                "model" => {
                    let device = pipeline::best_device();
                    let (_varmap, model) = combiner::load_combiner(&combiner_model, &device)?;

                    for i in 0..count {
                        let grid = combiner::generate_scene(&model, None, temperature, &device)?;
                        let img = grid.render_placeholder();
                        let out = if count == 1 {
                            output.clone()
                        } else {
                            format!("{}-{}.png", output.trim_end_matches(".png"), i + 1)
                        };
                        img.save(&out)?;
                        println!("scene {}/{}: {} sprites, saved {}", i + 1, count, grid.sprite_count(), out);
                    }
                }
                _ => anyhow::bail!("unknown scene mode '{mode}' — use 'bootstrap' or 'model'"),
            }
        }
        Cmd::TrainCombiner { num_scenes, epochs, output } => {
            let device = pipeline::best_device();

            println!("generating {} bootstrap scenes...", num_scenes);
            let scenes = scene::generate_bootstrap(num_scenes);

            let config = combiner::CombinerTrainConfig {
                epochs,
                batch_size: 16,
                lr: 1e-3,
                mask_frac_min: 0.2,
                mask_frac_max: 0.5,
            };

            let (varmap, _model, result) = combiner::train_combiner(&scenes, &config, &device)?;
            combiner::save_combiner(&varmap, &output)?;

            println!("{result}");
            let file_size = std::fs::metadata(&output)?.len();
            println!("saved: {} ({:.1} KB)", output, file_size as f64 / 1024.0);
        }
        Cmd::TrainLora { model, judge: judge_path, output, class } => {
            let device = pipeline::best_device();

            let cond = class_cond::lookup(&class.to_lowercase());

            // Load base Generator
            let mut base_varmap = candle_nn::VarMap::new();
            let base_vb = candle_nn::VarBuilder::from_varmap(&base_varmap, candle_core::DType::F32, &device);
            let _base_model = tiny_unet::TinyUNet::new(base_vb)?;
            base_varmap.load(&model)?;
            println!("loaded base model: {}", model);

            // Load Judge
            let (_jvm, judge_model) = judge::load_judge(&judge_path, &device)?;
            println!("loaded judge: {}", judge_path);

            // Create or load LoRA
            let mut lora_varmap = candle_nn::VarMap::new();
            let lora_vb = candle_nn::VarBuilder::from_varmap(&lora_varmap, candle_core::DType::F32, &device);
            let lora_set = lora::LoraSet::new(lora_vb)?;

            if std::path::Path::new(&output).exists() {
                lora_varmap.load(&output)?;
                println!("loaded existing lora: {}", output);
            }

            let loss = lora::train_lora_step(
                &base_varmap, &lora_varmap, &lora_set, &judge_model, &device, 8, &cond,
            )?;
            println!("lora training: avg_loss={:.4}", loss);

            lora::save_lora(&lora_varmap, &output)?;
            let file_size = std::fs::metadata(&output)?.len();
            println!("saved: {} ({:.1} KB)", output, file_size as f64 / 1024.0);
        }
        Cmd::SwipeStats { store } => {
            let swipe_store = swipe_store::SwipeStore::load_or_default(&store, 200);
            println!("{}", swipe_store.stats());
            if swipe_store.judge_ready() {
                println!("judge: ready to train");
            } else {
                let need = 10 - swipe_store.len().min(10);
                println!("judge: need {} more swipes", need);
            }
        }
        Cmd::Pipeline { model, judge: judge_path, combiner: combiner_path, lora: _lora_path, palette: palette_name, output } => {
            let device = pipeline::best_device();
            let pal = palette::load_palette(&palette_name)?;

            // 1. GENERATE — produce candidates per class
            println!("step 1: generating sprite candidates...");
            let mut all_sprites: Vec<(u32, Vec<f32>, image::RgbaImage)> = Vec::new();
            let classes_to_gen = ["character", "terrain", "enemy", "tree"];
            let count_per_class = 5u32;

            for &class_name in &classes_to_gen {
                let cond = class_cond::lookup(class_name);
                let cid = cond.super_id;
                let raw_images = train::sample(&model, &cond, 16, count_per_class, 40)?;
                for img in raw_images {
                    let snapped = grid::snap_to_grid(&img, 16);
                    let quantized = palette::quantize(&snapped, &pal);

                    // Convert to channel-first f32 for Judge
                    let n = 256usize;
                    let mut pixels = vec![0.0f32; 768];
                    for y in 0..16u32 {
                        for x in 0..16u32 {
                            let p = quantized.get_pixel(x, y);
                            let idx = (y * 16 + x) as usize;
                            pixels[idx] = p[0] as f32 / 255.0;
                            pixels[n + idx] = p[1] as f32 / 255.0;
                            pixels[2 * n + idx] = p[2] as f32 / 255.0;
                        }
                    }
                    all_sprites.push((cid, pixels, quantized));
                }
            }
            println!("  generated {} sprites across {} classes",
                all_sprites.len(), classes_to_gen.len());

            // 2. JUDGE — filter if Judge model exists
            let accepted: Vec<(u32, Vec<f32>, image::RgbaImage)>;
            if std::path::Path::new(&judge_path).exists() {
                println!("step 2: judging sprites...");
                let (_jvm, judge_model) = judge::load_judge(&judge_path, &device)?;
                accepted = all_sprites.into_iter().filter(|(_, pixels, _)| {
                    judge_model.score_one(pixels, &device).unwrap_or(0.0) > 0.5
                }).collect();
                println!("  accepted {}/{} sprites", accepted.len(),
                    classes_to_gen.len() * count_per_class as usize);
            } else {
                println!("step 2: no judge model, accepting all sprites");
                accepted = all_sprites;
            }

            if accepted.is_empty() {
                anyhow::bail!("no sprites passed the Judge — need more swipe data or lower threshold");
            }

            // 3. COMBINE — arrange into scene
            println!("step 3: arranging scene...");
            let grid = if std::path::Path::new(&combiner_path).exists() {
                let (_cvm, combiner_model) = combiner::load_combiner(&combiner_path, &device)?;
                combiner::generate_scene(&combiner_model, None, 0.8, &device)?
            } else {
                println!("  no combiner model, using bootstrap layout");
                scene::generate_seeded(scene::Biome::Forest)
            };

            // 4. RENDER — composite sprites into grid
            println!("step 4: rendering scene...");
            let mut rng = rand::thread_rng();
            use rand::seq::SliceRandom;
            let mut sprites_by_class: std::collections::HashMap<u32, Vec<&Vec<f32>>> =
                std::collections::HashMap::new();
            for (cid, pixels, _) in &accepted {
                sprites_by_class.entry(*cid).or_default().push(pixels);
            }

            let scene_img = grid.render(|class_id| {
                sprites_by_class.get(&class_id)
                    .and_then(|sprites| sprites.choose(&mut rng))
                    .map(|px| px.to_vec())
            });

            scene_img.save(&output)?;
            println!("saved scene: {} ({}×{} px, {} sprites placed)",
                output, scene::SCENE_PX, scene::SCENE_PX, grid.sprite_count());
        }
        Cmd::Auto { class, size, count, steps, palette: palette_name, output } => {
            let pal = palette::load_palette(&palette_name)?;
            println!("auto-detect: picking best model for this device...");

            let cond = class_cond::lookup(&class.to_lowercase());
            let raw_images = device_cap::auto_sample(
                &cond,
                size,
                count,
                steps,
            )?;

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, size);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if count == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&processed, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output}");
        }
        Cmd::Probe { json } => {
            let profile = device_cap::reprobe()?;
            if json {
                // Machine-readable output for cluster probing
                println!("{}", serde_json::to_string(&profile)?);
            } else {
                println!("backend: {}", profile.backend);
                println!("ram: {} MB", profile.ram_mb);
                println!("cinder: {:.1} ms/step", profile.cinder_ms);
                if let Some(q) = profile.quench_ms {
                    println!("quench: {:.1} ms/step", q);
                }
                println!("selected tier: {}", profile.tier);
                println!("model file: {}", profile.tier.model_file());
            }
        }
        Cmd::ClusterProbe => {
            let state = cluster::probe_cluster()?;
            println!("\ncluster summary:");
            let nodes = state.active_nodes();
            for n in &nodes {
                let tier = n.profile.as_ref().map(|p| p.tier.to_string()).unwrap_or("?".into());
                let backend = n.profile.as_ref().map(|p| p.backend.to_string()).unwrap_or("?".into());
                println!("  {:6} | {:5} | {:6} | {:.1} spr/s",
                    n.name, backend, tier, n.throughput);
            }
            println!("  total: {:.1} sprites/sec", state.total_throughput);
        }
        Cmd::ClusterDeploy => {
            cluster::deploy_all()?;
        }
        Cmd::ClusterSync => {
            cluster::sync_models_all()?;
        }
        Cmd::ClusterGenerate { class, count, steps, palette: palette_name, output } => {
            let state = cluster::probe_cluster()?;
            let pngs = cluster::cluster_generate(&class, count, steps, &palette_name, &state)?;

            if pngs.is_empty() {
                anyhow::bail!("no sprites generated");
            }

            // Decode PNGs into images for grid packing
            let images: Vec<image::RgbaImage> = pngs.iter()
                .filter_map(|bytes| {
                    image::load_from_memory(bytes).ok().map(|img| img.to_rgba8())
                })
                .collect();

            if images.len() == 1 {
                images[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&images, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {} ({} sprites)", output, images.len());
        }
        Cmd::Gpu { args } => {
            // Delegate to kova c2 gpu
            let mut cmd = std::process::Command::new("kova");
            cmd.args(["c2", "gpu"]);
            cmd.args(&args);
            let status = cmd.status()?;
            if !status.success() {
                anyhow::bail!("kova c2 gpu failed");
            }
        }
        Cmd::Forge { class, count, steps, threshold, max_attempts, disc, palette: palette_name, output, lat, lon, cascade: use_cascade } => {
            let cond = class_cond::lookup(&class.to_lowercase());

            let pal = palette::load_palette(&palette_name)?;
            let signing_key = poa::load_or_create_keypair()?;

            println!("forge: class={class} threshold={threshold} cascade={use_cascade}");

            // Generate with quality gate
            let sprites = if use_cascade {
                let config = moe::CascadeConfig::default();
                moe::cascade_with_gate(
                    "pixel-forge-quench.safetensors",
                    "pixel-forge-cinder.safetensors",
                    Some("experts.safetensors"),
                    &disc, &cond, 32, count, &config, threshold, max_attempts,
                )?
            } else {
                let device = pipeline::best_device();
                let tier = device_cap::best_available();
                discriminator::generate_with_gate(
                    tier, &cond, count, steps, threshold, max_attempts, &disc, &device,
                )?
            };

            // Sign each sprite, save images, produce Ghost Fabric packets
            let device = pipeline::best_device();
            let mut packet_count = 0usize;
            let sprite_count = sprites.len();
            let mut processed = Vec::new();

            for sprite in sprites {
                let (_dvm, disc_model) = discriminator::load(&disc, &device)?;
                let (_, score) = discriminator::quality_gate(&disc_model, &sprite, 0.0, &device)?;
                let packet = poa::sign_artifact(&sprite, cond.super_id, score, lat, lon, &signing_key);
                assert!(packet.verify(), "PoA self-verification failed");
                println!("  {}", packet.summary());
                let wire = packet.to_bytes();
                println!("  wire: {} bytes (LoRa ready: {})", wire.len(), wire.len() <= 255);
                packet_count += 1;

                let snapped = grid::snap_to_grid(&sprite, 16);
                processed.push(palette::quantize(&snapped, &pal));
            }

            if count == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&processed, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output} ({} sprites, {} PoA packets)", sprite_count, packet_count);
        }
        Cmd::Sign { image, class, score, disc, lat, lon } => {
            let cond = class_cond::lookup(&class.to_lowercase());

            let img = image::open(&image)?.to_rgba8();
            let signing_key = poa::load_or_create_keypair()?;

            let quality = if let Some(s) = score {
                s
            } else {
                let device = pipeline::best_device();
                let (_dvm, disc_model) = discriminator::load(&disc, &device)?;
                let (_, s) = discriminator::quality_gate(&disc_model, &img, 0.0, &device)?;
                println!("discriminator score: {:.3}", s);
                s
            };

            let packet = poa::sign_artifact(&img, cond.super_id, quality, lat, lon, &signing_key);
            let wire = packet.to_bytes();
            println!("{}", packet.summary());
            println!("verified: {}", packet.verify());
            println!("wire: {} bytes (LoRa ready: {})", wire.len(), wire.len() <= 255);

            // Write packet to .poa file alongside the image
            let poa_path = format!("{}.poa", image.trim_end_matches(".png"));
            std::fs::write(&poa_path, wire)?;
            println!("saved: {poa_path}");
        }
        Cmd::NodeKey => {
            let sk = poa::load_or_create_keypair()?;
            let vk = sk.verifying_key();
            let pk_hex: String = vk.to_bytes().iter().map(|b| format!("{b:02x}")).collect();
            println!("node public key: {pk_hex}");
            println!("key file: ~/.pixel-forge/node.key");
        }
        Cmd::Plugin { r#loop } => {
            if r#loop {
                plugin::run_loop()?;
            } else {
                plugin::run()?;
            }
        }
        Cmd::TrainExperts { quench, data, output, epochs, batch_size } => {
            expert_train::train_experts(&quench, &data, &output, epochs, batch_size)?;
        }
        Cmd::Quantize { input, output } => {
            let out = if let Some(o) = output {
                let size = quantize::quantize_f32_to_f16(&input, &o)?;
                println!("wrote {} ({:.1} MB)", o, size as f64 / 1_048_576.0);
                o
            } else {
                quantize::quantize_in_place(&input)?
            };
            println!("done: {out}");
        }
        Cmd::Cascade { class, quench, cinder, experts, count, quench_steps, cinder_steps, palette: palette_name, output, cfg, seed } => {
            let cond = class_cond::lookup(&class.to_lowercase());

            let pal = palette::load_palette(&palette_name)?;
            let experts_path = if std::path::Path::new(&experts).exists() {
                Some(experts.as_str())
            } else {
                None
            };

            let config = moe::CascadeConfig {
                quench_steps,
                cinder_steps,
            };

            let raw_images = moe::cascade_sample(&quench, &cinder, experts_path, &cond, 32, count, &config, cfg, seed)?;

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, 32);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if count == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&processed, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output}");
        }
        Cmd::Upscale { input, output } => {
            let classes = [
                "character", "weapon", "potion", "terrain", "enemy",
                "tree", "building", "animal", "effect", "food",
                "armor", "tool", "vehicle", "ui", "misc",
            ];
            let mut total = 0u32;
            for class in &classes {
                let src_dir = std::path::Path::new(&input).join(class);
                let dst_dir = std::path::Path::new(&output).join(class);
                if !src_dir.is_dir() {
                    continue;
                }
                std::fs::create_dir_all(&dst_dir)?;
                let mut count = 0u32;
                for entry in std::fs::read_dir(&src_dir)? {
                    let entry = entry?;
                    let path = entry.path();
                    if path.extension().and_then(|e| e.to_str()) != Some("png") {
                        continue;
                    }
                    let fname = path.file_name().unwrap();
                    let dst_path = dst_dir.join(fname);
                    if dst_path.exists() {
                        continue;
                    }
                    let img = image::open(&path)?.to_rgba8();
                    let resized = image::imageops::resize(
                        &img, 32, 32,
                        image::imageops::FilterType::Nearest,
                    );
                    resized.save(&dst_path)?;
                    count += 1;
                }
                if count > 0 {
                    println!("{class}: {count} upscaled");
                }
                total += count;
            }
            println!("total: {total} tiles upscaled to 32x32");
        }
        Cmd::Govdocs { doc } => {
            match doc.as_deref() {
                Some("sbom") => print!("{GOVDOC_SBOM}"),
                Some("security") => print!("{GOVDOC_SECURITY}"),
                Some("ssdf") => print!("{GOVDOC_SSDF}"),
                Some("fips") => print!("{GOVDOC_FIPS}"),
                Some("cmmc") => print!("{GOVDOC_CMMC}"),
                Some("privacy") => print!("{GOVDOC_PRIVACY}"),
                Some("supply-chain") => print!("{GOVDOC_SUPPLY_CHAIN}"),
                Some("accessibility") => print!("{GOVDOC_ACCESSIBILITY}"),
                Some("fedramp") => print!("{GOVDOC_FEDRAMP}"),
                Some("itar-ear") => print!("{GOVDOC_ITAR_EAR}"),
                Some("federal-use-cases") => print!("{GOVDOC_FEDERAL_USE}"),
                Some(other) => println!("unknown doc: {other}. Run `pixel-forge govdocs` for index."),
                None => {
                    println!("Pixel Forge v{} — Federal Compliance Docs", env!("CARGO_PKG_VERSION"));
                    println!("=========================================");
                    println!();
                    println!("Available documents:");
                    println!("  pixel-forge govdocs sbom              Software Bill of Materials (EO 14028)");
                    println!("  pixel-forge govdocs security           Security posture");
                    println!("  pixel-forge govdocs ssdf               NIST SP 800-218 mapping");
                    println!("  pixel-forge govdocs fips               FIPS 140-2/3 status");
                    println!("  pixel-forge govdocs cmmc               CMMC Level 1-2 practices");
                    println!("  pixel-forge govdocs privacy            Privacy impact assessment");
                    println!("  pixel-forge govdocs supply-chain       Supply chain integrity");
                    println!("  pixel-forge govdocs accessibility      Section 508 / WCAG");
                    println!("  pixel-forge govdocs fedramp            FedRAMP applicability");
                    println!("  pixel-forge govdocs itar-ear           Export control (ITAR/EAR)");
                    println!("  pixel-forge govdocs federal-use-cases  Federal agency use cases");
                }
            }
        }
        Cmd::TrainSilo {
            class, data, output, epochs, batch_size, lr, channels, config,
            checkpoint_every, no_ema, resume,
            #[cfg(feature = "vulkan")]
            vulkan,
        } => {
            // Resolve class filter from config if available
            let class_filter = if class_router::ClassRouter::exists(&config) {
                let router = class_router::ClassRouter::load(&config)?;
                let classes = router.classes_for_silo(&class);
                if classes.is_empty() {
                    vec![class.clone()]
                } else {
                    classes.into_iter().map(|s| s.to_string()).collect()
                }
            } else {
                vec![class.clone()]
            };

            // Resolve channels
            let micro_channels = if let Some(ref ch_str) = channels {
                ch_str.split(',').map(|s| s.trim().parse::<usize>()).collect::<std::result::Result<Vec<_>, _>>()?
            } else if class_router::ClassRouter::exists(&config) {
                let router = class_router::ClassRouter::load(&config)?;
                router.route(&class).map(|dc| dc.channels.clone()).unwrap_or_else(|| vec![16, 32])
            } else {
                vec![16, 32]
            };

            let output = output.unwrap_or_else(|| format!("models/{class}.safetensors"));
            // Ensure models/ directory exists
            if let Some(parent) = std::path::Path::new(&output).parent() {
                std::fs::create_dir_all(parent)?;
            }

            #[cfg(feature = "vulkan")]
            if vulkan {
                // Route to any-gpu Vulkan backend — bypasses candle entirely.
                println!("train-silo: using any-gpu Vulkan backend (class={class})");
                pixel_forge::vulkan_backend::vulkan_micro_train(
                    &data, &output,
                    &micro_channels,
                    &class_filter,
                    epochs,
                    batch_size as u32,
                    lr,
                    32,
                )?;
                println!("silo vulkan training complete: {output}");
                return Ok(());
            }

            let train_config = train::TrainConfig {
                data_dir: data,
                output,
                epochs,
                batch_size,
                lr,
                img_size: 32,
                micro: true,
                class_filter,
                micro_channels,
                ema: !no_ema,
                checkpoint_every,
                resume,
                cfg_dropout: 0.0, // no CFG for unconditional micro models
                ..Default::default()
            };
            let stop = train::train(&train_config)?;
            match stop {
                train::TrainStop::Complete { final_loss } => {
                    println!("silo training complete, final loss: {final_loss:.6}");
                }
                train::TrainStop::NanLoss(epoch) => {
                    eprintln!("training aborted: NaN loss at epoch {}", epoch + 1);
                    std::process::exit(1);
                }
                train::TrainStop::Plateau { epoch, loss } => {
                    eprintln!("training aborted: plateau at epoch {} (loss={loss:.8})", epoch + 1);
                    std::process::exit(1);
                }
            }
        }
        Cmd::Silo { class, config, count, palette: palette_name, output, seed, steps } => {
            let router = class_router::ClassRouter::load(&config)?;
            let dc = router.route(&class)
                .ok_or_else(|| anyhow::anyhow!("no silo found for class '{class}' in {config}"))?;

            if !std::path::Path::new(&dc.model_path).exists() {
                anyhow::bail!("model not found: {}. Train it first: pixel-forge train-silo --class {}", dc.model_path, class);
            }

            let steps = steps.unwrap_or(dc.steps);
            let pal = palette::load_palette(&palette_name)?;

            let raw_images = train::sample_micro(&dc.model_path, &dc.channels, 32, count, steps, seed)?;
            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, 32);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if count == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&processed, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output}");
        }
        Cmd::Anvil { class, anvil, count, steps, palette: palette_name, output, cfg } => {
            let cond = class_cond::lookup(&class.to_lowercase());

            let pal = palette::load_palette(&palette_name)?;
            let raw_images = moe::anvil_sample(&anvil, &cond, 32, count, steps, cfg)?;

            let processed: Vec<image::RgbaImage> = raw_images
                .into_iter()
                .map(|img| {
                    let snapped = grid::snap_to_grid(&img, 32);
                    palette::quantize(&snapped, &pal)
                })
                .collect();

            if count == 1 {
                processed[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&processed, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output}");
        }
        Cmd::TrainFleet { data, epochs, lr, nodes, dry_run } => {
            // Distributed training: each GPU trains a different model tier.
            // lf (RTX 3070 8GB)    → Anvil (16.9M params)
            // bt (RX 5700 XT 8GB)  → Quench (5.83M params) via Vulkan
            // gd (RTX 3050 Ti 4GB) → Cinder (1.09M params)

            struct FleetJob {
                node: &'static str,
                host: &'static str,
                model_tier: &'static str,
                model_flag: &'static str,
                output: &'static str,
                batch_size: u32,
                features: &'static str,
                device_flag: &'static str,
            }

            let jobs = vec![
                FleetJob {
                    node: "lf", host: "lf", model_tier: "Anvil",
                    model_flag: "--anvil", output: "pixel-forge-anvil.safetensors",
                    batch_size: 4, features: "cuda", device_flag: "",
                },
                FleetJob {
                    node: "bt", host: "bt", model_tier: "Quench",
                    model_flag: "--medium", output: "pixel-forge-quench.safetensors",
                    batch_size: 8, features: "vulkan", device_flag: "--device vulkan",
                },
                FleetJob {
                    node: "gd", host: "gd", model_tier: "Cinder",
                    model_flag: "", output: "pixel-forge-cinder.safetensors",
                    batch_size: 16, features: "cuda", device_flag: "",
                },
            ];

            let filter: Option<Vec<&str>> = nodes.as_ref().map(|n| n.split(',').collect());

            println!("train-fleet: dispatching to {} nodes", jobs.len());
            println!("  data={data}  epochs={epochs}  lr={lr:.1e}");
            println!();

            let mut handles = Vec::new();

            for job in &jobs {
                if let Some(ref f) = filter {
                    if !f.contains(&job.node) { continue; }
                }

                let cmd = format!(
                    "cd ~/pixel-forge && ./target/release/pixel-forge train \
                     --data {data} {} --epochs {epochs} --lr {lr} \
                     --batch-size {} --no-ema --sponge \
                     -o {} {device} 2>&1",
                    job.model_flag, job.batch_size, job.output,
                    device = job.device_flag,
                );

                println!("  {} ({}):", job.node, job.model_tier);
                println!("    gpu: {}", match job.node {
                    "lf" => "RTX 3070 8GB (CUDA)",
                    "bt" => "RX 5700 XT 8GB (Vulkan/any-gpu)",
                    "gd" => "RTX 3050 Ti 4GB (CUDA)",
                    _ => "unknown",
                });
                println!("    bs={} features={}", job.batch_size, job.features);
                println!("    cmd: {cmd}");
                println!();

                if dry_run { continue; }

                let host = job.host.to_string();
                let ssh_cmd = cmd.clone();
                let node_name = job.node.to_string();
                let tier = job.model_tier.to_string();

                let handle = std::thread::spawn(move || {
                    println!("[{node_name}] starting {tier} training...");
                    let result = std::process::Command::new("ssh")
                        .args([&host, &ssh_cmd])
                        .stdout(std::process::Stdio::inherit())
                        .stderr(std::process::Stdio::inherit())
                        .status();
                    match result {
                        Ok(status) => println!("[{node_name}] {tier} exited: {status}"),
                        Err(e) => eprintln!("[{node_name}] {tier} SSH failed: {e}"),
                    }
                });
                handles.push(handle);
            }

            if dry_run {
                println!("dry run complete. Add --no-dry-run to start training.");
            } else {
                for h in handles {
                    let _ = h.join();
                }
                println!("train-fleet: all nodes complete");
            }
        }
        Cmd::Tiered {
            class, silo_dir, detail_model, count, palette: palette_name,
            silo_steps, detail_steps, refine_t, cfg, seed, palette_model, output,
        } => {
            let cond = class_cond::lookup(&class.to_lowercase());
            let config = tiered_pipeline::TieredConfig {
                silo_dir: std::path::PathBuf::from(&silo_dir),
                detail_model: std::path::PathBuf::from(&detail_model),
                palette_name,
                silo_steps,
                detail_steps,
                refine_from_t: refine_t,
                cfg_scale: cfg,
                seed,
                palette_model: palette_model.map(std::path::PathBuf::from),
                palette_colors: 8,
            };
            let images = tiered_pipeline::tiered_generate(&cond, &config, count)?;
            if images.len() == 1 {
                images[0].save(&output)?;
            } else {
                let sheet_img = sheet::pack_grid(&images, 8);
                sheet_img.save(&output)?;
            }
            println!("saved: {output}");
        }
        #[cfg(feature = "vulkan")]
        Cmd::TrainCinderVk { data, condition, output, resume, epochs, batch_size, lr } => {
            let cfg = pixel_forge::vulkan_tiny::VulkanTinyTrainConfig {
                data_dir: data,
                cond_dir: condition,
                output,
                resume,
                epochs,
                batch_size,
                lr,
                img_size: 32,
            };
            pixel_forge::vulkan_tiny::vulkan_tiny_train(&cfg)?;
        }
        Cmd::PrepSiloCond { data, output } => {
            use image::imageops::FilterType;

            // Load source dataset (uses cache — guaranteed same count as training).
            let src = train::preprocess(&data, 32)?;
            let n = src.super_ids.len();
            let sz = 32usize;
            let n_px = sz * sz; // 1024 pixels per channel

            println!("prep-silo-cond: coarsening {n} samples (32→16→32 nearest-neighbor)...");

            let mut cond_pixels = Vec::with_capacity(n * src.stride);

            for i in 0..n {
                let base = i * src.stride;
                // Build RGBA image from channel-first float pixels
                let mut img = image::RgbaImage::new(sz as u32, sz as u32);
                for y in 0..sz {
                    for x in 0..sz {
                        let idx = y * sz + x;
                        let r = (src.pixels[base + idx]             * 255.0) as u8;
                        let g = (src.pixels[base + n_px + idx]     * 255.0) as u8;
                        let b = (src.pixels[base + 2 * n_px + idx] * 255.0) as u8;
                        img.put_pixel(x as u32, y as u32, image::Rgba([r, g, b, 255]));
                    }
                }
                // 32x32 → 16x16 → 32x32: captures coarse shape/color like silo outputs
                let small  = image::imageops::resize(&img,  16, 16, FilterType::Nearest);
                let coarse = image::imageops::resize(&small, 32, 32, FilterType::Nearest);
                // Convert back to channel-first f32
                let mut px = vec![0.0f32; src.stride];
                for y in 0..sz {
                    for x in 0..sz {
                        let p = coarse.get_pixel(x as u32, y as u32);
                        let idx = y * sz + x;
                        px[idx]             = p[0] as f32 / 255.0;
                        px[n_px + idx]     = p[1] as f32 / 255.0;
                        px[2 * n_px + idx] = p[2] as f32 / 255.0;
                    }
                }
                cond_pixels.extend_from_slice(&px);
            }

            // Build conditioning dataset with same metadata as source
            let cond_ds = train::PackedDataset {
                version: src.version,
                img_size: src.img_size,
                pixels: cond_pixels,
                labels: src.labels.clone(),
                super_ids: src.super_ids.clone(),
                tags: src.tags.clone(),
                stride: src.stride,
            };

            // Write cache directly — preprocess() will load it without scanning dirs
            std::fs::create_dir_all(&output)?;
            let encoded = bincode::serialize(&cond_ds)?;
            let compressed = zstd::encode_all(encoded.as_slice(), 3)?;
            let cache_path = format!("{output}/dataset.bin.zst");
            std::fs::write(&cache_path, &compressed)?;
            println!("prep-silo-cond: {n} samples → {cache_path} ({:.1} MB)",
                compressed.len() as f64 / 1_048_576.0);
        }
        Cmd::InspectCond { model, class, count, seed } => {
            use candle_core::{DType, Tensor};
            use candle_nn::{VarBuilder, VarMap};

            let in_ch = quantize::detect_in_channels(&model);
            anyhow::ensure!(in_ch == 6, "InspectCond needs a 6ch model; got {in_ch}ch");

            let device = pipeline::best_device();
            println!("inspect-cond: {model} ({in_ch}ch), class={class}, count={count}, device={device:?}");

            let mut vm = VarMap::new();
            let vb = VarBuilder::from_varmap(&vm, DType::F32, &device);
            let net = tiny_unet::TinyUNet::with_channels(vb, 6)?;
            quantize::load_varmap(&mut vm, &model)?;

            let cond = class_cond::lookup(&class.to_lowercase());
            let (super_t, tags_t, _, _) = train::cond_tensors(&cond, 1, &device)?;

            // Resolve a real silo for the class to produce a "real" conditioning hint.
            let silo = tiered_pipeline::resolve_silo(&cond.name, std::path::Path::new("models"));
            let real_hint = match &silo {
                Some((silo_path, channels)) => {
                    println!("  real cond from silo: {}", silo_path.display());
                    let img = train::sample_micro(
                        silo_path.to_str().unwrap_or(""),
                        channels, 32, 1, 20, Some(seed),
                    )?;
                    let (w, h) = img[0].dimensions();
                    let n = (w * h) as usize;
                    let mut px = vec![0.0f32; 3 * n];
                    for y in 0..h {
                        for x in 0..w {
                            let p = img[0].get_pixel(x, y);
                            let i = (y * w + x) as usize;
                            px[i]         = p[0] as f32 / 255.0;
                            px[n + i]     = p[1] as f32 / 255.0;
                            px[2 * n + i] = p[2] as f32 / 255.0;
                        }
                    }
                    Tensor::from_vec(px, (1, 3, h as usize, w as usize), &device)?
                }
                None => {
                    println!("  no silo for class — using deterministic gradient as real cond");
                    let n = 32 * 32;
                    let px: Vec<f32> = (0..3 * n).map(|i| (i as f32) / (3.0 * n as f32)).collect();
                    Tensor::from_vec(px, (1usize, 3usize, 32usize, 32usize), &device)?
                }
            };

            // Random "fake" hint at same magnitude — Gaussian noise scaled to roughly [0,1] range.
            let fake_hint = train::seeded_noise(Some(seed.wrapping_add(999)), cond.super_id, 0, 32, &device)?
                .broadcast_mul(&Tensor::new(&[0.5f32], &device)?)?
                .broadcast_add(&Tensor::new(&[0.5f32], &device)?)?
                .clamp(0.0, 1.0)?;

            let mut total_diff = 0.0f32;
            let mut total_real_mag = 0.0f32;
            for i in 0..count {
                let noisy = train::seeded_noise(Some(seed + i as u64), cond.super_id, i, 32, &device)?;
                let amount = 0.45f32; // mid-noise level — where conditioning matters most
                let t = Tensor::new(&[amount], &device)?;

                let in_real = Tensor::cat(&[&real_hint, &noisy], 1)?;
                let in_fake = Tensor::cat(&[&fake_hint, &noisy], 1)?;

                let pred_real = net.forward(&in_real, &t, &super_t, &tags_t)?;
                let pred_fake = net.forward(&in_fake, &t, &super_t, &tags_t)?;

                let diff = (&pred_real - &pred_fake)?.abs()?.mean_all()?.to_scalar::<f32>()?;
                let mag = pred_real.abs()?.mean_all()?.to_scalar::<f32>()?;
                total_diff += diff;
                total_real_mag += mag;
                println!("  sample {i}: |Δ|={diff:.5}  |pred|={mag:.5}  ratio={:.3}", diff / mag.max(1e-8));
            }

            let avg_diff = total_diff / count as f32;
            let avg_mag = total_real_mag / count as f32;
            let ratio = avg_diff / avg_mag.max(1e-8);
            println!("  avg |Δ|={avg_diff:.5}  avg |pred|={avg_mag:.5}  ratio={ratio:.3}");
            println!();
            if ratio < 0.01 {
                println!("VERDICT: model is IGNORING the conditioning channel (Δ < 1% of pred magnitude).");
            } else if ratio < 0.05 {
                println!("VERDICT: model is barely using conditioning ({:.1}% effect).", ratio * 100.0);
            } else if ratio < 0.20 {
                println!("VERDICT: model is using conditioning ({:.1}% effect — moderate).", ratio * 100.0);
            } else {
                println!("VERDICT: model is strongly using conditioning ({:.1}% effect).", ratio * 100.0);
            }
        }
        Cmd::TrainPaletteNet { data, output, epochs, batch_size, lr, palette_colors } => {
            palette_net::train_palette_net(&data, &output, epochs, batch_size, lr, palette_colors)?;
        }
    }

    Ok(())
}
