#!/usr/bin/env python3
"""Start and close Herdr agents for the HUD. herdr-monitor stays read-only."""
import json
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

TIMEOUT = 15

# Herdr's supported kinds (`herdr agent start --kind`), with display names.
KINDS = {
    "claude": "Claude Code", "codex": "Codex", "gemini": "Gemini CLI", "cursor": "Cursor",
    "opencode": "opencode", "copilot": "Copilot CLI", "pi": "pi", "omp": "omp", "grok": "Grok",
    "hermes": "Hermes", "amp": "Amp", "droid": "Droid", "cline": "Cline", "kimi": "Kimi",
    "kiro": "Kiro", "qwen": "Qwen Code", "kilo": "Kilo", "devin": "Devin", "agy": "Antigravity",
    "mastracode": "Mastra Code", "qodercli": "Qoder", "maki": "Maki",
}
EXECUTABLE = {"cursor": "cursor-agent"}


class AgentsError(RuntimeError):
    def __init__(self, message: str, code: str = ""):
        super().__init__(message)
        self.code = code


def herdr(*args: str, timeout: float = TIMEOUT) -> dict:
    executable = shutil.which("herdr")
    if not executable:
        raise AgentsError("Herdr is not installed or not on PATH")
    try:
        done = subprocess.run([executable, *args], capture_output=True, text=True, timeout=timeout)
    except subprocess.TimeoutExpired as error:
        raise AgentsError("Herdr did not respond in time") from error
    if done.returncode != 0:
        # Herdr prints {"error": {"code", "message"}} on stderr.
        try:
            error = json.loads(done.stderr).get("error") or {}
        except (json.JSONDecodeError, AttributeError):
            error = {}
        raise AgentsError(error.get("message") or done.stderr.strip() or f"herdr {args[0]} failed",
                          error.get("code", ""))
    return json.loads(done.stdout or "{}").get("result", {})


def installed_kinds() -> list[str]:
    # Agents run in Herdr's login shell, which may have PATH entries (mise, npm)
    # that the desktop shell lacks, so ask that shell.
    executables = [EXECUTABLE.get(kind, kind) for kind in KINDS]
    script = "for e in " + " ".join(executables) + '; do command -v "$e" >/dev/null 2>&1 && echo "$e"; done'
    try:
        done = subprocess.run([os.environ.get("SHELL", "/bin/bash"), "-lic", script],
                              capture_output=True, text=True, timeout=TIMEOUT, stdin=subprocess.DEVNULL)
        found = set(done.stdout.split())
    except subprocess.TimeoutExpired:
        found = {e for e in executables if shutil.which(e)}
    return [kind for kind in KINDS if EXECUTABLE.get(kind, kind) in found]


def options() -> dict:
    agents = herdr("agent", "list").get("agents", [])
    spaces = herdr("workspace", "list").get("workspaces", [])
    folders = {str(a.get("cwd")) for a in agents if a.get("cwd")}
    work = Path.home() / "Work"
    if work.is_dir():
        folders |= {str(p) for p in work.iterdir() if p.is_dir() and not p.name.startswith(".")}
    # A workspace's folder is where its first agent runs; used to pick a default.
    space_cwd = {}
    for agent in agents:
        space_cwd.setdefault(str(agent.get("workspace_id", "")), agent.get("cwd"))
    return {
        "kinds": [{"kind": kind, "label": KINDS[kind]} for kind in installed_kinds()],
        "folders": sorted(folders),
        "workspaces": [{"id": s.get("workspace_id"), "label": s.get("label") or s.get("workspace_id"),
                        "cwd": space_cwd.get(str(s.get("workspace_id")))} for s in spaces],
    }


def slug(text: str, fallback: str) -> str:
    # Herdr agent names: [a-z][a-z0-9_-]{0,31}, unique among live agents.
    base = re.sub(r"[^a-z0-9_-]+", "-", text.lower())
    base = re.sub(r"^[^a-z]+|-+$", "", base)[:24]
    return base or fallback


