#!/usr/bin/env python3
"""Clean, greppable transcript logs for agent-orchestrator agents (claude backend).

The claude CLI writes a structured JSONL transcript of every session under
~/.claude/projects/<cwd-slug>/<session-uuid>.jsonl. This renders that into a readable,
greppable one-event-per-line log — WITHOUT touching the agent (read-only on a file it writes
anyway, so no slowdown and zero extra tokens). The raw `tmux pipe-pane` logs are TUI-escape
soup; use these instead.

Agents are discovered from the harness config (the same agents.toml the driver reads), so there
is nothing project-specific in this file. An agent's transcript directory is derived from its
working `dir`; when several agents share a dir, give them a `log_signature = "..."` (a substring
of their kickoff) to disambiguate.

Usage:
  agent-log.py [--config PATH] render <agent>     print the clean transcript to stdout
  agent-log.py [--config PATH] tail <agent> [N]   print the last N (default 40) rendered events
  agent-log.py [--config PATH] follow-all         keep <agent>.clean.log current for every agent

Output logs (follow-all): <log_dir>/<agent>.clean.log
Each line: `HH:MM:SS [kind] ...` — kinds: user, asst, tool:<Name>, result, think (skipped by
default). Long text/results are truncated (full detail stays in the JSONL); newlines → ⏎.
"""
import json, os, re, sys, time, tomllib
from pathlib import Path

PROJ = os.environ.get("CLAUDE_PROJECTS", os.path.expanduser("~/.claude/projects"))
MAXLEN = 800  # truncate any single rendered block to this many chars


def _cfg_path(argv):
    if "--config" in argv:
        return Path(argv[argv.index("--config") + 1])
    cwd_cfg = Path.cwd() / "agents.toml"
    return cwd_cfg if cwd_cfg.exists() else Path(__file__).resolve().parent / "agents.toml"


def load_agents(cfg_path):
    """Return {agent_name: {"slug": <claude proj dir slug>, "sig": <optional signature>}} and the
    log_dir, derived from the harness config."""
    with open(cfg_path, "rb") as f:
        raw = tomllib.load(f)
    defaults = raw.get("defaults", {})
    base = Path(cfg_path).resolve().parent
    proj_dir = (base / defaults.get("project_dir", ".")).resolve()
    log_dir = os.path.join(str((proj_dir / defaults.get("log_dir", ".ao-state"))), "")
    agents = {}
    for a in raw.get("agent", []):
        d = a.get("dir", defaults.get("dir", "."))
        dp = Path(os.path.expanduser(d))
        dp = dp if dp.is_absolute() else (proj_dir / dp)
        slug = re.sub(r"[^a-zA-Z0-9]", "-", str(dp.resolve()))
        agents[a["name"]] = {"slug": slug, "sig": a.get("log_signature")}
    return agents, str((proj_dir / defaults.get("log_dir", ".ao-state")))


def _first_user_text(path, limit=4000):
    try:
        with open(path) as f:
            for _ in range(60):
                ln = f.readline()
                if not ln:
                    break
                try:
                    o = json.loads(ln)
                except Exception:
                    continue
                if o.get("type") == "user":
                    c = (o.get("message") or {}).get("content")
                    if isinstance(c, str):
                        return c[:limit]
                    if isinstance(c, list):
                        return " ".join(b.get("text", "") for b in c
                                        if isinstance(b, dict) and b.get("type") == "text")[:limit]
    except FileNotFoundError:
        return ""
    return ""


def active_jsonl(meta):
    """The agent's current transcript: newest *.jsonl in its slug dir (optionally filtered by the
    kickoff signature, to disambiguate agents that share a working dir)."""
    d = os.path.join(PROJ, meta["slug"])
    try:
        files = [os.path.join(d, f) for f in os.listdir(d) if f.endswith(".jsonl")]
    except FileNotFoundError:
        return None
    files.sort(key=lambda p: os.path.getmtime(p), reverse=True)
    for p in files:
        if not meta["sig"] or meta["sig"] in _first_user_text(p):
            return p
    return None


