"""Launch the grading panel (and the stage-1 extractor) on the head Spark.

Each role agent is a ONE-SHOT, hardened, read-only container running
sandbox/grader_agent.py: it reads the deck text mounted at /docs, runs its model
(through the on-Spark proxy) under its persona + the BDEF rubric, writes exactly
one output file to /out, and exits. There is no tool loop and no shared writable
workspace — agents cannot alter the deck text or each other's reports.

Per-deck mount contract (remote paths under {remoteWorkDir}/jobs/<job>/<company>/<deck>):
    docs/                          -> /docs:ro
    BDEF.md                        -> /BDEF.md:ro
    schemas/<role>.schema.json     -> /schema.json:ro   (extractor|grades)
    personas/<rid>.md              -> /persona/PERSONA.md:ro
    out/                           -> /out:rw

Sandbox (per the operator's confidentiality requirement):
  * non-root, --cap-drop ALL, --security-opt no-new-privileges
  * read-only rootfs + small writable tmpfs (/tmp, /home/rev)
  * NO docker socket, cpu/mem/pid caps
  * attached to the per-job network — in airgapped mode that network is
    --internal, so the agent can reach ONLY the model proxy, never the internet

Agents hold no credentials beyond a dummy proxy key.
"""
from __future__ import annotations

import os
import re
import shlex
import shutil

import bm_config
import prompts
import serving
import spark_client as sc

ORCH_DIR = os.path.dirname(os.path.abspath(__file__))
SANDBOX_SRC = os.path.join(ORCH_DIR, "sandbox")
SCHEMAS_SRC = os.path.join(ORCH_DIR, "schemas")
SCHEMA_FILES = ("extraction.schema.json", "grades.schema.json")

HARDEN = (
    "--user 1000:1000 --security-opt no-new-privileges --cap-drop ALL "
    "--read-only --tmpfs /tmp:size=256m --tmpfs /home/rev:size=128m "
    "--pids-limit 256 --memory 6g --cpus 4"
)


# ------------------------------------------------------------------ image
def ensure_grader_image(cfg: dict, log) -> None:
    """Build the grader image on the head Spark if missing (aarch64, native)."""
    head = sc.head(cfg)
    image = cfg["graderImage"]
    q = shlex.quote
    r = sc.run(head, f"docker image inspect {q(image)} >/dev/null 2>&1 && echo PRESENT || echo MISSING",
               timeout=30)
    if "PRESENT" in (r.stdout or ""):
        log(f"[graders] image {image} already present on {head.host}")
        return
    if not os.path.isdir(SANDBOX_SRC):
        raise RuntimeError(f"grader build context missing at {SANDBOX_SRC} (image not baked in?)")
    remote_dir = f"{cfg['remoteWorkDir']}/sandbox-build"
    log(f"[graders] building grader image {image} on {head.host} (first run; a few minutes)…")
    push = sc.push_dir(head, SANDBOX_SRC, remote_dir)
    if push.returncode != 0:
        raise RuntimeError(f"rsync grader build context to {head.host} failed: {push.stderr}")
    b = sc.run(head, f"cd {q(remote_dir)} && IMAGE={q(image)} bash build.sh", timeout=1800)
    if b.returncode != 0:
        raise RuntimeError(f"grader image build failed on {head.host}: {b.stderr or b.stdout}")
    log(f"[graders] grader image built: {image}")


# ------------------------------------------------------------------ roster
def slug(name: str) -> str:
    s = re.sub(r"[^A-Za-z0-9]+", "-", (name or "").strip().lower()).strip("-")
    return s or "grader"


