kibble 0.1.0

use std::io::Write;
use std::path::Path;
use scraper::{Html, Selector};
use crate::config::load_config;
use crate::git::{cache_slug, clone_repo};
use crate::net::{build_client, resolve_proxy};

pub async fn download_to_file(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<u64> {
    let mut resp = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("download failed for {url}: {e}")))?;

    if let Some(len) = resp.content_length() {
        if len > max_bytes {
            return Err(std::io::Error::other(format!(
                "{url}: content-length {len} exceeds max_bytes {max_bytes}"
            )));
        }
    }
    if let Some(parent) = dest.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let mut file = std::fs::File::create(dest)?;
    let mut total: u64 = 0;
    loop {
        let chunk = resp
            .chunk()
            .await
            .map_err(|e| std::io::Error::other(format!("download read failed: {e}")))?;
        let Some(chunk) = chunk else { break };
        total += chunk.len() as u64;
        if total > max_bytes {
            drop(file);
            let _ = std::fs::remove_file(dest);
            return Err(std::io::Error::other(format!(
                "{url}: download exceeded max_bytes {max_bytes}"
            )));
        }
        file.write_all(&chunk)?;
    }
    Ok(total)
}

fn is_safe_rel(path: &Path) -> bool {
    !path.is_absolute()
        && path.components().all(|c| !matches!(c, std::path::Component::ParentDir | std::path::Component::Prefix(_) | std::path::Component::RootDir))
}

fn extract_zip(archive: &Path, dest: &Path) -> std::io::Result<usize> {
    let file = std::fs::File::open(archive)?;
    let mut zip = zip::ZipArchive::new(file).map_err(|e| std::io::Error::other(e.to_string()))?;
    let mut count = 0;
    for i in 0..zip.len() {
        let mut entry = zip.by_index(i).map_err(|e| std::io::Error::other(e.to_string()))?;
        let name = entry.enclosed_name().ok_or_else(|| {
            std::io::Error::new(std::io::ErrorKind::InvalidData, "unsafe zip entry path")
        })?;
        if !is_safe_rel(&name) {
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "unsafe zip entry path"));
        }
        let out = dest.join(&name);
        if entry.is_dir() {
            std::fs::create_dir_all(&out)?;
        } else {
            if let Some(p) = out.parent() {
                std::fs::create_dir_all(p)?;
            }
            let mut f = std::fs::File::create(&out)?;
            std::io::copy(&mut entry, &mut f)?;
            count += 1;
        }
    }
    Ok(count)
}

fn extract_tar<R: std::io::Read>(reader: R, dest: &Path) -> std::io::Result<usize> {
    let mut archive = tar::Archive::new(reader);
    archive.set_preserve_permissions(false);
    let mut count = 0;
    for entry in archive.entries()? {
        let mut entry = entry?;
        let etype = entry.header().entry_type();
        if etype.is_symlink() || etype.is_hard_link() {
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "unsafe tar link entry"));
        }
        let path = entry.path()?.into_owned();
        if !is_safe_rel(&path) {
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "unsafe tar entry path"));
        }
        let is_file = etype.is_file();
        if entry.unpack_in(dest)? && is_file {
            count += 1;
        }
    }
    Ok(count)
}

pub fn extract_archive(archive: &Path, dest: &Path) -> std::io::Result<usize> {
    std::fs::create_dir_all(dest)?;
    let name = archive.file_name().and_then(|n| n.to_str()).unwrap_or("").to_lowercase();
    if name.ends_with(".zip") {
        extract_zip(archive, dest)
    } else if name.ends_with(".tar.gz") || name.ends_with(".tgz") {
        let f = std::fs::File::open(archive)?;
        extract_tar(flate2::read::GzDecoder::new(f), dest)
    } else if name.ends_with(".tar") {
        let f = std::fs::File::open(archive)?;
        extract_tar(f, dest)
    } else {
        Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, format!("not an archive: {name}")))
    }
}

fn is_archive(name: &str) -> bool {
    let n = name.to_lowercase();
    n.ends_with(".zip") || n.ends_with(".tar.gz") || n.ends_with(".tgz") || n.ends_with(".tar")
}

pub struct FetchResult {
    pub handler: String,
    pub files: usize,
    pub bytes: u64,
}

fn filename_from_url(url: &str) -> String {
    let path = url.split(['?', '#']).next().unwrap_or(url);
    let name = path.rsplit('/').next().unwrap_or("");
    if name.is_empty() || name == "." || name == ".." || name.contains('/') || name.contains('\\') {
        "download".to_string()
    } else {
        name.to_string()
    }
}

pub async fn fetch_http_file(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    std::fs::create_dir_all(dest)?;
    let name = filename_from_url(url);
    let target = dest.join(&name);
    let bytes = download_to_file(client, url, &target, max_bytes).await?;
    if is_archive(&name) {
        let n = extract_archive(&target, dest)?;
        let _ = std::fs::remove_file(&target);
        return Ok(FetchResult { handler: "http-archive".to_string(), files: n, bytes });
    }
    Ok(FetchResult { handler: "http-file".to_string(), files: 1, bytes })
}

pub async fn fetch_http_dir(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    std::fs::create_dir_all(dest)?;
    let html = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("list failed for {url}: {e}")))?
        .text()
        .await
        .map_err(|e| std::io::Error::other(format!("list read failed: {e}")))?;

    let base = if url.ends_with('/') { url.to_string() } else { format!("{url}/") };
    // Collect hrefs into owned Strings, then drop the DOM before the await loop (keeps the future Send).
    let hrefs: Vec<String> = {
        let doc = Html::parse_document(&html);
        let sel = Selector::parse("a").unwrap();
        doc.select(&sel)
            .filter_map(|a| a.value().attr("href").map(|h| h.to_string()))
            .collect()
    };
    let mut files = 0;
    let mut bytes = 0u64;
    for href in hrefs {
        if href.starts_with("..") || href.starts_with('?') || href.starts_with('#')
            || href.starts_with('/') || href.contains("://") || href.ends_with('/')
        {
            continue;
        }
        let file_url = format!("{base}{href}");
        let target = dest.join(filename_from_url(&href));
        match download_to_file(client, &file_url, &target, max_bytes).await {
            Ok(n) => { files += 1; bytes += n; }
            Err(e) => eprintln!("kibble: skipping {file_url}: {e}"),
        }
    }
    Ok(FetchResult { handler: "http-dir".to_string(), files, bytes })
}

pub async fn fetch_hf_dataset(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    // derive scheme://host base and the dataset id (everything after /datasets/)
    let (scheme, rest) = url.split_once("://").unwrap_or(("https", url));
    let host = rest.split('/').next().unwrap_or("");
    let base = format!("{scheme}://{host}");
    let id = rest
        .split_once("/datasets/")
        .map(|(_, r)| r.trim_end_matches('/'))
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidInput, "not a hf dataset url"))?;

    std::fs::create_dir_all(dest)?;
    let api = format!("{base}/api/datasets/{id}");
    let meta = client
        .get(&api)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("hf api failed for {api}: {e}")))?
        .text()
        .await
        .map_err(|e| std::io::Error::other(format!("hf api read failed: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&meta)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

    let mut files = 0;
    let mut bytes = 0u64;
    if let Some(siblings) = v.get("siblings").and_then(|s| s.as_array()) {
        for sib in siblings {
            let Some(rfile) = sib.get("rfilename").and_then(|r| r.as_str()) else { continue };
            // safety: skip unsafe relative paths
            if !is_safe_rel(Path::new(rfile)) {
                continue;
            }
            let file_url = format!("{base}/datasets/{id}/resolve/main/{rfile}");
            let target = dest.join(rfile);
            match download_to_file(client, &file_url, &target, max_bytes).await {
                Ok(n) => { files += 1; bytes += n; }
                Err(e) => eprintln!("kibble: skipping {file_url}: {e}"),
            }
        }
    }
    Ok(FetchResult { handler: "huggingface".to_string(), files, bytes })
}

fn scheme_host(url: &str) -> String {
    let (scheme, rest) = url.split_once("://").unwrap_or(("https", url));
    let host = rest.split('/').next().unwrap_or("");
    format!("{scheme}://{host}")
}

fn looks_like_gdrive(url: &str) -> bool {
    let l = url.to_lowercase();
    l.contains("drive.google.com") || l.contains("docs.google.com")
}

fn gdrive_extract_id(url: &str) -> Option<String> {
    if let Some(rest) = url.split_once("/file/d/") {
        let id = rest.1.split('/').next().unwrap_or("");
        if !id.is_empty() { return Some(id.to_string()); }
    }
    if let Some(rest) = url.split_once("id=") {
        let id = rest.1.split(['&', '#']).next().unwrap_or("");
        if !id.is_empty() { return Some(id.to_string()); }
    }
    None
}

pub async fn fetch_gdrive(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    let id = gdrive_extract_id(url)
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidInput, format!("no drive file id in {url}")))?;
    let base = scheme_host(url);
    // confirm=t clears the large-file virus-scan interstitial for public files.
    let dl = format!("{base}/uc?export=download&id={id}&confirm=t");
    std::fs::create_dir_all(dest)?;
    let target = dest.join(&id);
    let bytes = download_to_file(client, &dl, &target, max_bytes).await?;
    Ok(FetchResult { handler: "gdrive".to_string(), files: 1, bytes })
}