def _clip(s):
    s = " ".join(str(s).split()) if s else ""
    return s if len(s) <= MAXLEN else s[:MAXLEN] + " …[+%d]" % (len(s) - MAXLEN)


def render_line(o, show_think=False):
    t = o.get("type")
    if t not in ("user", "assistant"):
        return []
    ts = (o.get("timestamp") or "")[11:19] or "--:--:--"
    m = o.get("message") or {}
    c = m.get("content")
    out = []
    if isinstance(c, str):
        if c.strip():
            out.append(f"{ts} [user] {_clip(c)}")
        return out
    if not isinstance(c, list):
        return out
    for b in c:
        if not isinstance(b, dict):
            continue
        bt = b.get("type")
        if bt == "text":
            txt = b.get("text", "")
            if txt.strip():
                out.append(f"{ts} [{'asst' if t=='assistant' else 'user'}] {_clip(txt)}")
        elif bt == "thinking":
            if show_think:
                out.append(f"{ts} [think] {_clip(b.get('thinking',''))}")
        elif bt == "tool_use":
            inp = b.get("input") or {}
            brief = (inp.get("command") or inp.get("file_path") or inp.get("path")
                     or inp.get("prompt") or json.dumps(inp)[:200])
            out.append(f"{ts} [tool:{b.get('name')}] {_clip(brief)}")
        elif bt == "tool_result":
            rc = b.get("content")
            if isinstance(rc, list):
                rc = " ".join(x.get("text", "") for x in rc if isinstance(x, dict))
            out.append(f"{ts} [result] {_clip(rc)}")
    return out


def render_file(path, show_think=False):
    lines = []
    with open(path) as f:
        for ln in f:
            try:
                o = json.loads(ln)
            except Exception:
                continue
            lines += render_line(o, show_think)
    return lines


def cmd_render(agents, agent, show_think=False):
    p = active_jsonl(agents[agent])
    if not p:
        print(f"(no transcript found for {agent})", file=sys.stderr); sys.exit(1)
    print(f"# {agent}  ←  {p}")
    for l in render_file(p, show_think):
        print(l)


def cmd_tail(agents, agent, n=40):
    p = active_jsonl(agents[agent])
    if not p:
        print(f"(no transcript found for {agent})", file=sys.stderr); sys.exit(1)
    for l in render_file(p)[-n:]:
        print(l)


def cmd_follow_all(agents, log_dir):
    os.makedirs(log_dir, exist_ok=True)
    state = {}  # agent -> (path, byte_offset)
    while True:
        for agent, meta in agents.items():
            p = active_jsonl(meta)
            if not p:
                continue
            prev_path, off = state.get(agent, (None, 0))
            if p != prev_path:
                off = 0
            try:
                size = os.path.getsize(p)
                if off > size:
                    off = 0
                with open(p) as f:
                    f.seek(off)
                    chunk = f.read()
                    new_off = f.tell()
            except FileNotFoundError:
                continue
            if chunk:
                rendered = []
                for ln in chunk.splitlines():
                    try:
                        o = json.loads(ln)
                    except Exception:
                        continue
                    rendered += render_line(o)
                if rendered:
                    with open(os.path.join(log_dir, f"{agent}.clean.log"), "a") as out:
                        out.write("\n".join(rendered) + "\n")
            state[agent] = (p, new_off)
        time.sleep(5)


def main():
    argv = sys.argv[1:]
    cfg_path = _cfg_path(argv)
    argv = [a for i, a in enumerate(argv)
            if a != "--config" and (i == 0 or argv[i-1] != "--config")]
    agents, log_dir = load_agents(cfg_path)
    cmd = argv[0] if argv else "follow-all"
    if cmd == "render":
        cmd_render(agents, argv[1], "--think" in argv)
    elif cmd == "tail":
        cmd_tail(agents, argv[1], int(argv[2]) if len(argv) > 2 and argv[2].isdigit() else 40)
    elif cmd == "follow-all":
        cmd_follow_all(agents, log_dir)
    else:
        print(__doc__); sys.exit(2)


if __name__ == "__main__":
    main()