#!/usr/bin/env python3
"""Re-balance the download queues (run on M1M32).

Contiguous CIK blocks are badly balanced (old CIKs = large filers).  Companies
already *started* stay on their queue (the downloader resumes them); the
never-started companies are dealt round-robin over all queues so every queue
mixes old and new filers.  Then the downloaders are restarted.
"""
from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from orchestrate import BHS64, BHS64_DL, CODE, MASTER, NODES, REMOTE_ROOT, log, ssh  # noqa: E402

SHARDS = CODE / "data" / "metadata" / "shards.json"


def listed(host: str, path: str) -> set[str]:
    r = ssh(host, f"ls {path} 2>/dev/null")
    return {l.replace(".json", "").strip() for l in r.stdout.splitlines() if l.strip().endswith(".json")}


def main() -> None:
    man = json.loads(SHARDS.read_text())
    workers = [n for n in NODES if n != MASTER]
    all_new: set[str] = set()
    for n in workers:
        all_new |= set(man["nodes"][n]["self"]) | set(man["nodes"][n]["bhs64"])

    started: dict[tuple[str, str], set[str]] = {}
    for n in workers:
        started[(n, "self")] = listed(NODES[n][0], f"{REMOTE_ROOT}/data/metadata/submissions") & all_new
        started[(n, "bhs64")] = listed(BHS64, f"{BHS64_DL}/{n}/submissions") & all_new
    taken = set().union(*started.values())
    unstarted = sorted(all_new - taken)
    log(f"rebalance: {len(all_new)} new CIKs, {len(taken)} started, {len(unstarted)} to redistribute")

    queues = [(n, k) for n in workers for k in ("self", "bhs64")]
    new_lists = {q: sorted(started[q]) for q in queues}
    for i, cik in enumerate(unstarted):
        new_lists[queues[i % len(queues)]].append(cik)

    # stop every downloader (files already on disk are kept; the log is append-only)
    for n in workers:
        ssh(NODES[n][0], "pkill -f 'scripts/download_filings.py'; rm -f "
                         f"{REMOTE_ROOT}/data/logs/stage.download_self.done {REMOTE_ROOT}/data/logs/stage.download_self.running")
    ssh(BHS64, "pkill -f 'scripts/download_filings.py'; rm -f " + " ".join(f"{BHS64_DL}/{n}/DONE" for n in workers))

    shard_dir = CODE / "data" / "metadata" / "shards"
    for (n, k), ciks in new_lists.items():
        (shard_dir / f"{n}.{k}.txt").write_text("\n".join(ciks) + ("\n" if ciks else ""))
        man["nodes"][n][k] = ciks
        man["nodes"][n]["n_total"] = len(man["nodes"][n]["self"]) + len(man["nodes"][n]["bhs64"])
        log(f"  {n:7s} {k:5s} -> {len(ciks):4d} CIKs ({len(started[(n, k)])} resumed)")
    man["rebalanced"] = True
    SHARDS.write_text(json.dumps(man, indent=2))

    for n in workers:
        subprocess.run(["rsync", "-a", f"{shard_dir}/", f"{NODES[n][0]}:{REMOTE_ROOT}/data/metadata/shards/"], check=True)
    subprocess.run(["rsync", "-a", f"{shard_dir}/", f"{BHS64}:{REMOTE_ROOT}/data/metadata/shards/"], check=True)
    subprocess.run(["rsync", "-a", str(SHARDS), f"{BHS64}:{REMOTE_ROOT}/data/metadata/shards.json"], check=True)
    log("rebalance: lists pushed; relaunching downloads")
    subprocess.run([sys.executable, str(Path(__file__).resolve().parent / "orchestrate.py"), "downloads"], check=True)


if __name__ == "__main__":
    main()