kibble 0.1.0

use clap::{Parser, Subcommand};
use kibble::{
    ask, bench, build, caps, clean, cluster, config, crawl, dotenv, eval, extract,
    fetch, index, ingest, init, mcp, pack, retrieve, serve, soul, train, tune,
};

#[derive(Parser)]
#[command(
    name = "kibble",
    version,
    about = "chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit",
    before_help = "  U・ᴥ・U  🦴 kibble 🦴 — fetch · chew · digest"
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Apply text cleaning rules to stdin, writing the result to stdout.
    Clean,
    /// Build the training-dataset JSONL from ingested raw docs.
    Build,
    /// Pack the build output into a Kaggle dataset bundle (reads [pack] in kibble.toml).
    Pack,
    /// Grade the built dataset's quality (SFT readiness + dedup/leakage/length/balance/diversity).
    Eval {
        /// Exit nonzero if the dataset fails the quality gate.
        #[arg(long)]
        strict: bool,
    },
    /// Ingest local sources (twitter, textfiles) into data/raw/local/.
    Ingest {
        /// Which source to ingest.
        #[arg(value_enum)]
        source: IngestSource,
    },
    /// Run the ingest HTTP API (POST /ingest, GET /jobs/{id}); needs KIBBLE_API_TOKEN.
    Serve,
    /// Fetch a URL (http file/dir, HuggingFace dataset, git, archive) into the ingest tree.
    Fetch {
        /// URL or `user/name` to fetch.
        url: String,
        /// Optional bundle name (default: a slug of the URL / the HF repo name).
        #[arg(long)]
        name: Option<String>,
        /// Column-map override for exotic dataset schemas: user=<col>,assistant=<col>[,system=<col>].
        #[arg(long)]
        map: Option<String>,
        /// Sample at most N rows from an HF dataset (giants stream/stop early instead of full download).
        #[arg(long)]
        max_rows: Option<usize>,
    },
    /// Extract text from artifacts (image OCR, audio transcription, html/text) into data/extracted/.
    Extract {
        /// File or directory of artifacts to extract.
        path: String,
        /// Output dir override (default: [extract].out).
        #[arg(long)]
        out: Option<String>,
        /// Skip artifacts whose type isn't supported yet instead of erroring.
        #[arg(long)]
        skip_unsupported: bool,
    },
    /// Manage agent capabilities (skills) ingested from sources.
    Caps {
        #[command(subcommand)]
        action: CapsAction,
    },
    /// Benchmark a served model (OpenAI-compatible endpoint) against the [[bench.benchmark]] suite.
    Bench {
        /// Exit nonzero if the model fails the benchmark gate.
        #[arg(long)]
        strict: bool,
        /// Stream the model's output to stderr during research-method runs.
        #[arg(long)]
        stream: bool,
    },
    /// Compile soul.toml into serving artifacts (generation_config, system_prompt, tools, chat_template, manifest).
    Soul {
        #[command(subcommand)]
        action: SoulAction,
    },
    /// Run an MCP (Model Context Protocol) server over stdio, exposing KIBBLE's tools.
    Mcp,
    /// Crawl a website (BFS, same-host, markdown) into the ingest tree for `build`.
    Crawl {
        /// Seed URL to crawl.
        url: String,
        #[arg(long)]
        depth: Option<usize>,
        #[arg(long = "max-pages")]
        max_pages: Option<usize>,
        #[arg(long)]
        out: Option<String>,
        /// Follow links to other hosts (default: same host only).
        #[arg(long = "all-hosts")]
        all_hosts: bool,
    },
    /// Scaffold a ready-to-run kibble project (kibble.toml + sample corpus) in the current dir.
    Init {
        /// Overwrite an existing kibble.toml.
        #[arg(long)]
        force: bool,
    },
    /// Build the retrieval index over the corpus (or an explicit path).
    Index {
        /// Optional dir/file to index (default: [index].sources).
        path: Option<String>,
    },
    /// Search the retrieval index (hybrid semantic + BM25).
    Search {
        /// The query text.
        query: String,
        /// Number of results (default 10).
        #[arg(long, default_value_t = 10)]
        k: usize,
        /// Emit JSON instead of a table.
        #[arg(long)]
        json: bool,
    },
    /// Answer a question grounded in the indexed corpus (agentic RAG, cited).
    Ask {
        /// The question.
        query: String,
        /// Passages retrieved per search (default [ask].k).
        #[arg(long)]
        k: Option<usize>,
        /// Emit JSON instead of prose.
        #[arg(long)]
        json: bool,
        /// If not found in the corpus, answer from general knowledge (flagged).
        #[arg(long)]
        chance: bool,
        /// Also print every retrieved passage.
        #[arg(long = "show-context")]
        show_context: bool,
        /// Allow web tools (web_search/fetch_page) for this query.
        #[arg(long)]
        web: bool,
        /// Force corpus-only (overrides [web].default).
        #[arg(long = "no-web")]
        no_web: bool,
        /// Force streaming the answer token-by-token (overrides [ask].stream).
        #[arg(long)]
        stream: bool,
        /// Force buffered (non-streamed) output.
        #[arg(long = "no-stream")]
        no_stream: bool,
    },
    /// Cluster the built dataset's rows into topics (embeddings; writes data/clusters.json).
    Cluster {
        /// Number of topics (default [cluster].k).
        #[arg(long)]
        k: Option<usize>,
        /// Emit JSON instead of a table.
        #[arg(long)]
        json: bool,
    },
    /// Generate the training sprint curriculum + config.yaml from the built dataset ([tune] in kibble.toml).
    Tune,
    /// Run the configured [train].command against the tune package (see docs/TRAIN.md).
    Train {
        /// Print the resolved command and working directory without running it.
        #[arg(long)]
        dry_run: bool,
        /// Extra args appended to the command (after `--`); placeholders are substituted.
        #[arg(last = true)]
        extra: Vec<String>,
    },
}