fn looks_like_dropbox(url: &str) -> bool {
    url.to_lowercase().contains("dropbox.com")
}

fn dropbox_direct_url(url: &str) -> String {
    if url.contains("dl=0") {
        url.replace("dl=0", "dl=1")
    } else if url.contains("dl=1") {
        url.to_string()
    } else if url.contains('?') {
        format!("{url}&dl=1")
    } else {
        format!("{url}?dl=1")
    }
}

pub async fn fetch_dropbox(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    let direct = dropbox_direct_url(url);
    let r = fetch_http_file(client, &direct, dest, max_bytes).await?;
    Ok(FetchResult { handler: "dropbox".to_string(), ..r })
}

fn looks_like_git(url: &str) -> bool {
    let l = url.to_lowercase();
    (l.starts_with("http://") || l.starts_with("https://"))
        && (l.contains("github.com") || l.ends_with(".git"))
}

fn looks_like_hf_dataset(url: &str) -> bool {
    let l = url.to_lowercase();
    l.contains("huggingface.co/datasets/")
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum HfKind {
    Dataset,
    Model,
}

#[derive(Debug, PartialEq)]
pub struct HfRef {
    pub base: String,
    pub id: String,
    pub hint: Option<HfKind>,
}

/// Recognise an HF dataset/model target: a full `huggingface.co` URL (dataset if under
/// `/datasets/`, else model), a localhost URL (for tests, same rules), or a bare `user/name`
/// shorthand (defaults base to `https://huggingface.co`, kind unknown → probe). Returns None for
/// anything not HF-shaped so the caller falls through to the existing handlers.
pub fn parse_hf_target(target: &str) -> Option<HfRef> {
    // Bare shorthand: exactly `segment/segment`, no scheme, no dot in the first segment (a dot
    // would signal a host/filename), no whitespace.
    if !target.contains("://") {
        let parts: Vec<&str> = target.split('/').collect();
        if parts.len() == 2
            && parts.iter().all(|p| !p.is_empty() && !p.contains('.') && !p.contains(' '))
        {
            return Some(HfRef {
                base: "https://huggingface.co".to_string(),
                id: target.to_string(),
                hint: None,
            });
        }
        return None;
    }

    let (scheme, rest) = target.split_once("://")?;
    let host = rest.split('/').next().unwrap_or("");
    let is_hf_host = host.contains("huggingface.co");
    let is_localhost = host.starts_with("127.0.0.1") || host.starts_with("localhost");
    if !is_hf_host && !is_localhost {
        return None;
    }
    let base = format!("{scheme}://{host}");
    let path = rest.strip_prefix(host).unwrap_or("").trim_start_matches('/').trim_end_matches('/');
    // `/datasets/` with no id (possibly a trailing slash, already trimmed) is not a valid target.
    if path == "datasets" {
        return None;
    }
    if let Some(ds) = path.strip_prefix("datasets/") {
        if ds.is_empty() {
            return None;
        }
        return Some(HfRef { base, id: ds.to_string(), hint: Some(HfKind::Dataset) });
    }
    // Non-/datasets/ path on an HF host → a model repo.
    if path.is_empty() {
        return None;
    }
    if is_localhost {
        // Localhost URLs are test-only; a dotted last segment is a filename (e.g. note.txt),
        // not a model repo. Real hosts keep dotted model names (Qwen2.5-7B, Mistral-7B-v0.1).
        let last = path.rsplit('/').next().unwrap_or(path);
        if last.contains('.') {
            return None;
        }
    }
    Some(HfRef { base, id: path.to_string(), hint: Some(HfKind::Model) })
}

async fn hf_api_exists(client: &reqwest::Client, base: &str, kind_path: &str, id: &str) -> bool {
    let url = format!("{base}/api/{kind_path}/{id}");
    matches!(
        client.get(&url).send().await.map(|r| r.status().is_success()),
        Ok(true)
    )
}

/// Decide dataset vs model. A `hint` from URL shape wins; otherwise probe the dataset API then the
/// model API.
pub async fn resolve_hf_kind(
    client: &reqwest::Client,
    base: &str,
    id: &str,
    hint: Option<HfKind>,
) -> std::io::Result<HfKind> {
    if let Some(k) = hint {
        return Ok(k);
    }
    if hf_api_exists(client, base, "datasets", id).await {
        return Ok(HfKind::Dataset);
    }
    if hf_api_exists(client, base, "models", id).await {
        return Ok(HfKind::Model);
    }
    Err(std::io::Error::new(
        std::io::ErrorKind::NotFound,
        format!("{id} is neither a dataset nor a model on {base}"),
    ))
}

/// Download every sibling file of an HF dataset or model repo into `dest`. Shared by dataset and
/// model fetches; `kind` selects the API path (`datasets` vs `models`) and the handler label.
pub async fn fetch_hf_siblings(
    client: &reqwest::Client,
    kind: HfKind,
    base: &str,
    id: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    let (api_seg, resolve_seg, handler) = match kind {
        HfKind::Dataset => ("datasets", "datasets", "hf-dataset"),
        HfKind::Model => ("models", "", "hf-model"),
    };
    std::fs::create_dir_all(dest)?;
    let api = format!("{base}/api/{api_seg}/{id}");
    let meta = client
        .get(&api)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("hf api failed for {api}: {e}")))?
        .text()
        .await
        .map_err(|e| std::io::Error::other(format!("hf api read failed: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&meta)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
    // resolve URL: datasets use /datasets/<id>/resolve/main/, models use /<id>/resolve/main/
    let resolve_base = if resolve_seg.is_empty() {
        format!("{base}/{id}/resolve/main")
    } else {
        format!("{base}/{resolve_seg}/{id}/resolve/main")
    };
    let mut files = 0;
    let mut bytes = 0u64;
    if let Some(siblings) = v.get("siblings").and_then(|s| s.as_array()) {
        for sib in siblings {
            let Some(rfile) = sib.get("rfilename").and_then(|r| r.as_str()) else { continue };
            if !is_safe_rel(Path::new(rfile)) {
                continue;
            }
            let file_url = format!("{resolve_base}/{rfile}");
            let target = dest.join(rfile);
            match download_to_file(client, &file_url, &target, max_bytes).await {
                Ok(n) => { files += 1; bytes += n; }
                Err(e) => eprintln!("kibble: skipping {file_url}: {e}"),
            }
        }
    }
    Ok(FetchResult { handler: handler.to_string(), files, bytes })
}

/// List a dataset's main-branch siblings as (rfilename, resolve-url) pairs (safe paths only).
async fn hf_dataset_siblings(
    client: &reqwest::Client,
    base: &str,
    id: &str,
) -> std::io::Result<Vec<(String, String)>> {
    let api = format!("{base}/api/datasets/{id}");
    let meta = client
        .get(&api)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("hf api failed for {api}: {e}")))?
        .text()
        .await
        .map_err(|e| std::io::Error::other(format!("hf api read failed: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&meta)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
    let mut out = Vec::new();
    if let Some(siblings) = v.get("siblings").and_then(|s| s.as_array()) {
        for sib in siblings {
            let Some(rfile) = sib.get("rfilename").and_then(|r| r.as_str()) else { continue };
            if !is_safe_rel(Path::new(rfile)) {
                continue;
            }
            out.push((rfile.to_string(), format!("{base}/datasets/{id}/resolve/main/{rfile}")));
        }
    }
    Ok(out)
}

fn is_data_file(name: &str) -> bool {
    matches!(
        Path::new(name).extension().and_then(|e| e.to_str()).map(|s| s.to_lowercase()).as_deref(),
        Some("jsonl") | Some("parquet")
    )
}

/// Download data-file (filename, url) pairs into `raw`, honoring `max_rows`: jsonl files stream and
/// stop early, parquet files download whole then contribute their footer row-count; the loop stops
/// pulling further files once the cap is reached. Returns (files, bytes).
async fn download_data_capped(
    client: &reqwest::Client,
    pairs: &[(String, String)],
    raw: &Path,
    max_bytes: u64,
    max_rows: Option<usize>,
) -> std::io::Result<(usize, u64)> {
    let mut pairs = pairs.to_vec();
    pairs.sort();
    let mut files = 0usize;
    let mut bytes = 0u64;
    let mut rows = 0usize;
    for (fname, url) in &pairs {
        if let Some(cap) = max_rows {
            if rows >= cap {
                break;
            }
        }
        let dest = raw.join(fname);
        let is_jsonl = matches!(
            Path::new(fname).extension().and_then(|e| e.to_str()).map(|s| s.to_lowercase()).as_deref(),
            Some("jsonl")
        );
        match (is_jsonl, max_rows) {
            (true, Some(cap)) => {
                let n = download_jsonl_sampled(client, url, &dest, cap - rows, max_bytes).await?;
                rows += n;
                files += 1;
                if let Ok(m) = dest.metadata() {
                    bytes += m.len();
                }
            }
            _ => {
                let b = download_to_file(client, url, &dest, max_bytes).await?;
                bytes += b;
                files += 1;
                if max_rows.is_some() {
                    match parquet_num_rows(&dest) {
                        Ok(nr) => rows += nr.max(0) as usize,
                        Err(e) => eprintln!(
                            "kibble: could not read parquet row count for {}: {e} (may pull more shards than the sample needs)",
                            dest.display()
                        ),
                    }
                }
            }
        }
    }
    Ok((files, bytes))
}

/// Acquire an HF dataset: download main-branch siblings (non-data normally, data capped); if no data
/// files landed, fall back to the parquet-convert branch; then normalize. Returns the fetch result,
/// the normalize report, and the source branch ("main" | "refs/convert/parquet").
pub async fn acquire_dataset(
    client: &reqwest::Client,
    base: &str,
    id: &str,
    slug_dir: &Path,
    max_bytes: u64,
    max_rows: Option<usize>,
    override_map: Option<&crate::normalize::ColumnMap>,
) -> std::io::Result<(FetchResult, crate::normalize::NormalizeReport, String)> {
    let raw = slug_dir.join("raw");
    std::fs::create_dir_all(&raw)?;

    let siblings = hf_dataset_siblings(client, base, id).await?;
    let (data_pairs, other_pairs): (Vec<_>, Vec<_>) =
        siblings.into_iter().partition(|(f, _)| is_data_file(f));

    let mut files = 0usize;
    let mut bytes = 0u64;
    // non-data files (README, .txt, code) download whole — these carry raw-text/code datasets
    for (fname, url) in &other_pairs {
        let dest = raw.join(fname);
        match download_to_file(client, url, &dest, max_bytes).await {
            Ok(b) => { files += 1; bytes += b; }
            Err(e) => eprintln!("kibble: skipping {url}: {e}"),
        }
    }
    let (df, db) = download_data_capped(client, &data_pairs, &raw, max_bytes, max_rows).await?;
    files += df;
    bytes += db;

    let mut source_branch = "main".to_string();
    if crate::normalize::collect_data_files(&raw).is_empty() {
        let urls = hf_parquet_urls(client, base, id).await?;
        if !urls.is_empty() {
            source_branch = "refs/convert/parquet".to_string();
            let pairs: Vec<(String, String)> =
                urls.into_iter().map(|u| (filename_from_url(&u), u)).collect();
            let (df2, db2) = download_data_capped(client, &pairs, &raw, max_bytes, max_rows).await?;
            files += df2;
            bytes += db2;
        }
    }

    let out = slug_dir.join("normalized").join("data.jsonl");
    let report = crate::normalize::normalize_dataset(&raw, &out, override_map, max_rows)?;
    Ok((FetchResult { handler: "hf-dataset".to_string(), files, bytes }, report, source_branch))
}

/// Stream a jsonl file, writing only the first `max_rows` complete lines to `dest`, then stop the
/// HTTP transfer (drops the response). Returns the number of lines written. A final line with no
/// trailing newline is captured if the stream ends before `max_rows` is reached. `max_bytes` caps
/// total bytes read.
pub async fn download_jsonl_sampled(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_rows: usize,
    max_bytes: u64,
) -> std::io::Result<usize> {
    use std::io::Write as _;
    if let Some(parent) = dest.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let mut file = std::io::BufWriter::new(std::fs::File::create(dest)?);
    if max_rows == 0 {
        file.flush()?;
        return Ok(0);
    }
    let mut resp = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("download failed for {url}: {e}")))?;
    let mut lines = 0usize;
    let mut total: u64 = 0;
    let mut buf: Vec<u8> = Vec::new();
    'outer: loop {
        let chunk = resp
            .chunk()
            .await
            .map_err(|e| std::io::Error::other(format!("download read failed: {e}")))?;
        let Some(chunk) = chunk else { break };
        total += chunk.len() as u64;
        if total > max_bytes {
            return Err(std::io::Error::other(format!("{url}: download exceeded max_bytes {max_bytes}")));
        }
        buf.extend_from_slice(&chunk);
        while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
            let line: Vec<u8> = buf.drain(..=pos).collect(); // includes the '\n'
            file.write_all(&line)?;
            lines += 1;
            if lines >= max_rows {
                break 'outer;
            }
        }
    }
    // trailing line with no newline (stream ended first) — only real content, not stray whitespace
    if lines < max_rows && buf.iter().any(|&b| !b.is_ascii_whitespace()) {
        file.write_all(&buf)?;
        file.write_all(b"\n")?;
        lines += 1;
    }
    file.flush()?;
    Ok(lines)
}

