#!/usr/bin/env python3

"""Headless executor for dispatch tasks.

Reads a task JSON file, spawns the configured agent binary in headless mode,

captures output within caps, and writes:

  done/<id>.result.json  - structured result (summary, exit code, timing)

  done/<id>.log          - full stdout/stderr for audit

Honors the killswitch - refuses to start if KILLSWITCH file exists.

Usage:

  dispatch-exec <path-to-task-json>

"""

from __future__ import annotations

import json

import os

import pathlib

import shutil

import signal

import subprocess

import sys

import time

from datetime import datetime, timezone

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))

import dispatch_lib as d

def _locate_agent_bin() -> str:

    env = os.environ.get("DISPATCH_AGENT_BIN")

    if env and os.access(env, os.X_OK):

        return env

    found = shutil.which("claude")

    if found:

        return found

    return "claude"

AGENT_BIN = _locate_agent_bin()

SIDE = os.environ.get("DISPATCH_SIDE", "exec")

SYSTEM_PROMPT = (

    "You are running as a headless dispatch worker. A peer agent session "

    "sent you this task via signed dispatch. Execute it end-to-end and "

    "report results. Be terse. The task request follows."

)

def _other(node: str) -> str:

    return d.NODE_B if node == d.NODE_A else d.NODE_A

def killswitch_check(task_id: str) -> bool:

    tripped, reason = d.killswitch_tripped()

    if tripped:

        d.log_event(SIDE, "exec_killswitch", f"refused task: {reason[:80]}", task_id)

        return True

    return False

def processing_path(task: dict) -> pathlib.Path:

    receiver = task["to"]

    return d.ROOT / f"{_other(receiver)}-to-{receiver}" / "processing" / f"{task['id']}.json"

def run_task(task: dict) -> int:

    task_id = task["id"]

    d.fill_defaults(task)

    if killswitch_check(task_id):

        return 99

    timeout = int(task.get("timeout_s", 600))

    max_bytes = int(task.get("max_output_bytes", 2_000_000))

    dangerous = bool(task.get("require_dangerous", False))

    perm_mode = "bypassPermissions" if dangerous else "acceptEdits"

    sender = task["from"]

    done = d.ROOT / f"{sender}-to-{_other(sender)}" / "done"

    done.mkdir(parents=True, exist_ok=True)

    log_path = done / f"{task_id}.log"

    result_path = done / f"{task_id}.result.json"

    cmd = [

        AGENT_BIN, "-p", task["request"],

        "--permission-mode", perm_mode,

        "--output-format", "stream-json",

        "--include-partial-messages",

        "--verbose",

        "--append-system-prompt", SYSTEM_PROMPT,

    ]

    if dangerous:

        cmd.append("--dangerously-skip-permissions")

    d.log_event(SIDE, "task_exec_start", task["request"][:80], task_id,

                perm_mode=perm_mode, timeout_s=timeout)

    started = time.time()

    out_bytes = bytearray()

    exit_code = 0

    try:

        with log_path.open("wb") as logf:

            logf.write(f"# task {task_id}\n# started {datetime.now(timezone.utc).isoformat()}\n".encode())

            logf.write(f"# perm_mode={perm_mode} timeout={timeout}s max_bytes={max_bytes}\n".encode())

            logf.write(b"# ---\n")

            logf.flush()

            proc = subprocess.Popen(

                cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,

            )

            try:

                while True:

                    rc = proc.poll()

                    chunk = proc.stdout.read(8192) if proc.stdout else b""

                    if chunk:

                        if len(out_bytes) + len(chunk) <= max_bytes:

                            out_bytes.extend(chunk)

                        logf.write(chunk)

                        logf.flush()

                    if rc is not None and not chunk:

                        break

                    if time.time() - started > timeout:

                        proc.send_signal(signal.SIGTERM)

                        try:

                            proc.wait(timeout=10)

                        except subprocess.TimeoutExpired:

                            proc.kill()

                        raise TimeoutError(f"exceeded {timeout}s")

                exit_code = proc.returncode

            except TimeoutError as e:

                logf.write(f"\n# TIMEOUT: {e}\n".encode())

                exit_code = 124

    except Exception as e:

        exit_code = 127

        err_msg = f"{type(e).__name__}: {e}"

        d.log_event(SIDE, "task_exec_err", err_msg[:120], task_id)

        try:

            with log_path.open("ab") as lf:

                lf.write(f"\n# EXEC ERROR: {err_msg}\n".encode())

        except Exception:

            pass

    elapsed = round(time.time() - started, 2)

    status = "ok" if exit_code == 0 else ("timeout" if exit_code == 124 else "error")

    final_text = ""

    try:

        for line in out_bytes.decode("utf-8", errors="replace").splitlines():

            if not line.strip():

                continue

            try:

                msg = json.loads(line)

            except Exception:

                continue

            if msg.get("type") == "result" and msg.get("result"):

                final_text = msg["result"]

    except Exception:

        pass

    result = {

        "id": task_id,

        "answered_by": SIDE,

        "answered_at": datetime.now(timezone.utc).isoformat(),

        "request": task["request"],

        "status": status,

        "exit_code": exit_code,

        "elapsed_s": elapsed,

        "log_path": str(log_path),

        "final_text": final_text[:8000] if final_text else "",

        "output_tail": out_bytes[-8000:].decode("utf-8", errors="replace") if out_bytes else "",

        "output_bytes": len(out_bytes),

        "perm_mode": perm_mode,

    }

    tmp = result_path.with_suffix(".json.tmp")

    tmp.write_text(json.dumps(result, indent=2))

    os.replace(tmp, result_path)

    proc_path = processing_path(task)

    if proc_path.exists():

        proc_path.unlink()

    d.log_event(SIDE, "task_exec_end",

                f"{status} exit={exit_code} elapsed={elapsed}s", task_id,

                exit_code=exit_code, status=status)

    return exit_code

def main() -> int:

    if len(sys.argv) != 2:

        print("usage: dispatch-exec <task.json>", file=sys.stderr)

        return 2

    path = pathlib.Path(sys.argv[1])

    if not path.exists():

        print(f"ERR: {path} not found", file=sys.stderr)

        return 2

    task = json.loads(path.read_text())

    return run_task(task)

if __name__ == "__main__":

    raise SystemExit(main())