//! The Rust indexer over stored chunks (REBUILD.md, part 2 and "Voyager 3"): a range captured by `edge:bench
//! --capture` or stored in the corpus, or a group's corpus ranges in order through one window (`--group`, the
//! replay harness behind `replay:corpus`). Same work as the edge-node worker: inflate each chunk, scan it,
//! aggregate files, finish releases through the window, and every N chunks judge and encode them. `--dump`
//! writes every output as JSON lines; `--wire` the batches as the worker posts them (`-`: stdout).
//!
//!   edge-bench <range dir> | --group <group dir> [--from article] [--to article] [--limit ranges] [--rules rules.json]
//!              [--batch-chunks 10] [--idle 5] [--no-finish] [--dump out.jsonl] [--wire batches.jsonl|-] [--releases out.jsonl]

use base64::Engine;
use edge::indexer::{self, Batch, RangeIndexer};
use edge::js::json_str;
use edge::{names, wire};
use rustc_hash::FxHashMap;
use std::io::Write;
use std::time::Instant;

// Many small strings (message-ids, names): the system allocator was a measurable share.
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

/// A chunk as the news server sent it (zlib; gzip also taken), whole: libdeflate wants the output
/// size, so it grows until the text fits (~3.6 times the compressed size on boneless).
fn inflate(raw: &[u8], d: &mut libdeflater::Decompressor) -> Vec<u8> {
    let mut out = vec![0u8; raw.len() * 4 + 4096];
    loop {
        let result = if raw.starts_with(&[0x1f, 0x8b]) { d.gzip_decompress(raw, &mut out) } else { d.zlib_decompress(raw, &mut out) };
        match result {
            Ok(n) => {
                out.truncate(n);
                return out;
            }
            Err(libdeflater::DecompressionError::InsufficientSpace) => out.resize(out.len() * 2, 0),
            Err(e) => panic!("inflate: {e}"),
        }
    }
}

fn json_field(json: &str, name: &str) -> String {
    let at = json.find(&format!("\"{name}\":")).unwrap() + name.len() + 3;
    let rest = json[at..].trim_start();
    if let Some(s) = rest.strip_prefix('"') {
        s[..s.find('"').unwrap()].to_string()
    } else {
        rest[..rest.find([',', '}']).unwrap()].trim().to_string()
    }
}

fn peak_rss_mb() -> u64 {
    let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
    status.lines().find(|l| l.starts_with("VmHWM:")).and_then(|l| l.split_whitespace().nth(1)).and_then(|k| k.parse::<u64>().ok()).unwrap_or(0) / 1024
}

fn opt_num(n: Option<i64>, out: &mut String) {
    match n {
        Some(n) => out.push_str(&n.to_string()),
        None => out.push_str("null"),
    }
}

/// The --releases dump: per release chunk its key, name, rule, counter total, poster, time, counts and first subjects.
fn dump_releases(batch: &Batch, group: &str, out: &mut impl Write) {
    let mut line = String::new();
    for r in &batch.releases {
        line.push_str("{\"group\":");
        json_str(group, &mut line);
        line.push_str(",\"key\":");
        json_str(&r.key, &mut line);
        line.push_str(",\"name\":");
        json_str(&r.name, &mut line);
        line.push_str(&format!(",\"source\":\"{}\",\"fileTotal\":", r.source));
        opt_num(r.file_total, &mut line);
        line.push_str(",\"poster\":");
        json_str(&r.poster, &mut line);
        line.push_str(",\"postedAt\":");
        opt_num(r.posted_at, &mut line);
        line.push_str(",\"postedLast\":");
        opt_num(r.posted_last, &mut line);
        line.push_str(&format!(",\"files\":{},\"segments\":{},\"bytes\":{},\"firstArticle\":{},\"hasPar2\":{},\"subjects\":[", r.files, r.segments, r.bytes, r.first_article, r.has_par2));
        for (i, s) in r.subjects.iter().enumerate() {
            if i > 0 {
                line.push(',');
            }
            json_str(s, &mut line);
        }
        line.push_str("]}\n");
        out.write_all(line.as_bytes()).unwrap();
        line.clear();
    }
}

