#!/usr/bin/env python3
import argparse
import json
import os
import shlex
import subprocess
import time
from pathlib import Path


MODEL = "deepseek-v4-flash"


def line_count(path: Path) -> int:
    if not path.exists():
        return 0
    with path.open("rb") as handle:
        return sum(1 for _ in handle)


def pid_alive(pid_path: Path) -> bool:
    if not pid_path.exists():
        return False
    try:
        pid = int(pid_path.read_text(encoding="utf-8").strip())
    except Exception:
        return False
    return subprocess.run(
        ["ps", "-p", str(pid)],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    ).returncode == 0


def recent_mtime(paths: list[Path]) -> float:
    mtimes = [p.stat().st_mtime for p in paths if p.exists()]
    return max(mtimes) if mtimes else 0


def start_worker(run_id: str, shard_dir: Path, run_dir: Path, cache_root: Path, worker: str):
    inp = shard_dir / f"{worker}.jsonl"
    out = run_dir / f"{worker}.outputs.jsonl"
    report = run_dir / f"{worker}.report.json"
    cache = cache_root / worker
    run_dir.mkdir(parents=True, exist_ok=True)
    cache.mkdir(parents=True, exist_ok=True)
    cmd = (
        "set -euo pipefail; "
        "if [ -f ~/.aisim_simulation_env ]; then set -a; source ~/.aisim_simulation_env; set +a; fi; "
        f"SIMULATION_RUN_ID={shlex.quote(run_id)} "
        "LLM_API_KEY_ENV=OLLAMA_CLOUD_API_KEY "
        'LLM_BASE_URL="${OLLAMA_CLOUD_BASE_URL:-https://ollama.com/v1}" '
        f"LLM_MODEL={shlex.quote(MODEL)} "
        f"scripts/start_simulation_worker_v1.sh {shlex.quote(worker)} "
        f"{shlex.quote(str(inp))} {shlex.quote(str(out))} "
        f"{shlex.quote(str(report))} {shlex.quote(str(cache))}"
    )
    return subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True)


def worker_status(shard_dir: Path, run_dir: Path, cache_root: Path, worker: str):
    inp = shard_dir / f"{worker}.jsonl"
    out = run_dir / f"{worker}.outputs.jsonl"
    report = run_dir / f"{worker}.report.json"
    pid = run_dir / f"{worker}.pid"
    log = run_dir / f"{worker}.log"
    target = line_count(inp)
    output_count = line_count(out)
    cache_count = len(list((cache_root / worker).glob("*.json"))) if (cache_root / worker).exists() else 0
    alive = pid_alive(pid)
    mtime = recent_mtime([out, report, log])
    age = int(time.time() - mtime) if mtime else None
    failures = None
    validation_problem_count = None
    if report.exists():
        try:
            data = json.loads(report.read_text(encoding="utf-8"))
            failures = len(data.get("failures", []))
            validation_problem_count = len(data.get("validation_problems", {}))
        except Exception:
            pass
    return {
        "worker": worker,
        "target": target,
        "outputs": output_count,
        "cache_files": cache_count,
        "alive": alive,
        "idle_seconds": age,
        "failures": failures,
        "validation_problem_count": validation_problem_count,
        "complete": target > 0 and output_count >= target,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--host-prefix", required=True)
    parser.add_argument("--workers", default="1,2,3,4")
    parser.add_argument("--shard-dir", default="")
    parser.add_argument("--run-dir", default="")
    parser.add_argument("--cache-root", default="")
    parser.add_argument("--interval", type=int, default=300)
    parser.add_argument("--stale-seconds", type=int, default=1800)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()

    shard_dir = Path(args.shard_dir or (Path("data/simulations/shards") / args.run_id))
    run_dir = Path(args.run_dir or (Path("data/simulations/runs") / args.run_id))
    cache_root = Path(args.cache_root or (Path("data/cache/simulation_runs") / args.run_id))
    worker_ids = [int(x.strip()) for x in args.workers.split(",") if x.strip()]
    workers = [f"{args.host_prefix}-{i:02d}" for i in worker_ids]
    state_path = run_dir / f"{args.host_prefix}.supervisor_state.json"
    log_path = run_dir / f"{args.host_prefix}.supervisor.log"
    run_dir.mkdir(parents=True, exist_ok=True)

    while True:
        statuses = []
        actions = []
        for worker in workers:
            status = worker_status(shard_dir, run_dir, cache_root, worker)
            statuses.append(status)
            if not (shard_dir / f"{worker}.jsonl").exists():
                actions.append({"worker": worker, "action": "missing_shard"})
                continue
            if status["complete"]:
                continue
            stale = status["idle_seconds"] is not None and status["idle_seconds"] > args.stale_seconds
            if (not status["alive"]) or stale:
                result = start_worker(args.run_id, shard_dir, run_dir, cache_root, worker)
                actions.append(
                    {
                        "worker": worker,
                        "action": "restart" if stale else "start",
                        "returncode": result.returncode,
                        "stdout_tail": result.stdout[-300:],
                        "stderr_tail": result.stderr[-300:],
                    }
                )
                time.sleep(2)
        state = {
            "run_id": args.run_id,
            "host_prefix": args.host_prefix,
            "updated_at_unix": time.time(),
            "statuses": statuses,
            "actions": actions,
        }
        state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        with log_path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(state, ensure_ascii=False, separators=(",", ":")) + "\n")
        print(json.dumps(state, ensure_ascii=False, indent=2))
        if args.once or all(s["complete"] for s in statuses):
            return
        time.sleep(args.interval)


if __name__ == "__main__":
    main()