/// Row count from a parquet file's footer metadata (no data scan).
pub fn parquet_num_rows(path: &Path) -> std::io::Result<i64> {
    use parquet::file::reader::{FileReader, SerializedFileReader};
    let file = std::fs::File::open(path)?;
    let reader = SerializedFileReader::new(file)
        .map_err(|e| std::io::Error::other(format!("parquet meta {}: {e}", path.display())))?;
    Ok(reader.metadata().file_metadata().num_rows())
}

/// The parquet-convert-branch data URLs for a dataset. GET `{base}/api/datasets/{id}/parquet`
/// (shape `{config: {split: [url,…]}}`), return the first config that has a `train` split → its
/// train URLs; else the first config's first non-empty split's URLs; empty vec if none / 404.
pub async fn hf_parquet_urls(
    client: &reqwest::Client,
    base: &str,
    id: &str,
) -> std::io::Result<Vec<String>> {
    let api = format!("{base}/api/datasets/{id}/parquet");
    let text = match client.get(&api).send().await.and_then(|r| r.error_for_status()) {
        Ok(r) => r.text().await.map_err(|e| std::io::Error::other(format!("parquet api read: {e}")))?,
        Err(_) => return Ok(vec![]),
    };
    let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) else { return Ok(vec![]) };
    let Some(configs) = v.as_object() else { return Ok(vec![]) };
    let collect = |arr: &serde_json::Value| -> Vec<String> {
        arr.as_array().map(|a| a.iter().filter_map(|u| u.as_str().map(String::from)).collect()).unwrap_or_default()
    };
    for splits in configs.values() {
        if let Some(train) = splits.get("train") {
            let urls = collect(train);
            if !urls.is_empty() {
                return Ok(urls);
            }
        }
    }
    for splits in configs.values() {
        if let Some(obj) = splits.as_object() {
            for arr in obj.values() {
                let urls = collect(arr);
                if !urls.is_empty() {
                    return Ok(urls);
                }
            }
        }
    }
    Ok(vec![])
}