def roster(cfg: dict) -> list[dict]:
    """Grader roster from config. Each: {rid, name, model alias, persona, temperature}."""
    out: list[dict] = []
    seen: dict[str, int] = {}
    for w in (cfg.get("graders") or []):
        name = (w.get("name") or "grader").strip()
        rid = slug(name)
        if rid in seen:
            seen[rid] += 1
            rid = f"{rid}-{seen[rid]}"
        else:
            seen[rid] = 1
        out.append({
            "rid": rid, "name": name, "model": (w.get("model") or "").strip(),
            "persona": (w.get("persona") or "").strip(),
            "temperature": w.get("temperature"),
        })
    return out


# ------------------------------------------------------------------ staging
def stage_deck_files(cfg: dict, local_deck_dir: str, panel: list[dict]) -> None:
    """Write everything the role containers need into the LOCAL per-deck staging
    dir (jobs.py rsyncs the whole dir to the Spark afterwards): the BDEF rubric,
    both schemas, and one persona file per role agent (extractor + each grader +
    adjudicator). Also pre-creates out/ and adjudicator-out/ so the rsync creates
    them remotely with the SSH user's ownership (uid 1000 = the container user)."""
    for sub in ("personas", "schemas", "out", "adjudicator-out"):
        os.makedirs(os.path.join(local_deck_dir, sub), exist_ok=True)
    with open(os.path.join(local_deck_dir, "BDEF.md"), "w") as f:
        f.write(bm_config.bdef_text(cfg))
    for fn in SCHEMA_FILES:
        shutil.copyfile(os.path.join(SCHEMAS_SRC, fn),
                        os.path.join(local_deck_dir, "schemas", fn))

    def persona(rid: str, text: str) -> None:
        with open(os.path.join(local_deck_dir, "personas", f"{rid}.md"), "w") as f:
            f.write((text or "").strip() + "\n")

    persona("extractor", prompts.extractor_persona())
    for r in panel:
        persona(r["rid"], r["persona"] or prompts.default_grader_persona(r["name"]))
    persona("adjudicator",
            (cfg.get("adjudicatorPersona") or "").strip() or prompts.adjudicator_persona())


# ------------------------------------------------------------------ containers
def container_suffix(deckdir: str) -> str:
    """Short, docker-safe name suffix from the deck dir (…/<company>/<deck_id>)."""
    parts = deckdir.rstrip("/").split("/")
    tail = "-".join(parts[-2:]) if len(parts) >= 2 else parts[-1]
    return slug(tail)[:48] or "deck"


def base_env(cfg: dict, rid: str, name: str, role: str, model: str, temperature) -> str:
    """The env-var block shared by every role container (also used by adjudicator.py)."""
    q = shlex.quote
    base = serving.reviewer_proxy_base(cfg)
    env = (
        f"-e BM_ROLE={q(role)} -e BM_GRADER_ID={q(rid)} -e BM_GRADER_NAME={q(name)} "
        f"-e BM_MODEL={q(model)} -e BM_LLM_BASE={q(base)} -e BM_LLM_KEY=sk-local "
        f"-e BM_MAX_MODEL_LEN={q(str(int(cfg.get('maxModelLen', 32768))))} "
        f"-e HOME=/home/rev "
    )
    if role == "extractor":
        env += "-e BM_TEMPERATURE=0.0 "
    elif temperature is not None:
        env += f"-e BM_TEMPERATURE={q(str(temperature))} "
    return env


def _container(cfg: dict, deckdir: str, rid: str, name: str, model: str,
               temperature, role: str) -> tuple[str, str]:
    """(container name, docker run command) for one extractor/grader container."""
    q = shlex.quote
    net = serving.net_name(cfg)
    schema_file = "extraction.schema.json" if role == "extractor" else "grades.schema.json"
    env = base_env(cfg, rid, name, role, model, temperature)
    mounts = (
        f"-v {q(deckdir)}/docs:/docs:ro "
        f"-v {q(deckdir)}/BDEF.md:/BDEF.md:ro "
        f"-v {q(deckdir)}/schemas/{schema_file}:/schema.json:ro "
        f"-v {q(deckdir)}/personas/{rid}.md:/persona/PERSONA.md:ro "
        f"-v {q(deckdir)}/out:/out "
    )
    cname = f"bm-grader-{rid}-{container_suffix(deckdir)}"
    cmd = (
        f"docker rm -f {cname} >/dev/null 2>&1; "
        f"docker run -d --name {cname} --network {q(net)} {HARDEN} {env} {mounts} "
        f"{q(cfg['graderImage'])}"
    )
    return cname, cmd


