#!/usr/bin/env python3
"""Plan and launch difficulty-aware Codex or Claude Code workers through Herdr."""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
from pathlib import Path
from typing import Any


ALLOWED_EFFORTS = ("low", "medium", "high", "xhigh")
ROUTING_PATH = Path("/opt/data/workspace/coordinator/model-routing.json")
HERDR_BIN = Path("/opt/data/tools/bin/herdr")
CODEX_AUTH = Path("/opt/data/home/.codex/auth.json")


def _load_routes(path: Path) -> dict[str, Any]:
    """Load the non-secret route status generated by the model steward."""
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise RuntimeError(f"routing status unavailable: {path}") from error
    if not isinstance(value, dict) or not isinstance(value.get("routes"), dict):
        raise RuntimeError(f"routing status is malformed: {path}")
    return value


def _split_route(route: str) -> tuple[str, str]:
    """Split a provider/model route without corrupting model punctuation."""
    provider, separator, model = route.partition("/")
    if not separator or not provider or not model:
        raise RuntimeError(f"invalid route: {route}")
    return provider, model


def select_plan(
    status: dict[str, Any], shape: str, effort: str, provider: str | None = None
) -> dict[str, Any]:
    """Select a capped route and retain the complete capacity fallback chain."""
    if effort not in ALLOWED_EFFORTS:
        raise ValueError(f"effort must be one of: {', '.join(ALLOWED_EFFORTS)}")
    default_provider = "codex" if shape == "implementation" else "claude"
    selected_provider = provider or default_provider
    if selected_provider not in {"codex", "claude"}:
        raise ValueError("provider must be codex or claude")
    profile = f"{selected_provider}-{effort}"
    chain = status["routes"].get(profile)
    if not isinstance(chain, list) or not chain:
        raise RuntimeError(f"route profile unavailable: {profile}")
    primary_provider, model = _split_route(str(chain[0]))
    return {
        "shape": shape,
        "effort": effort,
        "profile": profile,
        "worker": selected_provider,
        "provider": primary_provider,
        "model": model,
        "fallback_chain": [str(route) for route in chain[1:]],
    }


def _slug(value: str, limit: int = 32) -> str:
    """Return a Herdr-safe stable label."""
    cleaned = re.sub(r"[^a-z0-9_-]+", "-", value.lower()).strip("-")
    if not cleaned or not cleaned[0].isalpha():
        cleaned = f"task-{cleaned}"
    return cleaned[:limit].rstrip("-")


def _run(command: list[str], env: dict[str, str]) -> dict[str, Any]:
    """Run one Herdr JSON command and surface a concise error."""
    completed = subprocess.run(
        command,
        env=env,
        text=True,
        capture_output=True,
        timeout=320,
        check=False,
    )
    if completed.returncode != 0:
        detail = completed.stderr.strip() or completed.stdout.strip()
        raise RuntimeError(detail or f"command failed with {completed.returncode}")
    try:
        value = json.loads(completed.stdout)
    except json.JSONDecodeError as error:
        raise RuntimeError("Herdr returned non-JSON output") from error
    if not isinstance(value, dict):
        raise RuntimeError("Herdr returned an unexpected response")
    return value


def launch_worker(
    plan: dict[str, Any], project: Path, task: str, prompt: str | None
) -> dict[str, Any]:
    """Create an isolated Herdr workspace and launch the selected worker."""
    if not HERDR_BIN.is_file():
        raise RuntimeError("Herdr is not installed in the agent tools volume")
    if not project.is_dir():
        raise RuntimeError(f"project workspace does not exist: {project}")
    if plan["worker"] == "codex" and not CODEX_AUTH.is_file():
        raise RuntimeError(
            "Codex CLI needs its one-time device login. Ask Hermes to run "
            "`codex login --device-auth`, complete the code, then retry."
        )

    env = os.environ.copy()
    env.update(
        {
            "HOME": "/opt/data/home",
            "CODEX_HOME": "/opt/data/home/.codex",
            "CLAUDE_CONFIG_DIR": "/opt/data/home/.claude",
            "HERDR_CONFIG_PATH": "/opt/data/home/.config/herdr/config.toml",
            "HERDR_SOCKET_PATH": "/opt/data/herdr/herdr.sock",
            "PATH": "/opt/data/tools/bin:" + env.get("PATH", ""),
        }
    )
    label = _slug(task)
    created = _run(
        [
            str(HERDR_BIN),
            "workspace",
            "create",
            "--cwd",
            str(project),
            "--label",
            label,
            "--no-focus",
        ],
        env,
    )
    try:
        pane = str(created["result"]["root_pane"]["pane_id"])
    except (KeyError, TypeError) as error:
        raise RuntimeError("Herdr workspace response omitted the root pane") from error

    agent_name = _slug(f"{plan['worker']}-{task}")
    command = [
        str(HERDR_BIN),
        "agent",
        "start",
        agent_name,
        "--kind",
        plan["worker"],
        "--pane",
        pane,
        "--timeout",
        "120000",
        "--",
    ]
    if plan["worker"] == "codex":
        command.extend(
            [
                "-m",
                plan["model"],
                "-c",
                f'model_reasoning_effort="{plan["effort"]}"',
                "-c",
                'approval_policy="on-request"',
                "-c",
                'sandbox_mode="workspace-write"',
            ]
        )
    else:
        command.extend(
            [
                "--model",
                plan["model"],
                "--effort",
                plan["effort"],
                "--permission-mode",
                "acceptEdits",
            ]
        )
    started = _run(command, env)
    result = {**plan, "agent": agent_name, "pane": pane, "started": started}
    if prompt:
        result["prompted"] = _run(
            [str(HERDR_BIN), "agent", "prompt", agent_name, prompt], env
        )
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--shape", choices=("implementation", "architecture", "review"), required=True)
    parser.add_argument("--effort", choices=ALLOWED_EFFORTS, required=True)
    parser.add_argument("--provider", choices=("codex", "claude"))
    parser.add_argument("--routes", type=Path, default=ROUTING_PATH)
    parser.add_argument("--start", action="store_true")
    parser.add_argument("--project", type=Path)
    parser.add_argument("--task", default="objective")
    parser.add_argument("--prompt")
    args = parser.parse_args()

    plan = select_plan(_load_routes(args.routes), args.shape, args.effort, args.provider)
    if args.start:
        if args.project is None:
            parser.error("--project is required with --start")
        plan = launch_worker(plan, args.project.resolve(), args.task, args.prompt)
    print(json.dumps(plan, indent=2, sort_keys=True))
    return 0


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