#[derive(Subcommand)]
enum SoulAction {
    /// Build the serving-artifact bundle from a soul.toml.
    Build {
        /// Path to the soul.toml (default ./soul.toml).
        #[arg(long)]
        soul: Option<String>,
        /// Output dir (default soul/dist).
        #[arg(long)]
        out: Option<String>,
    },
}

#[derive(Subcommand)]
enum CapsAction {
    /// Detect skills in a source (path or URL) and install them into the registry.
    Install {
        src: String,
        #[arg(long)]
        force: bool,
        #[arg(long)]
        project: bool,
    },
    /// Dry-run: list skills a source WOULD install, writing nothing.
    Scan {
        src: String,
        #[arg(long)]
        project: bool,
    },
    /// List installed skills.
    List {
        #[arg(long)]
        project: bool,
    },
    /// Remove an installed skill by name.
    Remove {
        name: String,
        #[arg(long)]
        project: bool,
    },
}

#[derive(clap::ValueEnum, Clone)]
enum IngestSource {
    Twitter,
    Textfiles,
}

/// Turn a command's `Err` into a clean `kibble: <msg>` on stderr + exit code 1,
/// instead of the Rust panic (`.expect` → exit 101 + backtrace) that external
/// callers can't parse. See issue #60.
trait OrExit<T> {
    fn or_exit(self) -> T;
}
impl<T, E: std::fmt::Display> OrExit<T> for Result<T, E> {
    fn or_exit(self) -> T {
        self.unwrap_or_else(|e| {
            eprintln!("kibble: {e}");
            std::process::exit(1);
        })
    }
}