fn dump_batch(batch: &Batch, out: &mut impl Write) {
    let b64 = base64::engine::general_purpose::STANDARD;
    let mut line = String::new();
    let mut emit = |line: &mut String| {
        line.push('\n');
        out.write_all(line.as_bytes()).unwrap();
        line.clear();
    };
    let opt = |t: Option<i64>| t.map_or("null".to_string(), |t| t.to_string());
    line.push_str(&format!(
        "{{\"t\":\"b\",\"from\":{},\"next\":{},\"headers\":{},\"postedLow\":{},\"postedHigh\":{}}}",
        batch.from, batch.next, batch.headers, opt(batch.posted_low), opt(batch.posted_high)
    ));
    emit(&mut line);
    let release = |r: &indexer::ReleaseChunk, kind: &str, line: &mut String| {
        line.push_str(&format!("{{\"t\":\"{kind}\",\"upload\":{},\"key\":", r.upload));
        json_str(&r.key, line);
        line.push_str(",\"name\":");
        json_str(&r.name, line);
        line.push_str(&format!(",\"source\":\"{}\",\"fileTotal\":", r.source));
        opt_num(r.file_total, line);
        line.push_str(",\"poster\":");
        json_str(&r.poster, line);
        line.push_str(",\"groups\":[");
        for (i, g) in r.groups.iter().enumerate() {
            if i > 0 {
                line.push(',');
            }
            json_str(g, line);
        }
        line.push_str("],\"postedAt\":");
        opt_num(r.posted_at, line);
        line.push_str(",\"postedLast\":");
        opt_num(r.posted_last, line);
        line.push_str(&format!(
            ",\"files\":{},\"segments\":{},\"bytes\":{},\"firstArticle\":{},\"lastArticle\":{},\"readable\":{},\"hasPar2\":{},\"finished\":{},\"edge\":{},\"partial\":{},\"payload\":\"{}\"",
            r.files,
            r.segments,
            r.bytes,
            r.first_article,
            r.last_article,
            r.readable,
            r.has_par2,
            r.finished,
            r.edge,
            r.partial,
            b64.encode(&r.payload)
        ));
    };
    for r in &batch.releases {
        release(r, "r", &mut line);
        line.push('}');
        emit(&mut line);
    }
    for (r, rule, key) in &batch.samples {
        release(r, "s", &mut line);
        line.push_str(&format!(",\"rule\":\"{rule}\",\"ruleKey\":"));
        json_str(key, &mut line);
        line.push('}');
        emit(&mut line);
    }
    for c in &batch.candidates {
        line.push_str(&format!("{{\"t\":\"c\",\"upload\":{},\"releaseKey\":", c.upload));
        json_str(&c.release_key, &mut line);
        line.push_str(&format!(",\"kind\":\"{}\",\"messageId\":", c.kind));
        json_str(&c.message_id, &mut line);
        line.push_str(",\"file\":");
        json_str(&c.file, &mut line);
        line.push('}');
        emit(&mut line);
    }
    for (rule, key, articles, files) in &batch.counts {
        line.push_str(&format!("{{\"t\":\"n\",\"rule\":\"{rule}\",\"key\":"));
        json_str(key, &mut line);
        line.push_str(&format!(",\"articles\":{articles},\"files\":{files}}}"));
        emit(&mut line);
    }
}

/// A stored range: its directory and article span (range.json).
struct Range {
    dir: String,
    start: i64,
    end: i64,
}

/// The range in a directory with a range.json, and its group.
fn range_in(dir: &str) -> Option<(String, Range)> {
    let json = std::fs::read_to_string(format!("{dir}/range.json")).ok()?;
    Some((json_field(&json, "group"), Range { dir: dir.to_string(), start: json_field(&json, "start").parse().ok()?, end: json_field(&json, "end").parse().ok()? }))
}

/// A group directory's complete ranges in article order (the corpus layout: `<group>/<start>-<end>/`).
fn ranges_in(group_dir: &str) -> (String, Vec<Range>) {
    let mut group = String::new();
    let mut ranges: Vec<Range> = Vec::new();
    for entry in std::fs::read_dir(group_dir).unwrap_or_else(|e| panic!("{group_dir}: {e}")) {
        let path = entry.unwrap().path();
        if !path.is_dir() {
            continue;
        }
        if let Some((g, range)) = range_in(path.to_str().unwrap()) {
            group = g;
            ranges.push(range);
        }
    }
    ranges.sort_by_key(|r| r.start);
    (group, ranges)
}

