"""Workspace composition (controller → node), collection (node → controller) and artefact hashing.
The workspace never contains the oracle; it contains the public bundle, COMPLETION_SCHEMA.json, PROMPT.md, SYSTEM.md,
(A3/A4) PORTAL.md, (A4) tools/selfcheck.py."""
from __future__ import annotations
import hashlib, json, os, shutil, tempfile
from .config import REPO, NODE_HOME, NODE_PYTHON
from .arms import ARMS, system_prompt, USER_PROMPT
from . import launcher

PUBLIC_EXCLUDES = ("generate.py", "__pycache__", ".DS_Store")

def scenario_dir(task_id: str) -> str:
    return os.path.join(REPO, "benchmark", "scenarios", task_id)

def load_task(task_id: str) -> dict:
    return json.load(open(os.path.join(scenario_dir(task_id), "task.json"), encoding="utf-8"))

def compose_local(task_id: str, arm: str, run_id: str, portal_url: str | None, portal_user: dict | None) -> str:
    """Build the workspace in a temp dir on the controller and return its path."""
    tmp = tempfile.mkdtemp(prefix=f"wp42ws_{run_id}_")
    src = scenario_dir(task_id)
    for name in os.listdir(src):
        if name in PUBLIC_EXCLUDES: continue
        p = os.path.join(src, name)
        if name == "portal" and not ARMS[arm]["portal"]: continue
        (shutil.copytree if os.path.isdir(p) else shutil.copy2)(p, os.path.join(tmp, name))
    shutil.copy2(os.path.join(REPO, "benchmark", "COMPLETION_SCHEMA.json"), os.path.join(tmp, "COMPLETION_SCHEMA.json"))
    # The worker's task.json must not reveal the design (planted errors, injection, depth, seed): keep only what a worker legitimately needs.
    t = json.load(open(os.path.join(tmp, "task.json"), encoding="utf-8"))
    public = {k: t[k] for k in ("task_id", "occupation", "organisation", "title", "version", "required_deliverables", "applications_involved") if k in t}
    json.dump(public, open(os.path.join(tmp, "task.json"), "w"), indent=1)
    os.makedirs(os.path.join(tmp, "out"), exist_ok=True)
    open(os.path.join(tmp, "PROMPT.md"), "w").write(USER_PROMPT)
    open(os.path.join(tmp, "SYSTEM.md"), "w").write(system_prompt(arm))
    if ARMS[arm]["portal"]:
        u = portal_user or {"username": "analyst", "password": "analyst"}
        open(os.path.join(tmp, "PORTAL.md"), "w").write(f"# Internal enterprise portal\n\nURL: {portal_url}\nUsername: {u['username']}\nPassword: {u['password']}\n\nPages: /documents (library, downloads), /records/<entity> and /api/<entity> (JSON), /forms/<form_id> (HTML form; or POST JSON to the same URL). Health: /health. Log in first (POST /login with username/password; the session cookie is then required).\n")
    if ARMS[arm]["recovery"]:
        os.makedirs(os.path.join(tmp, "tools"), exist_ok=True)
        shutil.copy2(os.path.join(REPO, "benchmark", "tools", "selfcheck.py"), os.path.join(tmp, "tools", "selfcheck.py"))
    # Public bundle hash manifest (what the worker saw)
    open(os.path.join(tmp, "BUNDLE_MANIFEST.json"), "w").write(json.dumps(hash_tree(tmp, skip=("out",)), indent=1))
    return tmp

def hash_tree(root: str, skip=()) -> dict:
    out = {}
    for dp, dns, fns in os.walk(root):
        dns[:] = [d for d in dns if d not in skip]
        for fn in sorted(fns):
            p = os.path.join(dp, fn); h = hashlib.sha256()
            with open(p, "rb") as f:
                for ch in iter(lambda: f.read(1 << 20), b""): h.update(ch)
            out[os.path.relpath(p, root)] = {"sha256": h.hexdigest(), "bytes": os.path.getsize(p)}
    return out

def node_workspace(run_id: str) -> str:
    return f"{NODE_HOME}/runs/{run_id}"

def push(node: str, local_ws: str, run_id: str) -> str:
    ws = node_workspace(run_id)
    launcher.ssh(node, f"mkdir -p {ws}", check=True)
    launcher.rsync_to(node, local_ws, ws)
    return ws

def collect(node: str, run_id: str, dest_dir: str) -> dict:
    """Copy out/ back, hash it, and verify the inputs were not modified (compare with the manifest)."""
    ws = node_workspace(run_id)
    launcher.rsync_from(node, f"{ws}/out", os.path.join(dest_dir, "out"))
    n_remote = launcher.ssh(node, f"find {ws}/out -type f | wc -l").stdout.strip()
    n_local = sum(len(f) for _, _, f in os.walk(os.path.join(dest_dir, "out")))
    if n_remote.isdigit() and int(n_remote) != n_local:
        raise launcher.LaunchError(f"collection mismatch on {node}: {n_remote} remote files vs {n_local} local")
    launcher.rsync_from(node, f"{ws}/inputs", os.path.join(dest_dir, "inputs_after"))
    r = launcher.ssh(node, f"cat {ws}/BUNDLE_MANIFEST.json")
    try: manifest = json.loads(r.stdout)
    except Exception: manifest = {}  # noqa: BLE001
    after = hash_tree(os.path.join(dest_dir, "inputs_after"))
    changed = [k for k, v in manifest.items() if k.startswith("inputs/") and after.get(k[len("inputs/"):], {}).get("sha256") != v["sha256"]]
    arts = hash_tree(os.path.join(dest_dir, "out"))
    shutil.rmtree(os.path.join(dest_dir, "inputs_after"), ignore_errors=True)
    return {"artifacts": arts, "inputs_changed": changed, "manifest": manifest}

def reset(node: str, run_id: str) -> None:
    ws = node_workspace(run_id)
    launcher.ssh(node, f"rm -rf {ws}")

def start_portal(node: str, run_id: str, port: int) -> str:
    """Start the portal on the node (background, loopback) and return the URL after a health check."""
    ws = node_workspace(run_id)
    cmd = (f"cd {ws} && (PYTHONPATH={NODE_HOME}/tools nohup {NODE_PYTHON} -m wp42.portal --workspace {ws} --port {port} "
           f"> {ws}/.portal.log 2>&1 & echo $! > {ws}/.portal.pid); sleep 2; curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{port}/health")
    r = launcher.ssh(node, cmd, timeout=60)
    if "200" not in r.stdout: raise launcher.LaunchError(f"portal did not start on {node}:{port}: {r.stdout[-200:]} {r.stderr[-300:]}")
    return f"http://127.0.0.1:{port}"

def stop_portal(node: str, run_id: str) -> None:
    ws = node_workspace(run_id)
    launcher.ssh(node, f"test -f {ws}/.portal.pid && kill $(cat {ws}/.portal.pid) 2>/dev/null; true")

def sync_tools(node: str) -> None:
    """Install the controller package (for the portal only) under ~/wp42/tools on the node."""
    launcher.ssh(node, f"mkdir -p {NODE_HOME}/tools/wp42", check=True)
    launcher.rsync_to(node, os.path.join(REPO, "src", "wp42"), f"{NODE_HOME}/tools/wp42", excludes=("__pycache__",))