pub async fn resolve(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
    proxy: Option<&str>,
) -> std::io::Result<FetchResult> {
    if looks_like_git(url) {
        clone_repo(url, dest, proxy)?;
        let files = count_files(dest);
        return Ok(FetchResult { handler: "git".to_string(), files, bytes: 0 });
    }
    if looks_like_hf_dataset(url) {
        return fetch_hf_dataset(client, url, dest, max_bytes).await;
    }
    if looks_like_dropbox(url) {
        return fetch_dropbox(client, url, dest, max_bytes).await;
    }
    if looks_like_gdrive(url) {
        return fetch_gdrive(client, url, dest, max_bytes).await;
    }
    if looks_like_mediafire(url) {
        return fetch_mediafire(client, url, dest, max_bytes).await;
    }
    if looks_like_gofile(url) {
        return fetch_gofile(client, url, dest, max_bytes).await;
    }
    if crate::mega::looks_like_mega(url) {
        return crate::mega::fetch_mega(client, url, dest, max_bytes).await;
    }
    // an open-directory URL conventionally ends with '/'
    if url.ends_with('/') {
        return fetch_http_dir(client, url, dest, max_bytes).await;
    }
    fetch_http_file(client, url, dest, max_bytes).await
}

fn looks_like_mediafire(url: &str) -> bool {
    url.to_lowercase().contains("mediafire.com")
}

pub async fn fetch_mediafire(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    let html = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("mediafire page failed for {url}: {e}")))?
        .text()
        .await
        .map_err(|e| std::io::Error::other(format!("mediafire read failed: {e}")))?;

    // Collect the download href as an owned String, dropping the DOM before any await (Send).
    let href: Option<String> = {
        let doc = Html::parse_document(&html);
        let sel = Selector::parse("a#downloadButton").unwrap();
        doc.select(&sel).next().and_then(|a| a.value().attr("href").map(|h| h.to_string()))
    };
    let href = href.ok_or_else(|| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("no download button on {url}"))
    })?;

    std::fs::create_dir_all(dest)?;
    let target = dest.join(filename_from_url(&href));
    let bytes = download_to_file(client, &href, &target, max_bytes).await?;
    Ok(FetchResult { handler: "mediafire".to_string(), files: 1, bytes })
}

fn looks_like_gofile(url: &str) -> bool {
    url.to_lowercase().contains("gofile.io")
}

fn gofile_content_id(url: &str) -> Option<String> {
    let id = url.split_once("/d/").map(|(_, r)| r.split(['/', '?', '#']).next().unwrap_or(""))?;
    if id.is_empty() { None } else { Some(id.to_string()) }
}

pub async fn fetch_gofile(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    max_bytes: u64,
) -> std::io::Result<FetchResult> {
    let id = gofile_content_id(url)
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidInput, format!("no gofile content id in {url}")))?;
    let base = scheme_host(url);

    // guest token
    let acct = client.post(format!("{base}/accounts"))
        .send().await.and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("gofile accounts failed: {e}")))?
        .text().await.map_err(|e| std::io::Error::other(format!("gofile accounts read: {e}")))?;
    let acct: serde_json::Value = serde_json::from_str(&acct)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
    let token = acct.get("data").and_then(|d| d.get("token")).and_then(|t| t.as_str()).unwrap_or("");

    let contents = client.get(format!("{base}/contents/{id}"))
        .header("authorization", format!("Bearer {token}"))
        .send().await.and_then(|r| r.error_for_status())
        .map_err(|e| std::io::Error::other(format!("gofile contents failed: {e}")))?
        .text().await.map_err(|e| std::io::Error::other(format!("gofile contents read: {e}")))?;
    let contents: serde_json::Value = serde_json::from_str(&contents)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

    // collect (name, link) for file children into owned values
    let mut items: Vec<(String, String)> = Vec::new();
    if let Some(children) = contents.get("data").and_then(|d| d.get("children")).and_then(|c| c.as_object()) {
        for child in children.values() {
            if child.get("type").and_then(|t| t.as_str()) == Some("file") {
                if let Some(link) = child.get("link").and_then(|l| l.as_str()) {
                    let name = child.get("name").and_then(|n| n.as_str())
                        .map(|s| s.to_string()).unwrap_or_else(|| filename_from_url(link));
                    items.push((name, link.to_string()));
                }
            }
        }
    }

    std::fs::create_dir_all(dest)?;
    let mut files = 0; let mut bytes = 0u64;
    for (name, link) in items {
        let safe = if name.contains('/') || name.contains('\\') || name == "." || name == ".." {
            filename_from_url(&link)
        } else { name };
        let target = dest.join(&safe);
        match download_to_file(client, &link, &target, max_bytes).await {
            Ok(n) => { files += 1; bytes += n; }
            Err(e) => eprintln!("kibble: skipping {link}: {e}"),
        }
    }
    Ok(FetchResult { handler: "gofile".to_string(), files, bytes })
}

fn count_files(dir: &Path) -> usize {
    let mut n = 0;
    if let Ok(rd) = std::fs::read_dir(dir) {
        for e in rd.flatten() {
            let p = e.path();
            if p.is_dir() {
                if p.file_name().and_then(|x| x.to_str()) != Some(".git") {
                    n += count_files(&p);
                }
            } else {
                n += 1;
            }
        }
    }
    n
}

/// Stage a URL into `<ingest_root>/<slug>/` using the configured proxy + cap. Shared by the
/// `kibble fetch` CLI and the serve job task.
pub async fn stage_url(
    repo_root: &Path,
    cfg: &crate::config::KibbleConfig,
    client: &reqwest::Client,
    url: &str,
    name: Option<&str>,
) -> std::io::Result<FetchResult> {
    let proxy = resolve_proxy(cfg.network.proxy.as_deref(), |k| std::env::var(k).ok());
    let slug = name.map(|s| s.to_string()).unwrap_or_else(|| cache_slug(url));
    let dest = repo_root.join(&cfg.fetch.ingest_root).join(&slug);
    resolve(client, url, &dest, cfg.fetch.max_bytes, proxy.as_deref()).await
}

#[derive(serde::Serialize)]
struct Manifest {
    source: String,
    id: String,
    source_branch: String,
    sampled: bool,
    max_rows: Option<usize>,
    fetched_at: u64,
    detected_schema: String,
    columns: Vec<String>,
    mapping: Option<String>,
    files: Vec<String>,
    rows_in: usize,
    rows_out: usize,
    rows_dropped: usize,
    drop_reasons: std::collections::BTreeMap<String, usize>,
}

fn epoch_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

fn drop_reason_label(r: &crate::normalize::DropReason) -> String {
    use crate::normalize::DropReason::*;
    match r {
        Empty => "empty".to_string(),
        MissingField(f) => format!("missing:{f}"),
        TooFewTurns => "too-few-turns".to_string(),
        NotChat => "not-chat".to_string(),
    }
}

fn source_stanza(slug: &str) -> String {
    format!(
        "    [[source]]\n    path = \"data/datasets/{slug}/normalized\"\n    name = \"{slug}\"\n    type = \"dataset\"",
    )
}

fn files_source_stanza(slug: &str) -> String {
    format!(
        "    [[source]]\n    path = \"data/datasets/{slug}/raw\"\n    name = \"{slug}\"\n    type = \"files\"",
    )
}

/// The HF slug: an explicit --name, else the repo name (last '/'-segment of the id).
fn hf_slug(id: &str, name: Option<&str>) -> String {
    name.map(|s| s.to_string())
        .unwrap_or_else(|| id.rsplit('/').next().unwrap_or(id).to_string())
}