fn date(ms: Option<i64>) -> String {
    match ms.and_then(chrono::DateTime::from_timestamp_millis) {
        Some(t) => t.format("%Y-%m-%d %H:%M").to_string(),
        None => "-".to_string(),
    }
}

/// Files waiting in finished releases after which a batch goes out whatever its chunks (edge-node's limit).
const MAX_PENDING_FILES: usize = 100_000;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let usage = "usage: edge-bench <range dir> | --group <group dir> [--from article] [--to article] [--limit ranges] [--rules rules.json] [--batch-chunks N] [--idle N] [--no-finish] [--dump file] [--wire file|-] [--releases file]";
    let flag = |name: &str| args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned();
    let has = |name: &str| args.iter().any(|a| a == name);
    let batch_chunks: usize = flag("--batch-chunks").map(|v| v.parse().unwrap()).unwrap_or(10);
    let idle: u64 = flag("--idle").map(|v| v.parse().unwrap()).unwrap_or(5);
    let mut dump = flag("--dump").map(|p| std::io::BufWriter::new(std::fs::File::create(p).unwrap()));
    // Each batch as the worker posts it (wire.rs), one JSON per line: a file, or stdout (`-`) for replay:corpus.
    let wire_path = flag("--wire");
    let wire_stdout = wire_path.as_deref() == Some("-");
    let mut wire_out: Option<Box<dyn Write>> = match wire_path.as_deref() {
        Some("-") => Some(Box::new(std::io::BufWriter::with_capacity(1 << 20, std::io::stdout().lock()))),
        Some(p) => Some(Box::new(std::io::BufWriter::new(std::fs::File::create(p).unwrap()))),
        None => None,
    };
    // One line per release chunk without its payload, with its first subjects (deploy/corpus_sweep.py).
    let mut releases_out = flag("--releases").map(|p| std::io::BufWriter::new(std::fs::File::create(p).unwrap()));

    // One range, or a group's ranges in order through one window (the window carries over from range to range;
    // the end is only finished after the last one, unless --no-finish: a replay to be continued).
    let (group, ranges) = match flag("--group") {
        Some(dir) => {
            let (group, mut ranges) = ranges_in(&dir);
            // --from: ranges ending before it are skipped (a replay continued); --to: ranges starting after it (one
            // window of the group's history, the corpus holding others); --limit: this many at most.
            if let Some(from) = flag("--from").map(|v| v.parse::<i64>().unwrap()) {
                ranges.retain(|r| r.end >= from);
            }
            if let Some(to) = flag("--to").map(|v| v.parse::<i64>().unwrap()) {
                ranges.retain(|r| r.start <= to);
            }
            if let Some(limit) = flag("--limit").map(|v| v.parse::<usize>().unwrap()) {
                ranges.truncate(limit);
            }
            (group, ranges)
        }
        None => {
            let positional = (1..args.len()).find(|&i| !args[i].starts_with("--") && !args[i - 1].starts_with("--")).map(|i| args[i].clone());
            let dir = positional.expect(usage);
            let (group, range) = range_in(&dir).unwrap_or_else(|| panic!("{dir}: no range.json"));
            (group, vec![range])
        }
    };
    if ranges.is_empty() {
        eprintln!("no complete range found");
        std::process::exit(2);
    }
    let finish = !has("--no-finish");
    let rules = match flag("--rules") {
        Some(path) => {
            let json = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("{path}: {e}"));
            serde_json::from_str::<edge::lease::RulesJson>(&json).unwrap_or_else(|e| panic!("{path}: {e}")).rules()
        }
        None => {
            let mut priorities = FxHashMap::default();
            priorities.insert(group.clone(), 1);
            names::Rules::defaults(priorities)
        }
    };
    // The report goes to stderr when the wire takes stdout.
    let say = |line: String| {
        if wire_stdout {
            eprintln!("{line}");
        } else {
            println!("{line}");
        }
    };

    let (start, end) = (ranges[0].start, ranges[ranges.len() - 1].end);
    let mut indexer = RangeIndexer::new(&group, start, end, start, rules, idle);

    let (mut inflate_s, mut index_s) = (0.0, 0.0);
    let (mut compressed, mut text_bytes, mut chunks) = (0usize, 0usize, 0usize);
    let mut out = (0u64, 0usize, 0usize, 0usize, 0usize, 0usize, 0usize); // headers, releases, files, segments, payload, candidates, samples
    let mut finished = 0usize;
    let mut keys = std::collections::HashSet::new();
    let mut decompressor = libdeflater::Decompressor::new();
    let wall = Instant::now();
    let mut since_batch = 0usize;
    for (ri, range) in ranges.iter().enumerate() {
        let range_wall = Instant::now();
        let headers_before = out.0;
        // After a gap between stored ranges the batch starts at this range (every range ends with a batch taken).
        if ri > 0 {
            indexer.skip_to(range.start);
        }
        let mut files: Vec<String> = std::fs::read_dir(&range.dir).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|f| f.ends_with(".gz")).collect();
        files.sort();
        for (i, file) in files.iter().enumerate() {
            let raw = std::fs::read(format!("{}/{file}", range.dir)).unwrap();
            compressed += raw.len();
            let t = Instant::now();
            let data = inflate(&raw, &mut decompressor);
            inflate_s += t.elapsed().as_secs_f64();
            text_bytes += data.len();
            chunks += 1;
            let t = Instant::now();
            let text = String::from_utf8_lossy(&data);
            let next: i64 = file.split('-').nth(2).unwrap().trim_end_matches(".gz").parse::<i64>().unwrap() + 1;
            indexer.add_text(&text, next);
            since_batch += 1;
            let range_end = i + 1 == files.len();
            let last = range_end && ri + 1 == ranges.len();
            if last && finish {
                indexer.finish();
            }
            if !range_end && since_batch < batch_chunks && indexer.pending() < MAX_PENDING_FILES {
                index_s += t.elapsed().as_secs_f64();
                continue;
            }
            let batch = indexer.take();
            since_batch = 0;
            index_s += t.elapsed().as_secs_f64();
            out.0 += batch.headers;
            out.1 += batch.releases.len();
            for r in &batch.releases {
                keys.insert(r.key.clone());
                out.2 += r.files;
                out.3 += r.segments;
                out.4 += r.payload.len();
                finished += r.finished as usize;
            }
            out.5 += batch.candidates.len();
            out.6 += batch.samples.len();
            if let Some(d) = dump.as_mut() {
                dump_batch(&batch, d);
            }
            if let Some(w) = wire_out.as_mut() {
                w.write_all(&wire::to_json(&batch)).unwrap();
                w.write_all(b"\n").unwrap();
                w.flush().unwrap();
            }
            if let Some(w) = releases_out.as_mut() {
                dump_releases(&batch, &group, w);
            }
        }
        if ranges.len() > 1 {
            eprintln!(
                "{} {}–{}: {} headers in {:.1} s; frontier {}; {} releases out ({} finished), {} releases and {} files open; RSS {} MB",
                group,
                range.start,
                range.end,
                out.0 - headers_before,
                range_wall.elapsed().as_secs_f64(),
                date(indexer.frontier()),
                out.1,
                finished,
                indexer.open_releases(),
                indexer.open_files(),
                peak_rss_mb()
            );
        }
    }
    let total = inflate_s + index_s;
    #[allow(static_mut_refs)]
    let p = unsafe { indexer::PROF };
    eprintln!("  of which aggregate {:.2} s, evict+cutoff {:.2} s, judge+encode {:.2} s", p[3] as f64 / 1e9, p[4] as f64 / 1e9, p[5] as f64 / 1e9);
    say(format!("{group} {start}–{end}: {} ranges, {chunks} chunks, {} MB compressed, {} MB text", ranges.len(), compressed / 1_000_000, text_bytes / 1_000_000));
    say(format!("{} headers in {:.2} s ({:.0}/s on one core; wall {:.2} s)", out.0, total, out.0 as f64 / total, wall.elapsed().as_secs_f64()));
    say(format!("  inflate {inflate_s:.2} s, index {index_s:.2} s; peak RSS {} MB", peak_rss_mb()));
    say(format!(
        "  kept: {} release chunks ({} releases, {} finished), {} files, {} segments, payload {:.1} MB ({:.1} B/segment); {} candidates, {} samples",
        out.1,
        keys.len(),
        finished,
        out.2,
        out.3,
        out.4 as f64 / 1e6,
        out.4 as f64 / out.3.max(1) as f64,
        out.5,
        out.6
    ));
}