def _wait_for(cfg: dict, cname: str, out_file: str, log, wait_timeout: int) -> tuple[str, bool]:
    """docker-wait a launched container, check its output file, remove it."""
    head = sc.head(cfg)
    q = shlex.quote
    w = sc.run(head, f"docker wait {cname}", timeout=wait_timeout)
    code = (w.stdout or "").strip()
    chk = sc.run(head, f"test -s {q(out_file)} && echo OK || echo MISSING", timeout=30)
    wrote = "OK" in (chk.stdout or "")
    sc.run(head, f"docker rm -f {cname} 2>/dev/null; true", timeout=30)
    return code, wrote


# ------------------------------------------------------------------ runs
def run_extractor(cfg: dict, remote_deck_dir: str, model_alias: str, log,
                  wait_timeout: int = 1800) -> dict:
    """Launch the stage-1 extractor for one deck and wait for out/extraction.json."""
    head = sc.head(cfg)
    cname, cmd = _container(cfg, remote_deck_dir, "extractor", "extractor",
                            model_alias, None, role="extractor")
    r = sc.run(head, cmd, timeout=120)
    if r.returncode != 0:
        log(f"[graders] extractor launch FAILED: {r.stderr or r.stdout}")
        return {"rid": "extractor", "model": model_alias, "ok": False,
                "error": (r.stderr or r.stdout)[:300], "report": False}
    log(f"[graders] up: extractor -> {model_alias}")
    code, wrote = _wait_for(cfg, cname, f"{remote_deck_dir}/out/extraction.json",
                            log, wait_timeout)
    log(f"[graders] extractor exited (code={code or '?'}), "
        f"extraction.json={'written' if wrote else 'MISSING'}")
    return {"rid": "extractor", "model": model_alias, "ok": True, "error": "",
            "exit": code, "report": wrote}


def run_wave_graders(cfg: dict, remote_deck_dir: str, wave_panel: list[dict], log,
                     wait_timeout: int = 1800) -> list[dict]:
    """Launch every grader in `wave_panel` (already filtered to this wave's
    models) against one deck, wait for them, and report status. Grade JSONs land
    in <remote_deck_dir>/out/<rid>.json."""
    head = sc.head(cfg)
    launched = []
    for r in wave_panel:
        cname, cmd = _container(cfg, remote_deck_dir, r["rid"], r["name"], r["model"],
                                r.get("temperature"), role="grader")
        res = sc.run(head, cmd, timeout=120)
        if res.returncode != 0:
            log(f"[graders] launch {r['rid']} FAILED: {res.stderr or res.stdout}")
            launched.append({**r, "ok": False, "error": (res.stderr or res.stdout)[:300],
                             "cname": cname})
            continue
        log(f"[graders] up: {r['rid']} -> {r['model']}")
        launched.append({**r, "ok": True, "error": "", "cname": cname})

    # Wait for each launched container to exit (they run in parallel; waiting
    # sequentially still finishes when the slowest does).
    results = []
    for r in launched:
        if not r["ok"]:
            results.append({k: v for k, v in r.items() if k != "cname"})
            continue
        code, wrote = _wait_for(cfg, r["cname"], f"{remote_deck_dir}/out/{r['rid']}.json",
                                log, wait_timeout)
        log(f"[graders] {r['rid']} exited (code={code or '?'}), "
            f"grades={'written' if wrote else 'MISSING'}")
        results.append({k: v for k, v in r.items() if k != "cname"}
                       | {"exit": code, "report": wrote})
    return results