pub async fn run_fetch(
    repo_root: &Path,
    url: &str,
    name: Option<&str>,
    map: Option<&str>,
    max_rows: Option<usize>,
) -> std::io::Result<FetchResult> {
    let cfg = load_config(&repo_root.join(crate::config::CONFIG_FILE));
    let proxy = resolve_proxy(cfg.network.proxy.as_deref(), |k| std::env::var(k).ok());
    let client = build_client(proxy.as_deref())
        .map_err(|e| std::io::Error::other(format!("http client: {e}")))?;

    // Parse the optional column-map override up front so a bad spec fails fast.
    let column_map = match map {
        Some(spec) => Some(
            crate::normalize::parse_column_map(spec)
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))?,
        ),
        None => None,
    };

    // HF dataset/model targets route to datasets_root / models_root.
    if let Some(hf) = parse_hf_target(url) {
        let kind = resolve_hf_kind(&client, &hf.base, &hf.id, hf.hint).await?;
        let slug = hf_slug(&hf.id, name);
        match kind {
            HfKind::Model => {
                let dest = repo_root.join(&cfg.fetch.models_root).join(&slug);
                let r = fetch_hf_siblings(&client, HfKind::Model, &hf.base, &hf.id, &dest, cfg.fetch.max_bytes).await?;
                println!(
                    "Fetched [{}]: {} → {} ({} files). Models are not normalized.",
                    r.handler, hf.id, dest.display(), r.files
                );
                return Ok(r);
            }
            HfKind::Dataset => {
                let slug_dir = repo_root.join(&cfg.fetch.datasets_root).join(&slug);
                let (r, report, source_branch) = acquire_dataset(
                    &client, &hf.base, &hf.id, &slug_dir, cfg.fetch.max_bytes, max_rows, column_map.as_ref(),
                ).await?;
                let out = slug_dir.join("normalized").join("data.jsonl");

                let rows_dropped: usize = report.dropped.iter().map(|(_, n)| n).sum();
                let drop_reasons: std::collections::BTreeMap<String, usize> = report
                    .dropped
                    .iter()
                    .map(|(reason, n)| (drop_reason_label(reason), *n))
                    .collect();
                let manifest = Manifest {
                    source: url.to_string(),
                    id: hf.id.clone(),
                    source_branch: source_branch.clone(),
                    sampled: max_rows.is_some(),
                    max_rows,
                    fetched_at: epoch_secs(),
                    detected_schema: report.schema.label(),
                    columns: match &report.schema {
                        crate::normalize::Schema::Unknown(cols) => cols.clone(),
                        _ => vec![],
                    },
                    mapping: map.map(|s| s.to_string()),
                    files: report.files.clone(),
                    rows_in: report.rows_in,
                    rows_out: report.rows_out,
                    rows_dropped,
                    drop_reasons,
                };
                std::fs::write(slug_dir.join("manifest.json"), serde_json::to_string_pretty(&manifest)?)?;

                let branch_note = if source_branch == "main" { String::new() } else { format!("  [via {source_branch}]\n") };
                if report.wrote_output {
                    let sampled_note = max_rows.map(|n| format!(" (sampled ≤{n})")).unwrap_or_default();
                    println!(
                        "Fetched [{}]: {}\n{}  schema: {} ({})   rows: {} in → {} out{}, {} dropped\n  normalized → {}\n  add to kibble.toml:\n{}",
                        r.handler, hf.id, branch_note,
                        report.schema.label(),
                        if map.is_some() { "override" } else { "auto" },
                        report.rows_in, report.rows_out, sampled_note, rows_dropped,
                        out.display(), source_stanza(&slug),
                    );
                } else if matches!(report.schema, crate::normalize::Schema::RawFiles) {
                    println!(
                        "Fetched [{}]: {} (raw → {})\n{}  no jsonl/parquet data files — this looks like raw text/code.\n  add it as a `files` source:\n{}",
                        r.handler, hf.id, slug_dir.join("raw").display(), branch_note, files_source_stanza(&slug),
                    );
                } else {
                    let cols = match &report.schema {
                        crate::normalize::Schema::Unknown(c) => format!("columns seen: {c:?}"),
                        _ => "looks like raw text/code, not chat".to_string(),
                    };
                    println!(
                        "Fetched [{}]: {} (raw → {})\n{}  ⚠ could not normalize: {} ({})\n  → re-run with --map user=<col>,assistant=<col>, or add as a `files` source if it's raw text/code.\n  No normalized output written.",
                        r.handler, hf.id, slug_dir.join("raw").display(), branch_note, report.schema.label(), cols,
                    );
                }
                return Ok(r);
            }
        }
    }

    // Non-HF: existing ingest-root staging behavior.
    let result = stage_url(repo_root, &cfg, &client, url, name).await?;
    let slug = name.map(|s| s.to_string()).unwrap_or_else(|| cache_slug(url));
    let dest = repo_root.join(&cfg.fetch.ingest_root).join(&slug);
    println!(
        "Fetched [{}]: {} files, {} bytes -> {}",
        result.handler, result.files, result.bytes, dest.display()
    );
    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpListener;

    fn serve_once(body: Vec<u8>, with_len: bool) -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            if let Ok((mut s, _)) = listener.accept() {
                let mut b = [0u8; 1024];
                let _ = s.read(&mut b);
                let mut head = String::from("HTTP/1.1 200 OK\r\nConnection: close\r\n");
                if with_len {
                    head.push_str(&format!("Content-Length: {}\r\n", body.len()));
                }
                head.push_str("\r\n");
                let _ = s.write_all(head.as_bytes());
                let _ = s.write_all(&body);
                let _ = s.flush();
            }
        });
        port
    }

    #[tokio::test]
    async fn downloads_small_file() {
        let port = serve_once(b"hello world".to_vec(), true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_dl_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let dest = dir.join("f.txt");
        let n = download_to_file(&client, &format!("http://127.0.0.1:{port}/"), &dest, 1000).await.unwrap();
        assert_eq!(n, 11);
        assert_eq!(std::fs::read_to_string(&dest).unwrap(), "hello world");
    }

    #[test]
    fn extracts_zip_and_rejects_zip_slip() {
        use std::io::Write as _;
        let dir = std::env::temp_dir().join(format!("kibble_zip_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        // a safe zip
        let zpath = dir.join("ok.zip");
        {
            let f = std::fs::File::create(&zpath).unwrap();
            let mut z = zip::ZipWriter::new(f);
            let o = zip::write::SimpleFileOptions::default();
            z.start_file("a/b.txt", o).unwrap();
            z.write_all(b"hi").unwrap();
            z.finish().unwrap();
        }
        let out = dir.join("out");
        let n = extract_archive(&zpath, &out).unwrap();
        assert_eq!(n, 1);
        assert_eq!(std::fs::read_to_string(out.join("a/b.txt")).unwrap(), "hi");

        // a zip-slip zip
        let evil = dir.join("evil.zip");
        {
            let f = std::fs::File::create(&evil).unwrap();
            let mut z = zip::ZipWriter::new(f);
            let o = zip::write::SimpleFileOptions::default();
            z.start_file("../escape.txt", o).unwrap();
            z.write_all(b"x").unwrap();
            z.finish().unwrap();
        }
        assert!(extract_archive(&evil, &dir.join("out2")).is_err());
    }

    #[test]
    fn extracts_tar_gz() {
        let dir = std::env::temp_dir().join(format!("kibble_tgz_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let tgz = dir.join("a.tar.gz");
        {
            let f = std::fs::File::create(&tgz).unwrap();
            let enc = flate2::write::GzEncoder::new(f, flate2::Compression::default());
            let mut tar = tar::Builder::new(enc);
            let data = b"tarred";
            let mut header = tar::Header::new_gnu();
            header.set_size(data.len() as u64);
            header.set_cksum();
            tar.append_data(&mut header, "dir/file.txt", &data[..]).unwrap();
            tar.into_inner().unwrap().finish().unwrap();
        }
        let out = dir.join("out");
        let n = extract_archive(&tgz, &out).unwrap();
        assert_eq!(n, 1);
        assert_eq!(std::fs::read_to_string(out.join("dir/file.txt")).unwrap(), "tarred");
    }

    #[test]
    fn rejects_tar_symlink_entry() {
        let dir = std::env::temp_dir().join(format!("kibble_tarsym_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        // Build a .tar.gz with a symlink entry pointing outside dest
        let tgz = dir.join("evil.tar.gz");
        {
            let f = std::fs::File::create(&tgz).unwrap();
            let enc = flate2::write::GzEncoder::new(f, flate2::Compression::default());
            let mut builder = tar::Builder::new(enc);
            let mut header = tar::Header::new_gnu();
            header.set_entry_type(tar::EntryType::Symlink);
            header.set_size(0);
            header.set_link_name("../escape").unwrap();
            header.set_cksum();
            builder.append_data(&mut header, "evil", std::io::empty()).unwrap();
            builder.into_inner().unwrap().finish().unwrap();
        }

        let out = dir.join("out");
        let res = extract_archive(&tgz, &out);
        assert!(res.is_err(), "symlink tar entry must be rejected");
        let msg = res.unwrap_err().to_string();
        assert!(msg.contains("unsafe tar link entry"), "unexpected error: {msg}");
    }

    #[tokio::test]
    async fn rejects_oversized_via_stream() {
        let port = serve_once(vec![b'x'; 100], false);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_dl_stream_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let dest = dir.join("f.bin");
        let res = download_to_file(&client, &format!("http://127.0.0.1:{port}/"), &dest, 10).await;
        assert!(res.is_err(), "oversized streaming download must return Err");
        assert!(!dest.exists(), "partial file must be removed");
    }

    #[tokio::test]
    async fn http_file_downloads_and_names_from_url() {
        let port = serve_once(b"filedata".to_vec(), true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_hf_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let r = fetch_http_file(&client, &format!("http://127.0.0.1:{port}/data/notes.txt"), &dir, 1000).await.unwrap();
        assert_eq!(r.files, 1);
        assert_eq!(std::fs::read_to_string(dir.join("notes.txt")).unwrap(), "filedata");
    }

    #[tokio::test]
    async fn http_dir_crawls_links() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        // server answers the index with two <a href> links, then serves each file
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..3 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("GET / ") || req.contains("GET /index") {
                        b"<html><body><a href=\"a.txt\">a</a><a href=\"../up\">up</a><a href=\"b.txt\">b</a></body></html>".to_vec()
                    } else if req.contains("a.txt") {
                        b"AAA".to_vec()
                    } else {
                        b"BBB".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_hd_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let r = fetch_http_dir(&client, &format!("http://127.0.0.1:{port}/"), &dir, 1000).await.unwrap();
        assert_eq!(r.files, 2); // a.txt + b.txt; ../up skipped
        assert_eq!(std::fs::read_to_string(dir.join("a.txt")).unwrap(), "AAA");
    }

    #[tokio::test]
    async fn hf_dataset_lists_and_downloads() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..3 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/datasets/myds") {
                        br#"{"siblings":[{"rfilename":"train.jsonl"},{"rfilename":".gitattributes"}]}"#.to_vec()
                    } else if req.contains("/resolve/main/train.jsonl") {
                        b"{\"x\":1}\n".to_vec()
                    } else {
                        b"attr".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_hfds_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let r = fetch_hf_dataset(&client, &format!("http://127.0.0.1:{port}/datasets/myds"), &dir, 10000).await.unwrap();
        assert_eq!(r.handler, "huggingface");
        assert!(r.files >= 1);
        assert_eq!(std::fs::read_to_string(dir.join("train.jsonl")).unwrap(), "{\"x\":1}\n");
    }

    #[test]
    fn filename_from_url_rejects_traversal() {
        assert_eq!(filename_from_url("http://h/a/.."), "download");
        assert_eq!(filename_from_url("http://h/a/."), "download");
        assert_eq!(filename_from_url("http://h/dir/"), "download");
        assert_eq!(filename_from_url("http://h/notes.txt?v=1"), "notes.txt");
        assert_eq!(filename_from_url("http://h/data/x.jsonl"), "x.jsonl");
    }

    #[tokio::test]
    async fn rejects_oversized_via_content_length() {
        let port = serve_once(vec![b'x'; 100], true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_dl_big_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let dest = dir.join("f.bin");
        let res = download_to_file(&client, &format!("http://127.0.0.1:{port}/"), &dest, 10).await;
        assert!(res.is_err());
        assert!(!dest.exists(), "partial file must be removed");
    }

    #[test]
    fn dropbox_url_forces_direct_download() {
        assert_eq!(dropbox_direct_url("https://www.dropbox.com/s/abc/f.txt?dl=0"),
                   "https://www.dropbox.com/s/abc/f.txt?dl=1");
        assert_eq!(dropbox_direct_url("https://www.dropbox.com/s/abc/f.txt"),
                   "https://www.dropbox.com/s/abc/f.txt?dl=1");
        assert_eq!(dropbox_direct_url("https://www.dropbox.com/scl/fi/x/f.txt?rlkey=k"),
                   "https://www.dropbox.com/scl/fi/x/f.txt?rlkey=k&dl=1");
    }

    #[tokio::test]
    async fn dropbox_downloads_direct_file() {
        let port = serve_once(b"dropbox body".to_vec(), true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_db_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        // the mock ignores the path; a dropbox-shaped URL on localhost
        let url = format!("http://127.0.0.1:{port}/s/abc/note.txt?dl=0");
        let r = fetch_dropbox(&client, &url, &dir, 10000).await.unwrap();
        assert_eq!(r.handler, "dropbox");
        assert_eq!(r.files, 1);
        assert_eq!(std::fs::read_to_string(dir.join("note.txt")).unwrap(), "dropbox body");
    }

    #[test]
    fn gdrive_extracts_file_id() {
        assert_eq!(gdrive_extract_id("https://drive.google.com/file/d/ABC123/view?usp=sharing").as_deref(), Some("ABC123"));
        assert_eq!(gdrive_extract_id("https://drive.google.com/uc?export=download&id=XYZ789").as_deref(), Some("XYZ789"));
        assert_eq!(gdrive_extract_id("https://drive.google.com/").as_deref(), None);
    }

    #[tokio::test]
    async fn gdrive_downloads_by_id() {
        // mock the uc endpoint returning the file body directly
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            if let Ok((mut s, _)) = listener.accept() {
                let mut b = [0u8; 1024]; let _ = s.read(&mut b);
                let body = b"drive file body";
                let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                let _ = s.write_all(head.as_bytes()); let _ = s.write_all(body);
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_gd_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let url = format!("http://127.0.0.1:{port}/file/d/FILEID/view");
        let r = fetch_gdrive(&client, &url, &dir, 10000).await.unwrap();
        assert_eq!(r.handler, "gdrive");
        assert_eq!(r.files, 1);
        assert_eq!(std::fs::read_to_string(dir.join("FILEID")).unwrap(), "drive file body");
    }

    #[tokio::test]
    async fn mediafire_scrapes_download_button() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..2 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024]; let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/file/") {
                        format!("<html><body><a id=\"downloadButton\" href=\"http://127.0.0.1:{port}/dl/real.bin\">Download</a></body></html>").into_bytes()
                    } else {
                        b"the real bytes".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes()); let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_mf_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let url = format!("http://127.0.0.1:{port}/file/abc/real.bin/file");
        let r = fetch_mediafire(&client, &url, &dir, 10000).await.unwrap();
        assert_eq!(r.handler, "mediafire");
        assert_eq!(std::fs::read_to_string(dir.join("real.bin")).unwrap(), "the real bytes");
    }

    #[tokio::test]
    async fn mediafire_errors_without_button() {
        let port = serve_once(b"<html><body>no button here</body></html>".to_vec(), true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_mf2_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let url = format!("http://127.0.0.1:{port}/file/x/y/file");
        assert!(fetch_mediafire(&client, &url, &dir, 10000).await.is_err());
    }

    #[tokio::test]
    async fn gofile_token_then_download() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..3 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024]; let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/accounts") {
                        br#"{"data":{"token":"tok123"}}"#.to_vec()
                    } else if req.contains("/contents/") {
                        format!(r#"{{"data":{{"children":{{"k":{{"type":"file","name":"g.bin","link":"http://127.0.0.1:{port}/store/g.bin"}}}}}}}}"#).into_bytes()
                    } else {
                        b"gofile bytes".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes()); let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_gf_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let url = format!("http://127.0.0.1:{port}/d/CONTENTID");
        let r = fetch_gofile(&client, &url, &dir, 10000).await.unwrap();
        assert_eq!(r.handler, "gofile");
        assert!(r.files >= 1);
        assert_eq!(std::fs::read_to_string(dir.join("g.bin")).unwrap(), "gofile bytes");
    }

    #[tokio::test]
    async fn run_fetch_stages_http_file_into_ingest_root() {
        let port = serve_once(b"staged body".to_vec(), true);
        let root = std::env::temp_dir().join(format!("kibble_runfetch_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        // default [fetch] ⇒ ingest_root data/ingest
        std::fs::write(root.join(crate::config::CONFIG_FILE), "").unwrap();
        let r = run_fetch(&root, &format!("http://127.0.0.1:{port}/x/note.txt"), Some("mydl"), None, None).await.unwrap();
        assert_eq!(r.files, 1);
        assert_eq!(std::fs::read_to_string(root.join("data/ingest/mydl/note.txt")).unwrap(), "staged body");
    }

    #[test]
    fn parse_hf_target_shapes() {
        // full dataset URL
        let r = parse_hf_target("https://huggingface.co/datasets/user/name").unwrap();
        assert_eq!(r.base, "https://huggingface.co");
        assert_eq!(r.id, "user/name");
        assert_eq!(r.hint, Some(HfKind::Dataset));
        // full model URL (no /datasets/)
        let r = parse_hf_target("https://huggingface.co/user/model").unwrap();
        assert_eq!(r.hint, Some(HfKind::Model));
        assert_eq!(r.id, "user/model");
        // real HF model repo names contain dots (version numbers) — must still resolve as Model
        assert_eq!(parse_hf_target("https://huggingface.co/Qwen/Qwen2.5-7B").unwrap().hint, Some(HfKind::Model));
        assert_eq!(parse_hf_target("https://huggingface.co/mistralai/Mistral-7B-v0.1").unwrap().hint, Some(HfKind::Model));
        // bare user/name shorthand -> defaults to hf.co, kind unknown (needs probe)
        let r = parse_hf_target("user/name").unwrap();
        assert_eq!(r.base, "https://huggingface.co");
        assert_eq!(r.id, "user/name");
        assert_eq!(r.hint, None);
        // localhost dataset URL (for tests) keeps its base
        let r = parse_hf_target("http://127.0.0.1:9999/datasets/name").unwrap();
        assert_eq!(r.base, "http://127.0.0.1:9999");
        assert_eq!(r.id, "name");
        assert_eq!(r.hint, Some(HfKind::Dataset));
        // not HF-shaped
        assert!(parse_hf_target("https://example.com/foo.zip").is_none());
        assert!(parse_hf_target("./local/path").is_none());
        assert!(parse_hf_target("a/b/c").is_none()); // too many segments for bare shorthand
        // /datasets/ with no id → None (not a model named "datasets")
        assert!(parse_hf_target("https://huggingface.co/datasets/").is_none());
        assert!(parse_hf_target("https://huggingface.co/datasets").is_none());
    }

    #[tokio::test]
    async fn resolve_hf_kind_probes_dataset_then_model() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..4 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    // /api/datasets/known -> 200; everything else -> 404
                    let (code, body): (&str, &str) = if req.contains("/api/datasets/known") {
                        ("200 OK", "{}")
                    } else {
                        ("404 Not Found", "no")
                    };
                    let head = format!("HTTP/1.1 {code}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(body.as_bytes());
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let base = format!("http://127.0.0.1:{port}");
        let hint = None;
        assert_eq!(resolve_hf_kind(&client, &base, "known", hint).await.unwrap(), HfKind::Dataset);
        // hint short-circuits without any probe
        assert_eq!(resolve_hf_kind(&client, &base, "whatever", Some(HfKind::Model)).await.unwrap(), HfKind::Model);
    }

    #[test]
    fn source_stanza_is_pasteable_toml() {
        let s = source_stanza("claude_mythos_distilled_25k");
        assert!(s.contains("[[source]]"));
        assert!(s.contains("data/datasets/claude_mythos_distilled_25k/normalized"));
        assert!(s.contains("type = \"dataset\""));
    }

    #[tokio::test]
    async fn run_fetch_dataset_normalizes_and_writes_manifest() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..4 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/datasets/myset") {
                        br#"{"siblings":[{"rfilename":"train.jsonl"}]}"#.to_vec()
                    } else {
                        b"{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(&body);
                }
            }
        });
        let root = std::env::temp_dir().join(format!("kibble_rf_ds_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join(crate::config::CONFIG_FILE), "").unwrap();
        let url = format!("http://127.0.0.1:{port}/datasets/myset");
        let r = run_fetch(&root, &url, None, None, None).await.unwrap();
        assert_eq!(r.handler, "hf-dataset");
        // normalized output the dataset source can read
        let norm = root.join("data/datasets/myset/normalized/data.jsonl");
        assert!(norm.exists(), "normalized data.jsonl missing");
        assert_eq!(std::fs::read_to_string(&norm).unwrap().lines().count(), 2);
        // manifest with provenance
        let manifest = root.join("data/datasets/myset/manifest.json");
        assert!(manifest.exists(), "manifest.json missing");
        let mj: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&manifest).unwrap()).unwrap();
        assert_eq!(mj["detected_schema"], "prompt-response");
        assert_eq!(mj["rows_in"], 2);
        assert_eq!(mj["rows_out"], 2);
        // raw kept under raw/
        assert!(root.join("data/datasets/myset/raw/train.jsonl").exists());
    }

    #[tokio::test]
    async fn run_fetch_model_lands_in_models_root_no_normalize() {
        use std::io::{Read, Write};
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..3 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/models/user/mymodel") {
                        br#"{"siblings":[{"rfilename":"config.json"}]}"#.to_vec()
                    } else {
                        b"{}".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(&body);
                }
            }
        });
        let root = std::env::temp_dir().join(format!("kibble_rf_m_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join(crate::config::CONFIG_FILE), "").unwrap();
        let url = format!("http://127.0.0.1:{port}/user/mymodel");
        let r = run_fetch(&root, &url, None, None, None).await.unwrap();
        assert_eq!(r.handler, "hf-model");
        assert!(root.join("data/models/mymodel/config.json").exists());
        // no dataset normalization/manifest for a model
        assert!(!root.join("data/datasets/mymodel").exists());
    }

    #[tokio::test]
    async fn run_fetch_non_hf_still_lands_in_ingest_root() {
        let port = serve_once(b"staged body".to_vec(), true);
        let root = std::env::temp_dir().join(format!("kibble_rf_ing_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join(crate::config::CONFIG_FILE), "").unwrap();
        let r = run_fetch(&root, &format!("http://127.0.0.1:{port}/x/note.txt"), Some("mydl"), None, None).await.unwrap();
        assert_eq!(r.files, 1);
        assert_eq!(std::fs::read_to_string(root.join("data/ingest/mydl/note.txt")).unwrap(), "staged body");
    }

    // serves `body` on every accepted connection, up to `n` times
    fn serve_n(body: Vec<u8>, n: usize, with_len: bool) -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..n {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024];
                    let _ = s.read(&mut b);
                    let mut head = String::from("HTTP/1.1 200 OK\r\nConnection: close\r\n");
                    if with_len { head.push_str(&format!("Content-Length: {}\r\n", body.len())); }
                    head.push_str("\r\n");
                    let _ = s.write_all(head.as_bytes());
                    let _ = s.write_all(&body);
                    let _ = s.flush();
                }
            }
        });
        port
    }

    #[tokio::test]
    async fn jsonl_sampled_stops_at_max_rows() {
        let body: String = (0..100).map(|i| format!("{{\"n\":{i}}}\n")).collect();
        let port = serve_n(body.into_bytes(), 1, true);
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_js_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let dest = dir.join("s.jsonl");
        let n = download_jsonl_sampled(&client, &format!("http://127.0.0.1:{port}/"), &dest, 10, 1_000_000).await.unwrap();
        assert_eq!(n, 10);
        assert_eq!(std::fs::read_to_string(&dest).unwrap().lines().count(), 10);
    }

    #[tokio::test]
    async fn jsonl_sampled_captures_trailing_line_without_newline() {
        let port = serve_n(b"{\"a\":1}\n{\"a\":2}".to_vec(), 1, true); // 2nd line has no trailing \n
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_js2_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let dest = dir.join("s.jsonl");
        let n = download_jsonl_sampled(&client, &format!("http://127.0.0.1:{port}/"), &dest, 100, 1_000_000).await.unwrap();
        assert_eq!(n, 2);
        assert_eq!(std::fs::read_to_string(&dest).unwrap().lines().count(), 2);
    }

    #[test]
    fn parquet_num_rows_reads_footer() {
        use parquet::data_type::{ByteArray, ByteArrayType};
        use parquet::file::properties::WriterProperties;
        use parquet::file::writer::SerializedFileWriter;
        use parquet::schema::parser::parse_message_type;
        use std::sync::Arc;
        let dir = std::env::temp_dir().join(format!("kibble_pq_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("t.parquet");
        let schema = parse_message_type("message s { REQUIRED BYTE_ARRAY x (UTF8); }").unwrap();
        let props = Arc::new(WriterProperties::builder().build());
        let file = std::fs::File::create(&path).unwrap();
        let mut w = SerializedFileWriter::new(file, Arc::new(schema), props).unwrap();
        let mut rg = w.next_row_group().unwrap();
        let mut c = rg.next_column().unwrap().unwrap();
        c.typed::<ByteArrayType>().write_batch(&[ByteArray::from("a"), ByteArray::from("b"), ByteArray::from("c")], None, None).unwrap();
        c.close().unwrap();
        rg.close().unwrap();
        w.close().unwrap();
        assert_eq!(parquet_num_rows(&path).unwrap(), 3);
    }

    #[tokio::test]
    async fn hf_parquet_urls_prefers_first_config_with_train() {
        // emotion-shaped: two configs, both with a "train" split; BTreeMap key order → "split" first
        let body = br#"{"split":{"test":["http://h/split/test/0.parquet"],"train":["http://h/split/train/0.parquet"]},"unsplit":{"train":["http://h/unsplit/train/0.parquet"]}}"#.to_vec();
        let port = serve_n(body, 1, true);
        let client = crate::net::build_client(None).unwrap();
        let urls = hf_parquet_urls(&client, &format!("http://127.0.0.1:{port}"), "u/d").await.unwrap();
        assert_eq!(urls, vec!["http://h/split/train/0.parquet".to_string()]);
    }

    #[tokio::test]
    async fn hf_parquet_urls_empty_on_404() {
        // server closes with 404
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            if let Ok((mut s, _)) = listener.accept() {
                let mut b = [0u8; 1024]; let _ = s.read(&mut b);
                let _ = s.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let urls = hf_parquet_urls(&client, &format!("http://127.0.0.1:{port}"), "u/d").await.unwrap();
        assert!(urls.is_empty());
    }

    // Multi-endpoint HF mock: /api/datasets/{id} lists siblings; /resolve/main/<f> serves file bodies;
    // /api/datasets/{id}/parquet serves the convert listing; convert file URLs serve parquet bytes.
    #[tokio::test]
    async fn acquire_dataset_main_branch_jsonl() {
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..6 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024]; let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/datasets/u/d/parquet") {
                        b"{}".to_vec()
                    } else if req.contains("/api/datasets/u/d") {
                        br#"{"siblings":[{"rfilename":"train.jsonl"},{"rfilename":"README.md"}]}"#.to_vec()
                    } else if req.contains("/resolve/main/train.jsonl") {
                        b"{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n".to_vec()
                    } else {
                        b"readme".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes()); let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_acq_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let (fr, report, branch) = acquire_dataset(&client, &format!("http://127.0.0.1:{port}"), "u/d", &dir, 1_000_000, None, None).await.unwrap();
        assert_eq!(branch, "main");
        assert_eq!(report.rows_out, 2);
        assert_eq!(report.schema.label(), "prompt-response");
        assert!(fr.files >= 1);
        assert!(dir.join("raw/train.jsonl").exists());
        assert!(dir.join("normalized/data.jsonl").exists());
    }

    #[tokio::test]
    async fn acquire_dataset_falls_back_to_parquet_convert() {
        use std::net::TcpListener;
        // pre-build the parquet bytes the convert URL will serve
        let pqbytes = {
            use parquet::data_type::{ByteArray, ByteArrayType};
            use parquet::file::properties::WriterProperties;
            use parquet::file::writer::SerializedFileWriter;
            use parquet::schema::parser::parse_message_type;
            use std::sync::Arc;
            let tmp = std::env::temp_dir().join(format!("kibble_acqpq_{}.parquet", std::process::id()));
            let schema = parse_message_type("message s { REQUIRED BYTE_ARRAY prompt (UTF8); REQUIRED BYTE_ARRAY response (UTF8); }").unwrap();
            let props = Arc::new(WriterProperties::builder().build());
            let f = std::fs::File::create(&tmp).unwrap();
            let mut w = SerializedFileWriter::new(f, Arc::new(schema), props).unwrap();
            let mut rg = w.next_row_group().unwrap();
            let mut c = rg.next_column().unwrap().unwrap();
            c.typed::<ByteArrayType>().write_batch(&[ByteArray::from("q1"), ByteArray::from("q2")], None, None).unwrap();
            c.close().unwrap();
            let mut c = rg.next_column().unwrap().unwrap();
            c.typed::<ByteArrayType>().write_batch(&[ByteArray::from("a1"), ByteArray::from("a2")], None, None).unwrap();
            c.close().unwrap();
            rg.close().unwrap(); w.close().unwrap();
            std::fs::read(&tmp).unwrap()
        };
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..6 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024]; let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/datasets/u/d/parquet") {
                        format!("{{\"default\":{{\"train\":[\"http://127.0.0.1:{port}/cvt/train/0.parquet\"]}}}}").into_bytes()
                    } else if req.contains("/api/datasets/u/d") {
                        br#"{"siblings":[{"rfilename":"README.md"}]}"#.to_vec()
                    } else if req.contains("/cvt/train/0.parquet") {
                        pqbytes.clone()
                    } else {
                        b"readme".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes()); let _ = s.write_all(&body);
                }
            }
        });
        let client = crate::net::build_client(None).unwrap();
        let dir = std::env::temp_dir().join(format!("kibble_acq2_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let (_fr, report, branch) = acquire_dataset(&client, &format!("http://127.0.0.1:{port}"), "u/d", &dir, 10_000_000, None, None).await.unwrap();
        assert_eq!(branch, "refs/convert/parquet");
        assert_eq!(report.rows_out, 2);
        assert!(dir.join("raw/0.parquet").exists());
    }

    #[test]
    fn files_source_stanza_points_at_raw() {
        let s = files_source_stanza("man_pages");
        assert!(s.contains("[[source]]"));
        assert!(s.contains("data/datasets/man_pages/raw"));
        assert!(s.contains("type = \"files\""));
    }

    #[tokio::test]
    async fn run_fetch_dataset_records_branch_and_sampling_in_manifest() {
        use std::net::TcpListener;
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        std::thread::spawn(move || {
            for _ in 0..6 {
                if let Ok((mut s, _)) = listener.accept() {
                    let mut b = [0u8; 1024]; let n = s.read(&mut b).unwrap_or(0);
                    let req = String::from_utf8_lossy(&b[..n]);
                    let body: Vec<u8> = if req.contains("/api/datasets/u/samp/parquet") {
                        b"{}".to_vec()
                    } else if req.contains("/api/datasets/u/samp") {
                        br#"{"siblings":[{"rfilename":"train.jsonl"}]}"#.to_vec()
                    } else {
                        b"{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n{\"prompt\":\"q3\",\"response\":\"a3\"}\n".to_vec()
                    };
                    let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
                    let _ = s.write_all(head.as_bytes()); let _ = s.write_all(&body);
                }
            }
        });
        let root = std::env::temp_dir().join(format!("kibble_rfsamp_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join(crate::config::CONFIG_FILE), "").unwrap();
        let url = format!("http://127.0.0.1:{port}/datasets/u/samp");
        let r = run_fetch(&root, &url, None, None, Some(2)).await.unwrap();
        assert_eq!(r.handler, "hf-dataset");
        let mj: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(root.join("data/datasets/samp/manifest.json")).unwrap()).unwrap();
        assert_eq!(mj["sampled"], true);
        assert_eq!(mj["max_rows"], 2);
        assert_eq!(mj["source_branch"], "main");
        assert_eq!(mj["rows_out"], 2);
        assert_eq!(std::fs::read_to_string(root.join("data/datasets/samp/normalized/data.jsonl")).unwrap().lines().count(), 2);
    }
}