#!/usr/bin/env python3
"""Live kanban worker E2E battery — real dispatcher + kernel worker + MLX.

Unlike ``kernel_live_agent_battery.py`` (foreground chat / planner tier), this
slice seeds kanban tasks and polls until workers reach a terminal state. Each
case runs under an isolated tenant so the board stays clean.

Usage:
  ~/.mugi/venv/bin/python scripts/kernel_live_worker_battery.py
  ~/.mugi/venv/bin/python scripts/kernel_live_worker_battery.py --cases worker_general_read_file
  ~/.mugi/venv/bin/python scripts/kernel_live_worker_battery.py --delay 15 --json /tmp/worker_e2e.json

Requires a listening mugi_backend (Swift app up). Reuses the app's MLX server —
never starts a second model.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
import uuid
from typing import Any

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

from utils.mcp_backend_client import (  # noqa: E402
    discover_backend,
    http_json,
    new_conversation,
    send_and_wait,
    _headers,
)
from scripts.kernel_live_worker_cases import WORKER_CASES  # noqa: E402
from scripts.kernel_live_worker_types import WorkerCase, WorkerCaseResult  # noqa: E402
from scripts.kernel_live_gate import (  # noqa: E402
    format_gate_summary_line,
    summarize_worker_results,
)

_TERMINAL = frozenset({"done", "blocked", "gave_up", "cancelled"})
_TASK_ID_RE = re.compile(r"kt_[a-f0-9]{12,16}", re.IGNORECASE)


def _kanban_get(conn, path: str, *, timeout: float = 30.0) -> tuple[int, Any]:
    return http_json("GET", f"{conn.base_url}{path}", _headers(conn), timeout=timeout)


def _kanban_post(conn, path: str, body: dict[str, Any] | None = None, *, timeout: float = 30.0) -> tuple[int, Any]:
    return http_json("POST", f"{conn.base_url}{path}", _headers(conn), body, timeout=timeout)


def kanban_create_task(
    conn,
    *,
    title: str,
    body: str,
    assignee: str,
    tenant: str,
    priority: int = 3,
    metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
    code, out = _kanban_post(
        conn,
        "/kanban/tasks",
        {
            "title": title,
            "body": body,
            "assignee": assignee,
            "tenant": tenant,
            "priority": priority,
            "metadata": metadata or {},
        },
    )
    if code != 200 or not isinstance(out, dict) or not out.get("id"):
        return {"error": f"POST /kanban/tasks failed ({code})", "body": out}
    return out


def kanban_create_batch(
    conn,
    *,
    items: list[dict[str, Any]],
    tenant: str,
) -> dict[str, Any]:
    code, out = _kanban_post(
        conn,
        "/kanban/tasks/batch",
        {"items": items, "tenant": tenant},
        timeout=60.0,
    )
    if code != 200:
        return {"error": f"POST /kanban/tasks/batch failed ({code})", "body": out}
    tasks = out if isinstance(out, list) else out.get("tasks") or out.get("items") or []
    ids = [str(t.get("id")) for t in tasks if isinstance(t, dict) and t.get("id")]
    if not ids:
        return {"error": "batch returned no task ids", "body": out}
    return {"task_ids": ids, "tasks": tasks}


def kanban_show_task(conn, task_id: str) -> dict[str, Any]:
    code, out = _kanban_get(conn, f"/kanban/tasks/{task_id}")
    if code != 200 or not isinstance(out, dict):
        return {"error": f"GET /kanban/tasks/{task_id} failed ({code})", "body": out}
    return out


def kanban_list_runs(conn, task_id: str, *, limit: int = 3) -> list[dict[str, Any]]:
    code, out = _kanban_get(conn, f"/kanban/tasks/{task_id}/runs?limit={limit}")
    if code != 200:
        return []
    if isinstance(out, list):
        return out
    return out.get("runs") or [] if isinstance(out, dict) else []


def kanban_nudge(conn) -> None:
    _kanban_post(conn, "/kanban/dispatcher/nudge", {})


def kanban_dispatcher_health(conn) -> dict[str, Any]:
    code, out = _kanban_get(conn, "/kanban/dispatcher/health", timeout=10.0)
    return out if code == 200 and isinstance(out, dict) else {}


def _trace_path(thread_id: str) -> str:
    safe = re.sub(r"[^A-Za-z0-9_\-]", "_", thread_id)
    return os.path.expanduser(f"~/.mugi/logs/traces/{safe}.jsonl")


def _load_trace_blob(thread_id: str) -> str:
    path = _trace_path(thread_id)
    if not os.path.isfile(path):
        return ""
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            return f.read()
    except OSError:
        return ""


def _task_ids_from_chat_trace(thread_id: str) -> list[str]:
    blob = _load_trace_blob(thread_id)
    return list(dict.fromkeys(_TASK_ID_RE.findall(blob)))


def _tenant_for_case(case_id: str) -> str:
    return f"worker_battery_{case_id}_{uuid.uuid4().hex[:8]}"


def _summarize_terminal(conn, task_id: str) -> dict[str, Any]:
    task = kanban_show_task(conn, task_id)
    if task.get("error"):
        return {"task_id": task_id, "error": task.get("error")}
    runs = kanban_list_runs(conn, task_id, limit=1)
    latest = runs[0] if runs else {}
    md = latest.get("metadata") if isinstance(latest.get("metadata"), dict) else {}
    summary = str(latest.get("summary") or "")
    return {
        "task_id": task_id,
        "status": task.get("status"),
        "failure_count": task.get("failure_count"),
        "run_outcome": latest.get("outcome"),
        "run_summary": summary,
        "typed_reason": latest.get("typed_reason") or md.get("failure_reason"),
        "error": latest.get("error"),
    }


def wait_for_tasks(
    conn,
    task_ids: list[str],
    *,
    timeout_sec: float,
    poll_sec: float,
) -> tuple[list[dict[str, Any]], bool]:
    """Poll until every task is terminal or timeout. Returns summaries + timed_out flag."""
    deadline = time.monotonic() + max(1.0, timeout_sec)
    last_nudge = 0.0
    while time.monotonic() < deadline:
        summaries = [_summarize_terminal(conn, tid) for tid in task_ids]
        if all(str(s.get("status") or "") in _TERMINAL for s in summaries):
            return summaries, False
        now = time.monotonic()
        if now - last_nudge >= 30.0:
            kanban_nudge(conn)
            last_nudge = now
        time.sleep(max(1.0, poll_sec))
    summaries = [_summarize_terminal(conn, tid) for tid in task_ids]
    timed_out = not all(str(s.get("status") or "") in _TERMINAL for s in summaries)
    return summaries, timed_out


def score_worker_case(
    case: WorkerCase,
    *,
    summaries: list[dict[str, Any]],
    timed_out: bool,
) -> list[str]:
    failures: list[str] = []
    if timed_out:
        failures.append("timed_out waiting for terminal task state")
    if not summaries:
        failures.append("no task summaries to score")
        return failures

    done_count = sum(1 for s in summaries if str(s.get("status") or "") == case.expected_status)
    if done_count < case.expect_tasks_done:
        statuses = [str(s.get("status") or "?") for s in summaries]
        failures.append(
            f"expected {case.expect_tasks_done} task(s) in {case.expected_status!r}, "
            f"got {done_count}; statuses={statuses}"
        )

    for s in summaries:
        st = str(s.get("status") or "")
        if st in case.forbidden_statuses:
            failures.append(f"task {s.get('task_id')} reached forbidden status {st!r}")
        if st == case.expected_status and case.expected_run_outcome:
            outcome = str(s.get("run_outcome") or "")
            if outcome != case.expected_run_outcome:
                failures.append(
                    f"task {s.get('task_id')} status={st!r} but run_outcome={outcome!r} "
                    f"(expected {case.expected_run_outcome!r})"
                )
        summary = str(s.get("run_summary") or "")
        if case.min_summary_chars and st == "done" and len(summary) < case.min_summary_chars:
            failures.append(
                f"task {s.get('task_id')} summary length {len(summary)} "
                f"< min {case.min_summary_chars}"
            )
        blob = summary.lower()
        for marker in case.require_markers_in_summary:
            if st == "done" and marker.lower() not in blob:
                failures.append(
                    f"task {s.get('task_id')} summary missing required marker {marker!r}"
                )
        if s.get("error") and st not in ("done", "blocked"):
            failures.append(f"task {s.get('task_id')} run error: {s.get('error')}")

    return failures


def seed_case(conn, case: WorkerCase, tenant: str) -> tuple[list[str], str, list[str]]:
    """Return (task_ids, chat_thread_id, seed_failures)."""
    failures: list[str] = []
    chat_tid = ""

    if case.seed == "http":
        created = kanban_create_task(
            conn,
            title=case.title,
            body=case.body,
            assignee=case.assignee,
            tenant=tenant,
            priority=case.priority,
            metadata=case.metadata,
        )
        if created.get("error"):
            failures.append(str(created.get("error")))
            return [], chat_tid, failures
        return [str(created["id"])], chat_tid, failures

    if case.seed == "batch":
        if not case.batch_items:
            failures.append("batch seed requires batch_items")
            return [], chat_tid, failures
        created = kanban_create_batch(conn, items=case.batch_items, tenant=tenant)
        if created.get("error"):
            failures.append(str(created.get("error")))
            return [], chat_tid, failures
        return list(created["task_ids"]), chat_tid, failures

    if case.seed == "chat":
        conv = new_conversation(conn)
        if conv.get("error"):
            failures.append(str(conv.get("error")))
            return [], chat_tid, failures
        chat_tid = str(conv["thread_id"])
        waited = send_and_wait(
            conn,
            text=case.chat_prompt,
            thread_id=chat_tid,
            timeout_sec=case.chat_timeout_sec,
            poll_sec=2.0,
            trace_path=_trace_path(chat_tid),
        )
        if waited.get("error"):
            failures.append(str(waited.get("error")))
            return [], chat_tid, failures
        if waited.get("timed_out"):
            failures.append("chat seed timed out")
        ids = _task_ids_from_chat_trace(chat_tid)
        if not ids:
            failures.append("chat seed did not produce a task id in trace")
        for req in case.chat_required_tools:
            blob = _load_trace_blob(chat_tid)
            if f'"tool_name": "{req}"' not in blob and f'"tool_name":"{req}"' not in blob:
                if req not in blob:
                    failures.append(f"chat seed missing required tool {req!r}")
        return ids[:1], chat_tid, failures

    failures.append(f"unknown seed mode {case.seed!r}")
    return [], chat_tid, failures


def run_worker_case(conn, case: WorkerCase) -> WorkerCaseResult:
    tenant = _tenant_for_case(case.id)
    t0 = time.monotonic()
    task_ids, chat_tid, seed_failures = seed_case(conn, case, tenant)
    if seed_failures or not task_ids:
        return WorkerCaseResult(
            case_id=case.id,
            passed=False,
            elapsed_sec=time.monotonic() - t0,
            tenant=tenant,
            task_ids=task_ids,
            failures=seed_failures or ["no tasks seeded"],
            chat_thread_id=chat_tid,
        )

    kanban_nudge(conn)
    summaries, timed_out = wait_for_tasks(
        conn,
        task_ids,
        timeout_sec=case.timeout_sec,
        poll_sec=case.poll_sec,
    )
    failures = seed_failures + score_worker_case(case, summaries=summaries, timed_out=timed_out)
    return WorkerCaseResult(
        case_id=case.id,
        passed=not failures,
        elapsed_sec=time.monotonic() - t0,
        tenant=tenant,
        task_ids=task_ids,
        terminal=summaries,
        failures=failures,
        chat_thread_id=chat_tid,
        timed_out=timed_out,
    )


_CONTRACT_FLAKE_MARKERS = (
    "contract_missing",
    "verifier_failed",
    "Kanban worker ended without a parseable kanban-result",
)


def _is_contract_flake(res: WorkerCaseResult) -> bool:
    """A failure caused by the model not emitting a kanban-result fence.

    This is model compliance variance, not an infrastructure regression.
    Worth one clean retry — genuine assertion failures are not retried.
    """
    if res.passed:
        return False
    for s in res.terminal:
        reason = str(s.get("typed_reason") or "")
        summary = str(s.get("run_summary") or "")
        blockers = s.get("metadata") or {}
        if isinstance(blockers, dict):
            blockers = blockers.get("blockers") or []
        blob = f"{reason} {summary} {blockers}"
        if any(m in blob for m in _CONTRACT_FLAKE_MARKERS):
            return True
    return False


def run_worker_case_with_retry(
    conn,
    case: WorkerCase,
    *,
    max_attempts: int = 2,
    cooldown_sec: float = 5.0,
    cleanup_fn=None,
    on_retry=None,
) -> WorkerCaseResult:
    """Run a worker case, retrying once on contract-missing model flake."""
    res = run_worker_case(conn, case)
    res.attempts = 1
    if max_attempts <= 1 or not _is_contract_flake(res):
        return res
    first_failures = list(res.failures)
    if cleanup_fn and res.tenant:
        cleanup_fn(conn, res.tenant)
    for attempt in range(2, max_attempts + 1):
        if on_retry is not None:
            on_retry(attempt)
        if cooldown_sec > 0:
            time.sleep(cooldown_sec)
        retry = run_worker_case(conn, case)
        retry.attempts = attempt
        retry.first_attempt_failures = first_failures
        if retry.passed:
            retry.flaky = True
            return retry
        if not _is_contract_flake(retry):
            retry.first_attempt_failures = first_failures
            return retry
        if cleanup_fn and retry.tenant:
            cleanup_fn(conn, retry.tenant)
        res = retry
    res.first_attempt_failures = first_failures
    return res


def cleanup_tenant(conn, tenant: str) -> int:
    """Cancel active tasks in tenant, then clear cancelled records."""
    code, tasks = _kanban_get(conn, f"/kanban/tasks?tenant={tenant}&limit=200")
    if code != 200 or not isinstance(tasks, list):
        return 0
    cancelled = 0
    for t in tasks:
        tid = t.get("id")
        st = t.get("status")
        if tid and st in ("todo", "ready", "in_progress", "blocked"):
            _kanban_post(conn, f"/kanban/tasks/{tid}/cancel", {"cascade": True})
            cancelled += 1
    for terminal in ("done", "cancelled"):
        _kanban_post(conn, "/kanban/tasks/clear", {"status": terminal, "tenant": tenant})
    return cancelled


def main() -> int:
    parser = argparse.ArgumentParser(description="Live kanban worker E2E battery")
    parser.add_argument("--cases", default="", help="Comma-separated case ids (default: all)")
    parser.add_argument(
        "--delay",
        type=float,
        default=10.0,
        help="Seconds between cases (default 10 — gentler on MLX/GPU).",
    )
    parser.add_argument("--json", default="", help="Write JSON report to path")
    parser.add_argument(
        "--cleanup",
        action="store_true",
        help="Cancel/clear worker_battery_* tenant tasks after each case.",
    )
    parser.add_argument(
        "--gate",
        action="store_true",
        help="Release-gate mode (same exit semantics as default; prints GATE summary).",
    )
    args = parser.parse_args()

    conn = discover_backend()
    if isinstance(conn, dict):
        print("BACKEND:", json.dumps(conn, indent=2), file=sys.stderr)
        return 2

    selected = list(WORKER_CASES)
    if args.cases.strip():
        want = {x.strip() for x in args.cases.split(",") if x.strip()}
        selected = [c for c in WORKER_CASES if c.id in want]
        if not selected:
            print("No matching worker cases", file=sys.stderr)
            return 2

    health = kanban_dispatcher_health(conn)
    inflight = (
        health.get("dispatcher", {}).get("inflight_workers", 0)
        if isinstance(health.get("dispatcher"), dict)
        else health.get("inflight_workers", 0)
    )
    print(
        f"backend pid={conn.pid} port={conn.port} cases={len(selected)} "
        f"delay={args.delay}s inflight_workers={inflight}\n"
    )
    if inflight:
        print("  warning: dispatcher already has inflight workers — board may not be idle\n")

    max_attempts = 2 if args.gate else 1
    results: list[WorkerCaseResult] = []
    for i, case in enumerate(selected, 1):
        print(f"[{i}/{len(selected)}] {case.id} (seed={case.seed}) …", flush=True)

        def _on_retry(attempt: int) -> None:
            print(f"    [retry] contract-missing flake; attempt {attempt}/{max_attempts}…")

        res = run_worker_case_with_retry(
            conn,
            case,
            max_attempts=max_attempts,
            cooldown_sec=args.delay,
            cleanup_fn=cleanup_tenant if args.cleanup else None,
            on_retry=_on_retry,
        )
        results.append(res)
        status = "PASS" if res.passed else "FAIL"
        if res.passed and res.flaky:
            status = f"PASS (flaky, {res.attempts} tries)"
        terminals = [
            f"{s.get('task_id')}:{s.get('status')}" for s in res.terminal
        ]
        print(f"  {status} ({res.elapsed_sec:.1f}s) tenant={res.tenant} tasks={terminals}")
        if res.failures:
            for f in res.failures:
                print(f"    - {f}")
        if res.passed and res.flaky and res.first_attempt_failures:
            print(f"    (first attempt: {'; '.join(res.first_attempt_failures)})")
        if args.cleanup and res.tenant:
            n = cleanup_tenant(conn, res.tenant)
            if n:
                print(f"    cleaned {n} active task(s) in tenant")
        print()
        if i < len(selected) and args.delay > 0:
            time.sleep(args.delay)

    passed = sum(1 for r in results if r.passed)
    flaky = [r for r in results if r.passed and r.flaky]
    failed = len(results) - passed
    gate_summary = summarize_worker_results(results) if args.gate else None
    if args.gate and gate_summary:
        print(format_gate_summary_line("worker", gate_summary))
    flaky_note = f" ({len(flaky)} flaky)" if flaky else ""
    print(f"SUMMARY: {passed}/{len(results)} passed, {failed} failed{flaky_note}")
    if flaky:
        print(f"  flaky (passed on retry): {', '.join(r.case_id for r in flaky)}")

    if args.json:
        report = {
            "backend": {"pid": conn.pid, "port": conn.port},
            "summary": {
                "passed": passed,
                "failed": failed,
                "total": len(results),
                "gate": gate_summary,
            },
            "results": [
                {
                    "case_id": r.case_id,
                    "passed": r.passed,
                    "flaky": r.flaky,
                    "attempts": r.attempts,
                    "elapsed_sec": round(r.elapsed_sec, 2),
                    "tenant": r.tenant,
                    "task_ids": r.task_ids,
                    "terminal": r.terminal,
                    "failures": r.failures,
                    "first_attempt_failures": r.first_attempt_failures,
                    "chat_thread_id": r.chat_thread_id,
                    "timed_out": r.timed_out,
                }
                for r in results
            ],
        }
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump(report, f, indent=2)
        print(f"Wrote {args.json}")

    exit_code = gate_summary["exit_code"] if gate_summary else (0 if failed == 0 else 1)
    return exit_code


if __name__ == "__main__":
    raise SystemExit(main())