#[tokio::main]
async fn main() {
    // Load secrets from a nearby `.env` before anything reads the environment.
    // Real env vars always win; this only fills in what isn't already set.
    dotenv::load();
    let cli = Cli::parse();
    match cli.command {
        Command::Clean => {
            use std::io::Read;
            let mut input = String::new();
            std::io::stdin().read_to_string(&mut input).or_exit();
            print!("{}", clean::clean_text(&input));
        }
        Command::Build => {
            let stats = build::run_build(std::path::Path::new(".")).await.or_exit();
            if stats.total_documents == 0 {
                eprintln!("kibble: no sources found — add [[source]] to kibble.toml or drop files in data/raw/.");
                eprintln!("  try: kibble init   (scaffolds a sample project you can build)");
                return;
            }
            println!(
                "Build complete: {} docs -> train {} / valid {} / test {} (dropped {} dup, {} near-dup, {} semantic, {} filtered, {} rebalanced)",
                stats.total_documents, stats.train, stats.valid, stats.test,
                stats.dropped_duplicates, stats.dropped_near_duplicates,
                stats.dropped_semantic_duplicates, stats.dropped_filtered,
                stats.dropped_topic_rebalanced
            );
            for (name, t, v, te) in &stats.sources {
                println!("  source {name}: train {t} / valid {v} / test {te}");
            }
        }
        Command::Pack => {
            pack::run_pack(std::path::Path::new(".")).or_exit();
        }
        Command::Eval { strict } => {
            let r = eval::run_eval(std::path::Path::new("."), strict).await.or_exit();
            println!(
                "Dataset quality: {}/100 — {} ({} rows · dup {:.1}% · leakage {:.1}% · short {:.1}%)",
                r.quality_score, if r.overall { "PASS" } else { "FAIL" },
                r.total_rows, r.duplicate_rate * 100.0, r.leakage_rate * 100.0, r.short_rate * 100.0
            );
            if strict && !r.overall { std::process::exit(1); }
        }
        Command::Serve => {
            serve::run_serve(std::path::Path::new(".")).await.or_exit();
        }
        Command::Fetch { url, name, map, max_rows } => {
            fetch::run_fetch(std::path::Path::new("."), &url, name.as_deref(), map.as_deref(), max_rows)
                .await
                .or_exit();
        }
        Command::Extract { path, out, skip_unsupported } => {
            let n = extract::run_extract(std::path::Path::new("."), std::path::Path::new(&path), out.as_deref(), skip_unsupported)
                .await
                .or_exit();
            println!("Extracted {n} artifact(s) -> data/extracted (or --out)");
        }
        Command::Caps { action } => {
            let repo_root = std::path::Path::new(".");
            let cfg = config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
            match action {
                CapsAction::Install { src, force, project } => {
                    let names = caps::install_source(repo_root, &cfg, &src, force, project, false)
                        .await.or_exit();
                    println!("Installed {} skill(s): {}", names.len(), names.join(", "));
                }
                CapsAction::Scan { src, project } => {
                    let names = caps::install_source(repo_root, &cfg, &src, false, project, true)
                        .await.or_exit();
                    println!("Would install {} skill(s): {}", names.len(), names.join(", "));
                }
                CapsAction::List { project } => {
                    let root = caps::caps_root(&cfg.caps, if project { Some(repo_root) } else { None });
                    for s in caps::list(&root) {
                        println!("{:<24} {}", s.name, s.description);
                    }
                }
                CapsAction::Remove { name, project } => {
                    let root = caps::caps_root(&cfg.caps, if project { Some(repo_root) } else { None });
                    let removed = caps::remove(&root, &name).or_exit();
                    println!("{}", if removed { format!("Removed {name}") } else { format!("{name} not found") });
                }
            }
        }
        Command::Bench { strict, stream } => {
            let r = bench::run_bench(std::path::Path::new("."), strict, stream).await.or_exit();
            println!("Benchmarks: {}/100 — {} ({} suites)", r.quality_score,
                if r.overall { "PASS" } else { "FAIL" }, r.benchmarks.len());
            for b in &r.benchmarks {
                println!("  {:<20} {:.3} (>= {:.3}) · {:.1} tok/s · {}", b.name, b.score, b.threshold, b.mean_tok_s, if b.ok {"ok"} else {"FAIL"});
            }
            if strict && !r.overall { std::process::exit(1); }
        }
        Command::Soul { action } => match action {
            SoulAction::Build { soul, out } => {
                let soul_path = std::path::PathBuf::from(soul.unwrap_or_else(|| "soul.toml".into()));
                let out_dir = std::path::PathBuf::from(out.unwrap_or_else(|| "soul/dist".into()));
                let m = soul::run_soul_build(&soul_path, &out_dir).or_exit();
                println!("Soul '{}' v{} → {} ({} files, sha {})", m.name, m.version, out_dir.display(), m.files.len(), &m.sha256[..12.min(m.sha256.len())]);
            }
        },
        Command::Ingest { source } => {
            use std::path::Path;
            match source {
                IngestSource::Twitter => {
                    let n = ingest::ingest_twitter(
                        Path::new("twitter/data"),
                        Path::new("data/raw/local/twitter"),
                    )
                    .or_exit();
                    println!("Wrote {n} tweet files to data/raw/local/twitter");
                }
                IngestSource::Textfiles => {
                    let n = ingest::ingest_textfiles(
                        Path::new("textfiles"),
                        Path::new("data/raw/local/textfiles"),
                    )
                    .or_exit();
                    println!("Wrote {n} textfiles to data/raw/local/textfiles");
                }
            }
        }
        Command::Mcp => {
            mcp::run_mcp(std::path::Path::new(".")).await.or_exit();
        }
        Command::Crawl { url, depth, max_pages, out, all_hosts } => {
            let st = crawl::run_crawl(std::path::Path::new("."), &url, depth, max_pages, out, all_hosts).await.or_exit();
            println!("Crawled {} → {} pages, {} skipped → {}", st.seed, st.pages, st.skipped, st.out_dir);
        }
        Command::Init { force } => {
            let r = init::run_init(std::path::Path::new("."), force).or_exit();
            if r.config_written { println!("  ✓ wrote kibble.toml"); }
            if !r.samples_written.is_empty() {
                println!("  ✓ wrote {} sample docs under examples/notes/", r.samples_written.len());
            }
            for s in &r.samples_skipped { println!("  · skipped {s} (already exists)"); }
            println!("\n  next: kibble index examples && kibble search \"ownership\"");
        }
        Command::Index { path } => {
            let root = std::path::Path::new(".");
            let explicit = path.as_ref().map(std::path::PathBuf::from);
            let n = index::build_index(root, explicit.as_deref()).await.or_exit();
            println!("Indexed {n} chunk(s) -> data/index (or [index].dir)");
        }
        Command::Search { query, k, json } => {
            let hits = retrieve::search(std::path::Path::new("."), &query, k).await.or_exit();
            if json {
                let arr: Vec<serde_json::Value> = hits.iter().map(|h| serde_json::json!({
                    "score": h.score, "source": h.chunk.source, "doc_id": h.chunk.doc_id,
                    "title": h.chunk.title, "text": h.chunk.text,
                })).collect();
                println!("{}", serde_json::to_string_pretty(&arr).unwrap());
            } else if hits.is_empty() {
                println!("No results.");
            } else {
                for (i, h) in hits.iter().enumerate() {
                    let snippet: String = h.chunk.text.chars().take(120).collect();
                    let snippet = snippet.replace('\n', " ");
                    println!("{:>2}. [{:.4}] {} — {}\n    {}", i + 1, h.score, h.chunk.source, h.chunk.title, snippet);
                }
            }
        }
        Command::Ask { query, k, json, chance, show_context, web, no_web, stream, no_stream } => {
            ask::run(std::path::Path::new("."), &query, k, json, chance, show_context, web, no_web, stream, no_stream).await.or_exit();
        }
        Command::Cluster { k, json } => {
            if let Some(r) = cluster::run_cluster(std::path::Path::new("."), k).await.or_exit() {
                let mut topics: Vec<(usize, &String, &String)> = (0..r.k)
                    .map(|i| (r.sizes[i], &r.labels[i], &r.samples[i]))
                    .collect();
                topics.sort_by_key(|b| std::cmp::Reverse(b.0));
                if json {
                    let arr: Vec<serde_json::Value> = topics.iter().map(|(size, label, sample)| serde_json::json!({
                        "label": label, "size": size, "sample": sample
                    })).collect();
                    println!("{}", serde_json::to_string_pretty(&serde_json::json!({ "k": r.k, "topics": arr })).unwrap());
                } else {
                    for (size, label, sample) in &topics {
                        let snip: String = sample.chars().take(60).collect();
                        println!("  [{label}]  {size} rows · {}", snip.replace('\n', " "));
                    }
                    println!("\n{} topics over {} rows -> clusters.json", r.k, r.sizes.iter().sum::<usize>());
                }
            }
        }
        Command::Tune => {
            tune::run_tune(std::path::Path::new(".")).or_exit();
        }
        Command::Train { dry_run, extra } => {
            if let Err(e) = train::run_train(std::path::Path::new("."), dry_run, &extra) {
                eprintln!("kibble: {e}");
                std::process::exit(1);
            }
        }
    }
}