#!/usr/bin/env python3
"""The grouping sweeps' corpus (REBUILD.md, corpus): stored XOVER chunks replayed with each build's rules.

Run on homelab from the deployment root (/mnt/bulk/usenet-rebuild):

  python3 src/deploy/corpus_sweep.py capture [N]   mark N more ranges to store (default 100), fetched first
  python3 src/deploy/corpus_sweep.py replay        edge-bench over every stored range with the deployed edge image
  python3 src/deploy/corpus_sweep.py report        what the replay grouped: suspects, and the change since the last build

The prefetcher stores the chunks of ranges marked `capture` in corpus/<group>/<start>-<end>/ (range.json once
complete) and reads them from there afterwards, so a reset doesn't fetch them again. Replays go to
corpus-out/<image id>/<group>/<start>-<end>.jsonl, one per build; `report` compares the newest two.
"""
import collections
import json
import os
import shutil
import subprocess
import sys
from datetime import datetime, timezone

ROOT = os.getcwd()
CORPUS = os.path.join(ROOT, 'corpus')
OUT = os.path.join(ROOT, 'corpus-out')
IMAGE = 'usenet-rebuild-edge:latest'
# The corpus may take this much of the HDD (user, 2026-10-12: up to 300 GB).
MAX_BYTES = 300e9
# Ranks in groups.tsv the new ranges come from first (the top 100), then the rest of what's indexed.
TOP = 100


def env(name):
    for line in open(os.path.join(ROOT, '.env')):
        if line.startswith(name + '='):
            return line.split('=', 1)[1].strip()
    raise SystemExit(f'{name} not in .env')


def psql(sql):
    """Runs SQL in the catalog; rows as lists of strings."""
    out = subprocess.run(
        ['docker', 'compose', 'exec', '-T', 'postgres', 'psql', '-U', env('DB_USER'), '-d', env('DB_DATABASE'), '-At', '-F', '\t', '-v', 'ON_ERROR_STOP=1'],
        input=sql, capture_output=True, text=True, check=True,
    ).stdout
    return [line.split('\t') for line in out.splitlines() if line]


def corpus_bytes():
    total = 0
    for dirpath, _, files in os.walk(CORPUS):
        total += sum(os.path.getsize(os.path.join(dirpath, f)) for f in files)
    return total


def capture(n):
    used = corpus_bytes()
    if used >= MAX_BYTES:
        print(f'corpus at {used / 1e9:.0f} GB of {MAX_BYTES / 1e9:.0f}: nothing more marked')
        return
    top = [line.split('\t')[1] for line in open(os.path.join(ROOT, 'groups.tsv')).read().splitlines()[:TOP]]
    names = ','.join("'" + g.replace("'", "''") + "'" for g in top)
    # Spread over the groups and their history: per group, the ranges farthest from those stored already
    # (at random in a group without any), the top-100 groups first. Only ranges not started (a range half
    # spooled would miss its first chunks; a done one is stored at the next reset, which re-reads them all).
    # Priority 0: fetched before the rest of the history (the tail stays first).
    rows = psql(f"""
WITH stored AS (SELECT group_id, array_agg(start_id) AS starts FROM ranges WHERE capture GROUP BY 1),
pick AS (
  SELECT r.id, g.name IN ({names}) AS top, row_number() OVER (PARTITION BY r.group_id ORDER BY
    coalesce((SELECT min(abs(r.start_id - s)) FROM unnest(st.starts) s), 0) DESC, random()) AS n
  FROM ranges r JOIN groups g ON g.id = r.group_id LEFT JOIN stored st ON st.group_id = r.group_id
  WHERE r.priority >= 0 AND NOT r.capture AND r.status = 'planned' AND r.spooled_chunks = 0 AND r.next_id = r.start_id
),
chosen AS (SELECT id FROM pick ORDER BY top DESC, n, random() LIMIT {int(n)})
UPDATE ranges r SET capture = true, priority = 0 FROM chosen WHERE r.id = chosen.id RETURNING r.id;
""")
    total = psql('SELECT count(*), count(*) FILTER (WHERE status = \'done\') FROM ranges WHERE capture;')[0]
    print(f'marked {len(rows)} ranges; {total[0]} marked in all, {total[1]} indexed; corpus {used / 1e9:.1f} GB')


def complete_ranges():
    for group in sorted(os.listdir(CORPUS)) if os.path.isdir(CORPUS) else []:
        for rng in sorted(os.listdir(os.path.join(CORPUS, group))):
            if os.path.exists(os.path.join(CORPUS, group, rng, 'range.json')):
                yield group, rng


def image_id():
    return subprocess.run(['docker', 'image', 'inspect', '-f', '{{.Id}}', IMAGE], capture_output=True, text=True, check=True).stdout.strip().split(':')[-1][:12]