def start(request: dict) -> dict:
    kind = str(request.get("kind", ""))
    if kind not in installed_kinds():
        raise AgentsError(f"{KINDS.get(kind, kind)} isn't installed")
    home = Path.home().resolve()
    folder = Path(os.path.expanduser(str(request.get("cwd", "")).strip())).resolve()
    if not folder.is_dir():
        raise AgentsError("That folder doesn't exist")
    if folder != home and home not in folder.parents:
        raise AgentsError("Pick a folder in your home directory")
    name = str(request.get("name", "")).strip()[:60]
    prompt = str(request.get("prompt", "")).strip()
    workspace = str(request.get("workspace", "") or "")
    label = (name or folder.name)[:24]

    if workspace:
        created = herdr("tab", "create", "--workspace", workspace, "--cwd", str(folder), "--label", label)
    else:
        created = herdr("workspace", "create", "--cwd", str(folder), "--label", folder.name, "--no-focus")
        herdr("tab", "rename", created["tab"]["tab_id"], label)
    pane = created["root_pane"]["pane_id"]
    tab = created["tab"]["tab_id"]

    taken = {a.get("name") for a in herdr("agent", "list").get("agents", [])}
    base = slug(name or folder.name, kind)
    agent_name, n = base, 2
    while agent_name in taken:
        agent_name, n = f"{base[:28]}-{n}", n + 1

    # Claude shows --name as its terminal title and session name.
    args = ["--", "--name", name] if kind == "claude" and name else []
    time.sleep(0.8)  # let the shell reach its prompt
    try:
        herdr("agent", "start", agent_name, "--kind", kind, "--pane", pane, "--timeout", "45000", *args, timeout=50)
    except AgentsError as error:
        # Blocked at startup (a trust-this-folder question, say) still counts:
        # the HUD terminal can answer it. Otherwise close the empty tab.
        if error.code == "agent_not_ready":
            return {"pane_id": pane, "blocked": True}
        try:
            herdr("tab", "close", tab)
        except AgentsError:
            pass
        raise AgentsError(f"{KINDS[kind]} didn't start: {error}") from None
    if prompt:
        # Herdr reports the start a moment before it takes prompts for the name.
        for attempt in range(20):
            try:
                herdr("agent", "prompt", pane, prompt)
                break
            except AgentsError as error:
                if error.code != "agent_not_ready" or attempt == 19:
                    return {"pane_id": pane, "blocked": False, "prompt_error": str(error)}
                time.sleep(0.5)
    return {"pane_id": pane, "blocked": False}


def claude_attach_id(pane: str) -> str:
    info = herdr("pane", "process-info", "--pane", pane).get("process_info", {})
    for process in info.get("foreground_processes", []):
        argv = process.get("argv") or []
        if len(argv) >= 3 and os.path.basename(argv[0]) == "claude" and argv[1] == "attach":
            return argv[2]
    return ""


def close(pane: str) -> dict:
    # A `claude attach` pane only views a background session; closing the pane
    # would leave the session running, so stop it (its conversation is kept).
    session = ""
    try:
        session = claude_attach_id(pane)
    except AgentsError:
        pass
    if session:
        claude = shutil.which("claude") or str(Path.home() / ".local/share/mise/shims/claude")
        done = subprocess.run([claude, "stop", session], capture_output=True, text=True, timeout=TIMEOUT)
        if done.returncode != 0:
            raise AgentsError((done.stderr or done.stdout).strip().splitlines()[0] if (done.stderr or done.stdout).strip()
                              else "claude stop failed")
    herdr("pane", "close", pane)
    return {"closed": pane, "stopped_session": session}


def main(argv: list[str]) -> int:
    if argv[1:] == ["options"]:
        result = options()
    elif len(argv) == 3 and argv[1] == "start":
        result = start(json.loads(argv[2]))
    elif len(argv) == 3 and argv[1] == "close":
        result = close(argv[2])
    else:
        raise AgentsError("usage: herdr-agents options|start JSON|close PANE")
    print(json.dumps(result, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main(sys.argv))
    except (AgentsError, OSError, ValueError, KeyError, subprocess.TimeoutExpired) as error:
        print(str(error)[:4096], file=sys.stderr)
        raise SystemExit(1)