def replay():
    build = image_id()
    out = os.path.join(OUT, build)
    todo = [(g, r) for g, r in complete_ranges() if not os.path.exists(os.path.join(out, g, r + '.jsonl'))]
    for g, _ in todo:
        os.makedirs(os.path.join(out, g), exist_ok=True)
    print(f'build {build}: {len(todo)} ranges to replay', flush=True)
    if not todo:
        return
    # One container, 6 at a time (the edge worker keeps its threads).
    script = 'xargs -P 6 -n 2 sh -c \'edge-bench "/corpus/$0/$1" --releases "/out/$0/$1.jsonl.part" >/dev/null 2>&1 && mv "/out/$0/$1.jsonl.part" "/out/$0/$1.jsonl"\''
    started = datetime.now()
    subprocess.run(
        ['docker', 'run', '--rm', '-i', '--user', '1000:1000', '-v', f'{CORPUS}:/corpus:ro', '-v', f'{out}:/out', '--entrypoint', 'sh', IMAGE, '-c', script],
        input=''.join(f'{g} {r}\n' for g, r in todo), text=True, check=True,
    )
    print(f'replayed in {(datetime.now() - started).total_seconds():.0f} s')
    # Older builds' replays beyond the last two go.
    builds = sorted((os.path.getmtime(os.path.join(OUT, b)), b) for b in os.listdir(OUT))
    for _, b in builds[:-2]:
        shutil.rmtree(os.path.join(OUT, b))


def load(build):
    """Releases of one build's replay, merged by (group, key) as the catalog does."""
    releases = {}
    for dirpath, _, files in os.walk(os.path.join(OUT, build)):
        for f in files:
            if not f.endswith('.jsonl'):
                continue
            for line in open(os.path.join(dirpath, f), encoding='utf-8', errors='replace'):
                r = json.loads(line)
                k = (r['group'], r['key'])
                have = releases.get(k)
                if have is None:
                    releases[k] = r
                else:
                    have['files'] += r['files']
                    have['segments'] += r['segments']
                    have['subjects'] = (have['subjects'] + r['subjects'])[:3]
                    have['fileTotal'] = have['fileTotal'] or r['fileTotal']
    return releases


def hour(r):
    return datetime.fromtimestamp((r['postedAt'] or 0) / 1000, timezone.utc).strftime('%Y-%m-%d %H')


def show(r, prefix='    '):
    print(f"{prefix}{r['files']} files, total {r['fileTotal']}, {r['source']}: {r['name'][:80]!r} ({r['poster'][:30]}, {r['group']})")
    for s in r['subjects'][:2]:
        print(f'{prefix}  {s[:150]}')


def report():
    builds = sorted((os.path.getmtime(os.path.join(OUT, b)), b) for b in os.listdir(OUT)) if os.path.isdir(OUT) else []
    if not builds:
        raise SystemExit('nothing replayed yet')
    build = builds[-1][1]
    rel = load(build)
    ranges = sum(1 for _ in complete_ranges())
    print(f'build {build}: {len(rel):,} releases, {sum(r["files"] for r in rel.values()):,} files from {ranges} stored ranges ({corpus_bytes() / 1e9:.1f} GB)')

    print('\n== same name, poster and hour: 4+ releases (under-grouping)')
    clusters = collections.defaultdict(list)
    for r in rel.values():
        clusters[(r['group'], r['poster'], r['name'], hour(r))].append(r)
    for (g, p, name, h), rs in sorted(clusters.items(), key=lambda kv: -len(kv[1]))[:25]:
        if len(rs) < 4:
            break
        print(f"  {len(rs)} releases, {sum(r['files'] for r in rs)} files: {name[:70]!r} ({p[:25]}, {g}, {h})")
        for s in rs[0]['subjects'][:1] + rs[-1]['subjects'][:1]:
            print(f'      {s[:150]}')

    print('\n== single files of a counter of 5+, per group and poster (under-grouping)')
    singles = collections.defaultdict(list)
    for r in rel.values():
        if r['files'] == 1 and (r['fileTotal'] or 0) >= 5:
            singles[(r['group'], r['poster'])].append(r)
    for (g, p), rs in sorted(singles.items(), key=lambda kv: -len(kv[1]))[:15]:
        if len(rs) < 20:
            break
        print(f'  {len(rs)} single files: {p[:30]} ({g})')
        for s in rs[0]['subjects'][:1] + rs[-1]['subjects'][:1]:
            print(f'      {s[:150]}')

    print('\n== 2000+ files, or 3x the counter (over-grouping)')
    big = [r for r in rel.values() if r['files'] > 2000 or (r['fileTotal'] and r['files'] > 3 * r['fileTotal'] and r['files'] > 50)]
    for r in sorted(big, key=lambda r: -r['files'])[:20]:
        show(r, '  ')

    if len(builds) >= 2:
        before = load(builds[-2][1])
        print(f'\n== change since build {builds[-2][1]}: {len(before):,} → {len(rel):,} releases')
        names = lambda rs: collections.Counter((r['group'], r['name']) for r in rs.values())
        a, b = names(before), names(rel)
        moved = sorted(set(a) | set(b), key=lambda k: -abs(a[k] - b[k]))[:20]
        for k in moved:
            if a[k] != b[k]:
                print(f'  {a[k]:>6} → {b[k]:<6} {k[1][:80]!r} ({k[0]})')


if __name__ == '__main__':
    command = sys.argv[1] if len(sys.argv) > 1 else 'report'
    if command == 'capture':
        capture(int(sys.argv[2]) if len(sys.argv) > 2 else 100)
    elif command == 'replay':
        replay()
    elif command == 'report':
        report()
    else:
        raise SystemExit(__doc__)