#!/usr/bin/env python3
"""
claude_code_tracer.py — OpenInference tracer for Claude Code hooks.

Ships Claude Code activities as OpenInference traces to Arthur Engine.

Trace model:
  - Each user prompt → one trace (CHAIN root + TOOL/RETRIEVER/AGENT children + LLM children)
  - Exiting the session completes the current in-progress trace
  - Traces are linked by arthur.session attribute

Hook events:
  user_prompt_submit — start a new trace; complete previous trace if in progress
  pre_tool           — record current tool start (fallback trace creation if UserPromptSubmit unavailable)
  post_tool          — send a TOOL/RETRIEVER/AGENT span for the completed tool call
  post_tool_failure  — send an error TOOL/RETRIEVER/AGENT span for a failed tool call
  stop               — complete the current trace and clean up session state

Config priority (highest first):
  1. Env vars: GENAI_ENGINE_API_KEY, GENAI_ENGINE_TASK_ID, GENAI_ENGINE_TRACE_ENDPOINT
  2. Project config: $CLAUDE_PROJECT_DIR/.claude/arthur_config.json
  3. Global config: ~/.claude/arthur_config.json
  4. Silent no-op if nothing configured
"""

import contextlib
import fcntl
import json
import logging
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

# ---------------------------------------------------------------------------
# Logging — stderr only so stdout stays clean for Claude hooks
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.WARNING,
    format="[claude_code_tracer] %(levelname)s: %(message)s",
    stream=sys.stderr,
)
log = logging.getLogger("claude_code_tracer")


# ---------------------------------------------------------------------------
# Config discovery
# ---------------------------------------------------------------------------


def _load_config_file(path: Path) -> dict:
    try:
        if path.exists():
            return json.loads(path.read_text())
    except Exception as e:
        log.debug("Failed to read config %s: %s", path, e)
    return {}


def discover_config() -> Optional[dict]:
    """Returns config dict with keys: api_key, task_id, endpoint.
    Returns None if not configured (silent no-op)."""
    api_key = os.environ.get("GENAI_ENGINE_API_KEY", "")
    task_id = os.environ.get("GENAI_ENGINE_TASK_ID", "")
    endpoint = os.environ.get("GENAI_ENGINE_TRACE_ENDPOINT", "")

    if not (api_key and task_id and endpoint):
        project_dir = os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd())
        project_cfg = _load_config_file(
            Path(project_dir) / ".claude" / "arthur_config.json",
        )
        api_key = api_key or project_cfg.get("api_key", "")
        task_id = task_id or project_cfg.get("task_id", "")
        endpoint = endpoint or project_cfg.get("endpoint", "")

    if not (api_key and task_id and endpoint):
        global_cfg = _load_config_file(Path.home() / ".claude" / "arthur_config.json")
        api_key = api_key or global_cfg.get("api_key", "")
        task_id = task_id or global_cfg.get("task_id", "")
        endpoint = endpoint or global_cfg.get("endpoint", "")

    if not (api_key and task_id and endpoint):
        return None

    return {"api_key": api_key, "task_id": task_id, "endpoint": endpoint}


# ---------------------------------------------------------------------------
# State file helpers
# ---------------------------------------------------------------------------

STATE_DIR = Path.home() / ".claude" / "tracer"
STATE_MAX_AGE_S = 48 * 3600


def _state_path(session_id: str) -> Path:
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    path = (STATE_DIR / f"{session_id}.json").resolve()
    if not path.is_relative_to(STATE_DIR.resolve()):
        raise ValueError(f"Invalid session_id: {session_id!r}")
    return path


def _load_state(session_id: str) -> dict:
    path = _state_path(session_id)
    try:
        if path.exists():
            return json.loads(path.read_text())
    except Exception as e:
        log.debug("Failed to load state for %s: %s", session_id, e)
    return {}


def _save_state(session_id: str, state: dict) -> None:
    try:
        _state_path(session_id).write_text(json.dumps(state))
    except Exception as e:
        log.warning("Failed to save state for %s: %s", session_id, e)


def _delete_state(session_id: str) -> None:
    try:
        p = _state_path(session_id)
        if p.exists():
            p.unlink()
    except Exception as e:
        log.debug("Failed to delete state for %s: %s", session_id, e)


def _cleanup_stale_states() -> None:
    try:
        now = time.time()
        for p in STATE_DIR.glob("*.json"):
            if now - p.stat().st_mtime > STATE_MAX_AGE_S:
                p.unlink()
    except Exception as e:
        log.debug("Stale state cleanup failed: %s", e)


def _new_trace_id() -> str:
    import secrets

    return secrets.token_hex(16)


def _new_span_id() -> str:
    import secrets

    return secrets.token_hex(8)


# ---------------------------------------------------------------------------
# Subagent trace-context propagation
# ---------------------------------------------------------------------------

# How long to wait for a subagent to claim the context file before giving up.
_PENDING_AGENT_TIMEOUT_S = 30
_PENDING_AGENT_DIR = STATE_DIR / "pending_agent"


def _write_pending_agent_context(
    parent_session_id: str,
    parent_trace_id: str,
    parent_span_id: str,
    agent_prompt: str = "",
) -> None:
    """Write a side-channel file so the next subagent can inherit this trace.

    Called from handle_pre_tool when an Agent tool invocation is detected.
    The file is keyed by session+span so parallel agent invocations in the
    same session each get their own file.  ``agent_prompt`` is stored so that
    ``_claim_pending_agent_context`` can match the exact context for a given
    subagent when multiple agents are running in parallel.
    """
    _PENDING_AGENT_DIR.mkdir(parents=True, exist_ok=True)
    path = _PENDING_AGENT_DIR / f"{parent_session_id}_{parent_span_id}.json"
    path.write_text(
        json.dumps(
            {
                "parent_session_id": parent_session_id,
                "parent_trace_id": parent_trace_id,
                "parent_span_id": parent_span_id,
                "created_ns": time.time_ns(),
                "agent_prompt": agent_prompt,
            },
        ),
    )


def _claim_pending_agent_context(prompt: str = "") -> Optional[dict]:
    """Claim the pending subagent context that matches this session, if any.

    Called from handle_user_prompt_submit when a new session is initialised.
    Returns the context dict (with ``parent_trace_id`` and ``parent_span_id``)
    or None.  Claiming removes the file so no other session can use it.

    When ``prompt`` is provided (the subagent's own prompt text) we first try
    to find an exact match against ``agent_prompt`` stored in each context
    file.  This prevents parallel agent invocations from claiming each other's
    contexts.  If no exact match is found we fall back to newest-first so that
    contexts written by older tracer versions (which lack ``agent_prompt``) are
    still claimed.
    """
    if not _PENDING_AGENT_DIR.exists():
        return None

    deadline_ns = time.time_ns() - _PENDING_AGENT_TIMEOUT_S * 1_000_000_000
    candidates: list[tuple[int, Path, dict]] = []
    for p in _PENDING_AGENT_DIR.glob("*.json"):
        try:
            ctx = json.loads(p.read_text())
            if ctx.get("created_ns", 0) >= deadline_ns:
                candidates.append((ctx["created_ns"], p, ctx))
        except Exception:
            continue

    if not candidates:
        return None

    # Sort newest-first once; used by both the prompt-match pass and fallback.
    candidates.sort(key=lambda x: x[0], reverse=True)

    # Pass 1: exact prompt match — eliminates parallel-agent races.
    if prompt:
        for _, path, ctx in candidates:
            if ctx.get("agent_prompt", "") == prompt:
                try:
                    path.unlink()
                    return ctx
                except FileNotFoundError:
                    # Another process claimed it first; keep scanning.
                    continue

    # Pass 2: fallback — newest unclaimed (handles legacy files without agent_prompt).
    for _, path, ctx in candidates:
        try:
            path.unlink()
            return ctx
        except FileNotFoundError:
            continue
    return None


@contextlib.contextmanager
def _session_lock(session_id: str):
    """Exclusive per-session file lock.

    Serialises concurrent handle_post_tool / handle_post_tool_failure
    processes so that _emit_pending_llm_spans reads a consistent
    emitted_llm_span_count and never emits the same LLM span twice.
    This is necessary because Claude Code can run multiple tools in parallel,
    which causes multiple PostToolUse hook processes to fire concurrently.
    """
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    lock_path = (STATE_DIR / f"{session_id}.lock").resolve()
    if not lock_path.is_relative_to(STATE_DIR.resolve()):
        raise ValueError(f"Invalid session_id: {session_id!r}")

    with open(lock_path, "w") as fh:
        fcntl.flock(fh, fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(fh, fcntl.LOCK_UN)


# ---------------------------------------------------------------------------
# Transcript helpers
# ---------------------------------------------------------------------------


def _is_real_transcript(path: str) -> bool:
    """Return True if the transcript has at least one human or assistant entry.

    Shadow transcripts created by ``gh pr create`` in a worktree contain only
    ``pr-link`` metadata entries and no actual conversation content.  Accepting
    one of these as the authoritative transcript would leave the tracer with
    nothing to extract LLM spans from.
    """
    try:
        with open(path) as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if entry.get("type") in ("user", "assistant"):
                    return True
        return False
    except OSError:
        return False


def _find_transcript_path(data: dict, session_id: str) -> Optional[str]:
    """Locate the session's transcript JSONL file.

    Searches three locations in priority order and returns the first match
    that contains real conversation content (at least one user/assistant entry).
    Callers that need a stable path across the lifetime of a session should
    cache the result in session state via ``_get_cached_transcript_path``
    rather than calling this function repeatedly — the ``transcript_path``
    value in hook ``data`` can change mid-session when Claude Code writes
    entries to a *different* project directory (e.g. after ``gh pr create``
    in a worktree context).
    """
    candidates: list[str] = []

    # 1. Provided directly in hook data
    tp = data.get("transcript_path", "")
    if tp and Path(tp).exists():
        candidates.append(tp)

    # 2. Construct from CLAUDE_PROJECT_DIR (path with / → -)
    project_dir = os.environ.get("CLAUDE_PROJECT_DIR", "")
    if project_dir:
        sanitized = project_dir.replace("/", "-")
        path = Path.home() / ".claude" / "projects" / sanitized / f"{session_id}.jsonl"
        if path.exists():
            candidates.append(str(path))

    # 3. Glob fallback across all project dirs (sorted largest-first so the
    #    real transcript wins over tiny shadow files)
    projects_dir = Path.home() / ".claude" / "projects"
    if projects_dir.exists():
        matches = list(projects_dir.glob(f"*/{session_id}.jsonl"))
        matches.sort(key=lambda p: p.stat().st_size, reverse=True)
        candidates.extend(str(m) for m in matches)

    # Return the first candidate that contains actual conversation entries
    seen: set[str] = set()
    for candidate in candidates:
        if candidate in seen:
            continue
        seen.add(candidate)
        if _is_real_transcript(candidate):
            return candidate

    return None


def _get_cached_transcript_path(
    data: dict,
    state: dict,
    session_id: str,
) -> Optional[str]:
    """Return the transcript path for this session, caching it in *state*.

    Once resolved, the path is stored under ``state["transcript_path"]`` so
    that all subsequent hook calls use the same file even if Claude Code later
    writes a new entry to a different project directory (e.g. a worktree dir
    after ``gh pr create``).  The cache is only accepted when the file still
    exists; if it has been deleted the function re-resolves.
    """
    cached = state.get("transcript_path", "")
    if cached and Path(cached).exists():
        return cached

    resolved = _find_transcript_path(data, session_id)
    if resolved:
        state["transcript_path"] = resolved
    return resolved


def _is_human_message(entry: dict) -> bool:
    """True for user-initiated messages (not tool results)."""
    if entry.get("type") != "user":
        return False
    content = entry.get("message", {}).get("content", "")
    # Tool result messages have content as a list; human messages have a string
    return isinstance(content, str)


def _count_human_messages(transcript_path: str) -> int:
    """Count human (non-tool-result) user messages to detect turn boundaries."""
    try:
        count = 0
        for line in Path(transcript_path).read_text().splitlines():
            if not line.strip():
                continue
            try:
                if _is_human_message(json.loads(line)):
                    count += 1
            except json.JSONDecodeError:
                pass
        return count
    except Exception:
        return 0


def _iso_to_ns(ts: str) -> int:
    """Convert ISO 8601 timestamp string to nanoseconds."""
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        return int(dt.timestamp() * 1_000_000_000)
    except Exception:
        return time.time_ns()


def _tool_result_text(content: list) -> str:
    """Extract readable text from a tool_result content list."""
    parts = []
    for item in content:
        if not isinstance(item, dict):
            continue
        if item.get("type") == "tool_result":
            inner = item.get("content", "")
            if isinstance(inner, str):
                parts.append(inner)
            elif isinstance(inner, list):
                for rc in inner:
                    if isinstance(rc, dict) and rc.get("type") == "text":
                        parts.append(rc.get("text", ""))
    return "\n".join(parts)


def _extract_llm_spans_for_turn(
    transcript_path: str,
    human_count_at_start: int,
    trace_id_hex: str,
    root_span_id_hex: str,
) -> list[dict]:
    """
    Extract LLM spans from transcript entries that belong to this turn.

    A turn starts after the `human_count_at_start`-th human message and ends
    at the next human message (or end of transcript).  Scanning stops as soon
    as a second human message is seen while in_turn is True — those subsequent
    turns belong to their own traces.
    Uses actual timestamps from the transcript.

    For each LLM call we use the immediately-preceding message as the input:
      - First call in turn  → the human prompt (text/plain)
      - Subsequent calls    → the tool result(s) that preceded them (application/json)

    Output is the assistant's text response when present; otherwise the
    tool_use blocks serialised as JSON so the span always has an output value.

    **Grouping by message.id**

    Claude Code (with extended thinking enabled) splits a single API response
    across multiple consecutive transcript entries that share the same
    ``message.id`` — one entry per block type (thinking / text / tool_use).
    Each entry duplicates the input-side token counts.  We MUST group these
    entries and emit exactly one LLM span per API response, otherwise:

      * Prompt tokens are reported N× too high (once per entry).
      * "Thinking-only" entries produce empty, noisy spans.
      * Intermediate text (the entry with ``type=text`` that precedes a tool
        call) becomes its own orphaned 8-token span instead of being part of
        the single span for that API call.

    Grouping strategy:
      - Input tokens  → first entry in the group (avoids duplication).
      - Output tokens → sum across all entries (each entry captures the tokens
                        for its own block).
      - Text output   → concatenate all ``text`` blocks from any entry.
      - Tool output   → JSON of tool_use blocks from the last entry (fall-back
                        when no text is present).
      - Timestamp     → first entry.
    """
    spans = []
    try:
        lines = Path(transcript_path).read_text().splitlines()
        human_count = 0
        in_turn = False

        # Tracks the input for the *next* LLM call we encounter
        last_input_value = ""
        last_input_mime = "text/plain"
        last_input_role = "user"
        last_input_content = ""
        last_input_tool_call_id = ""

        # Accumulator for the current message.id group
        current_group_id: Optional[str] = None
        group_first_usage: dict = {}
        group_total_output_tokens: int = 0
        group_text_parts: list = []
        group_tool_use_parts: list = []
        group_model: str = "claude"
        group_ts: str = ""
        group_stop_reason: str = ""
        group_input_snapshot: dict = {}  # last_input* captured at group start

        def _flush_group() -> None:
            """Emit one LLM span for the accumulated group, if non-empty."""
            nonlocal current_group_id, group_first_usage, group_total_output_tokens
            nonlocal group_text_parts, group_tool_use_parts, group_model, group_ts
            nonlocal group_input_snapshot, group_stop_reason

            if not current_group_id:
                return

            usage = group_first_usage
            input_tokens = usage.get("input_tokens", 0)
            cache_read = usage.get("cache_read_input_tokens", 0)
            cache_create = usage.get("cache_creation_input_tokens", 0)
            output_tokens = group_total_output_tokens

            if group_text_parts:
                output_value = _truncate("".join(group_text_parts))
                output_mime = "text/plain"
            elif group_tool_use_parts:
                output_value = _truncate(json.dumps(group_tool_use_parts))
                output_mime = "application/json"
            else:
                output_value = ""
                output_mime = "text/plain"

            # Per OpenInference spec, message.content carries text only; tool calls
            # go in structured tool_calls attributes (not serialised into content).
            output_message_content = (
                _truncate("".join(group_text_parts)) if group_text_parts else ""
            )

            start_ns = _iso_to_ns(group_ts)
            end_ns = start_ns + max(output_tokens * 10_000_000, 1_000_000)

            snap = group_input_snapshot
            attrs: dict[str, Any] = {
                "openinference.span.kind": "LLM",
                "llm.system": "anthropic",
                "llm.model_name": group_model,
                "llm.token_count.prompt": input_tokens + cache_read + cache_create,
                "llm.token_count.completion": output_tokens,
                "llm.token_count.total": input_tokens
                + cache_read
                + cache_create
                + output_tokens,
                "llm.token_count.prompt_details.cache_read": cache_read,
                "llm.token_count.prompt_details.cache_write": cache_create,
                "llm.input_messages.0.message.role": snap.get("role", "user"),
                "llm.input_messages.0.message.content": snap.get("content", ""),
                "input.value": snap.get("value", ""),
                "input.mime_type": snap.get("mime", "text/plain"),
                "llm.output_messages.0.message.role": "assistant",
                "llm.output_messages.0.message.content": output_message_content,
                "output.value": output_value,
                "output.mime_type": output_mime,
            }

            # Structured tool_calls attributes (OpenInference spec)
            for i, tc in enumerate(group_tool_use_parts):
                attrs[f"llm.output_messages.0.message.tool_calls.{i}.tool_call.id"] = (
                    tc.get("id", "")
                )
                attrs[
                    f"llm.output_messages.0.message.tool_calls.{i}.tool_call.function.name"
                ] = tc.get("name", "")
                attrs[
                    f"llm.output_messages.0.message.tool_calls.{i}.tool_call.function.arguments"
                ] = _truncate(json.dumps(tc.get("input", {})))

            # tool_call_id on input message when input is a tool result
            if snap.get("tool_call_id"):
                attrs["llm.input_messages.0.message.tool_call_id"] = snap[
                    "tool_call_id"
                ]

            # Map stop_reason to OpenInference llm.finish_reason
            _FINISH_REASON_MAP = {"end_turn": "stop", "max_tokens": "length"}
            finish_reason = _FINISH_REASON_MAP.get(group_stop_reason, group_stop_reason)
            if finish_reason:
                attrs["llm.finish_reason"] = finish_reason

            spans.append(
                {
                    "trace_id_hex": trace_id_hex,
                    "span_id_hex": _new_span_id(),
                    "parent_span_id_hex": root_span_id_hex,
                    "name": f"claude/{group_model}",
                    "kind": None,
                    "start_ns": start_ns,
                    "end_ns": end_ns,
                    "attributes": attrs,
                    "force_span_id": False,
                },
            )

            # Reset accumulator
            current_group_id = None
            group_first_usage = {}
            group_total_output_tokens = 0
            group_text_parts = []
            group_tool_use_parts = []
            group_model = "claude"
            group_ts = ""
            group_stop_reason = ""
            group_input_snapshot = {}

        for line in lines:
            if not line.strip():
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError:
                continue

            entry_type = entry.get("type", "")

            # ── Human message: marks the start of this turn ──────────────────
            if _is_human_message(entry):
                _flush_group()
                human_count += 1
                if in_turn:
                    # Next user turn has started — stop here.  Its spans belong
                    # to a different trace.
                    break
                if human_count > human_count_at_start:
                    in_turn = True
                    human_text = entry.get("message", {}).get("content", "")
                    last_input_value = json.dumps(
                        {"role": "user", "content": human_text[:500]},
                    )
                    last_input_mime = "application/json"
                    last_input_role = "user"
                    last_input_content = _truncate(human_text)
                    last_input_tool_call_id = ""
                continue

            if not in_turn:
                continue

            # ── Tool result (user message with list content) ──────────────────
            if entry_type == "user":
                _flush_group()
                content = entry.get("message", {}).get("content", "")
                if isinstance(content, list):
                    text = _tool_result_text(content)
                    payload = text if text else json.dumps(content)
                    # Extract tool_use_id from first tool_result for message.tool_call_id
                    tool_use_id = ""
                    for item in content:
                        if isinstance(item, dict) and item.get("type") == "tool_result":
                            tool_use_id = item.get("tool_use_id", "")
                            break
                    last_input_value = json.dumps(
                        {"role": "tool", "content": payload[:500]},
                    )
                    last_input_mime = "application/json"
                    last_input_role = "tool"
                    last_input_content = _truncate(payload)
                    last_input_tool_call_id = tool_use_id
                continue

            # ── Non-assistant entries (progress, system, …) ───────────────────
            if entry_type != "assistant":
                continue

            msg = entry.get("message", {})
            usage = msg.get("usage", {})
            if not usage:
                continue

            msg_id = msg.get("id", "") or id(entry)  # fall back to object id
            model = msg.get("model", "claude")
            content_blocks = msg.get("content", [])
            ts = entry.get("timestamp", "")
            stop_reason = msg.get("stop_reason", "")

            # ── Start a new group or extend the current one ───────────────────
            if msg_id != current_group_id:
                _flush_group()
                current_group_id = msg_id
                group_first_usage = usage
                group_total_output_tokens = 0
                group_text_parts = []
                group_tool_use_parts = []
                group_model = model
                group_ts = ts
                group_stop_reason = stop_reason
                # Snapshot the input context at the start of this API call
                group_input_snapshot = {
                    "role": last_input_role,
                    "content": last_input_content,
                    "value": last_input_value,
                    "mime": last_input_mime,
                    "tool_call_id": last_input_tool_call_id,
                }
            else:
                # Extending the group: last non-empty stop_reason wins
                if stop_reason:
                    group_stop_reason = stop_reason

            # Accumulate output tokens and content blocks for this group
            group_total_output_tokens += usage.get("output_tokens", 0)
            for block in content_blocks:
                if not isinstance(block, dict):
                    continue
                if block.get("type") == "text":
                    group_text_parts.append(block.get("text", ""))
                elif block.get("type") == "tool_use":
                    group_tool_use_parts.append(
                        {
                            "id": block.get("id", ""),
                            "name": block.get("name", ""),
                            "input": block.get("input", {}),
                        },
                    )

        _flush_group()

    except Exception as e:
        log.warning("Failed to extract LLM spans: %s", e)

    return spans


# ---------------------------------------------------------------------------
# OpenTelemetry imports (lazy)
# ---------------------------------------------------------------------------


def _otel_imports():
    from opentelemetry import trace
    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
    from opentelemetry.sdk.resources import Resource
    from opentelemetry.sdk.trace import SpanContext, TracerProvider
    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
    from opentelemetry.trace import NonRecordingSpan, SpanKind, StatusCode, TraceFlags

    return (
        trace,
        Resource,
        TracerProvider,
        SpanContext,
        SimpleSpanProcessor,
        OTLPSpanExporter,
        SpanKind,
        TraceFlags,
        NonRecordingSpan,
        StatusCode,
    )


# ---------------------------------------------------------------------------
# Tool schemas
# ---------------------------------------------------------------------------

TOOL_SCHEMAS: dict[str, dict] = {
    "Bash": {
        "type": "object",
        "properties": {
            "command": {"type": "string"},
            "description": {"type": "string"},
            "timeout": {"type": "number"},
            "run_in_background": {"type": "boolean"},
        },
        "required": ["command"],
    },
    "Read": {
        "type": "object",
        "properties": {
            "file_path": {"type": "string"},
            "offset": {"type": "number"},
            "limit": {"type": "number"},
        },
        "required": ["file_path"],
    },
    "Edit": {
        "type": "object",
        "properties": {
            "file_path": {"type": "string"},
            "old_string": {"type": "string"},
            "new_string": {"type": "string"},
            "replace_all": {"type": "boolean"},
        },
        "required": ["file_path", "old_string", "new_string"],
    },
    "Write": {
        "type": "object",
        "properties": {
            "file_path": {"type": "string"},
            "content": {"type": "string"},
        },
        "required": ["file_path", "content"],
    },
    "Glob": {
        "type": "object",
        "properties": {
            "pattern": {"type": "string"},
            "path": {"type": "string"},
        },
        "required": ["pattern"],
    },
    "Grep": {
        "type": "object",
        "properties": {
            "pattern": {"type": "string"},
            "path": {"type": "string"},
            "glob": {"type": "string"},
            "type": {"type": "string"},
            "output_mode": {
                "type": "string",
                "enum": ["content", "files_with_matches", "count"],
            },
            "context": {"type": "number"},
        },
        "required": ["pattern"],
    },
    "Agent": {
        "type": "object",
        "properties": {
            "description": {"type": "string"},
            "prompt": {"type": "string"},
            "subagent_type": {"type": "string"},
            "run_in_background": {"type": "boolean"},
            "isolation": {"type": "string"},
        },
        "required": ["description", "prompt"],
    },
    "Task": {
        "type": "object",
        "properties": {
            "description": {"type": "string"},
            "prompt": {"type": "string"},
            "subagent_type": {"type": "string"},
            "run_in_background": {"type": "boolean"},
        },
        "required": ["description", "prompt", "subagent_type"],
    },
    "WebFetch": {
        "type": "object",
        "properties": {"url": {"type": "string"}, "prompt": {"type": "string"}},
        "required": ["url", "prompt"],
    },
    "WebSearch": {
        "type": "object",
        "properties": {
            "query": {"type": "string"},
            "allowed_domains": {"type": "array", "items": {"type": "string"}},
            "blocked_domains": {"type": "array", "items": {"type": "string"}},
        },
        "required": ["query"],
    },
    "NotebookEdit": {
        "type": "object",
        "properties": {
            "notebook_path": {"type": "string"},
            "new_source": {"type": "string"},
            "cell_id": {"type": "string"},
            "cell_type": {"type": "string"},
            "edit_mode": {"type": "string"},
        },
        "required": ["notebook_path", "new_source"],
    },
    "Skill": {
        "type": "object",
        "properties": {"skill": {"type": "string"}, "args": {"type": "string"}},
        "required": ["skill"],
    },
    "AskUserQuestion": {
        "type": "object",
        "properties": {"questions": {"type": "array"}},
        "required": ["questions"],
    },
    "EnterPlanMode": {"type": "object", "properties": {}, "required": []},
    "ExitPlanMode": {"type": "object", "properties": {}, "required": []},
    "TaskCreate": {
        "type": "object",
        "properties": {
            "subject": {"type": "string"},
            "description": {"type": "string"},
            "activeForm": {"type": "string"},
        },
        "required": ["subject", "description"],
    },
    "TaskUpdate": {
        "type": "object",
        "properties": {
            "taskId": {"type": "string"},
            "status": {"type": "string"},
            "subject": {"type": "string"},
            "description": {"type": "string"},
        },
        "required": ["taskId"],
    },
    "TaskGet": {
        "type": "object",
        "properties": {"taskId": {"type": "string"}},
        "required": ["taskId"],
    },
    "TaskList": {"type": "object", "properties": {}, "required": []},
    "TaskStop": {
        "type": "object",
        "properties": {"task_id": {"type": "string"}},
        "required": [],
    },
    "TaskOutput": {
        "type": "object",
        "properties": {
            "task_id": {"type": "string"},
            "block": {"type": "boolean"},
            "timeout": {"type": "number"},
        },
        "required": ["task_id"],
    },
}

# Tools that map to RETRIEVER span kind (web retrieval operations)
RETRIEVER_TOOLS = {"WebSearch", "WebFetch"}

_MAX_ATTR_BYTES = 8192


def _truncate(value: Any) -> str:
    s = json.dumps(value) if not isinstance(value, str) else value
    encoded = s.encode("utf-8")
    if len(encoded) > _MAX_ATTR_BYTES:
        s = (
            encoded[:_MAX_ATTR_BYTES].decode("utf-8", errors="ignore")
            + "...[truncated]"
        )
    return s


# ---------------------------------------------------------------------------
# FixedSpanIdGenerator — forces the pre-generated root span ID
# ---------------------------------------------------------------------------


def _make_fixed_id_generator(forced_span_id_hex: str):
    import secrets

    from opentelemetry.sdk.trace.id_generator import IdGenerator

    forced_int = int(forced_span_id_hex, 16)

    class FixedSpanIdGenerator(IdGenerator):
        def __init__(self):
            self._used = False

        def generate_span_id(self) -> int:
            if not self._used:
                self._used = True
                return forced_int
            return int(secrets.token_hex(8), 16)

        def generate_trace_id(self) -> int:
            return int(secrets.token_hex(16), 16)

    return FixedSpanIdGenerator()


# ---------------------------------------------------------------------------
# OTLP export helper
# ---------------------------------------------------------------------------


def _build_and_export_spans(
    config: dict,
    session_id: str,
    username: str,
    span_records: list[dict],
) -> None:
    """Create spans from records and export via OTLP HTTP."""
    (
        trace,
        Resource,
        TracerProvider,
        SpanContext,
        SimpleSpanProcessor,
        OTLPSpanExporter,
        SpanKind,
        TraceFlags,
        NonRecordingSpan,
        StatusCode,
    ) = _otel_imports()

    resource = Resource.create(
        {
            "service.name": "claude-code",
            "arthur.task": config["task_id"],
            "arthur.session": session_id,
            "arthur.user": username,
        },
    )

    for rec in span_records:
        # Each span gets its own exporter so that provider.shutdown() on span N
        # doesn't mark the shared exporter as closed before span N+1 is sent.
        exporter = OTLPSpanExporter(
            endpoint=config["endpoint"],
            headers={"Authorization": f"Bearer {config['api_key']}"},
        )
        # Resolve None kind (used for LLM spans set by caller)
        kind = rec.get("kind") or SpanKind.CLIENT

        id_generator = None
        if rec.get("force_span_id"):
            id_generator = _make_fixed_id_generator(rec["span_id_hex"])

        kwargs = {"resource": resource}
        if id_generator:
            kwargs["id_generator"] = id_generator

        provider = TracerProvider(**kwargs)
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        tracer = provider.get_tracer("claude-code-tracer")

        ctx = None
        if rec.get("parent_span_id_hex"):
            parent_sc = SpanContext(
                trace_id=int(rec["trace_id_hex"], 16),
                span_id=int(rec["parent_span_id_hex"], 16),
                is_remote=True,
                trace_flags=TraceFlags(TraceFlags.SAMPLED),
            )
            ctx = trace.set_span_in_context(NonRecordingSpan(parent_sc))

        span = tracer.start_span(
            name=rec["name"],
            context=ctx,
            kind=kind,
            start_time=rec["start_ns"],
        )

        # Patch trace_id if provider generated a different one
        try:
            sc = span.get_span_context()
            desired = int(rec["trace_id_hex"], 16)
            if sc.trace_id != desired:
                span._context = SpanContext(
                    trace_id=desired,
                    span_id=sc.span_id,
                    is_remote=sc.is_remote,
                    trace_flags=sc.trace_flags,
                )
        except Exception:
            pass

        span.set_attribute("arthur_span_version", "arthur_span_v1")
        span.set_attribute("session.id", session_id)
        span.set_attribute("user.id", username)
        for k, v in rec.get("attributes", {}).items():
            span.set_attribute(k, v)

        if rec.get("error"):
            span.set_status(StatusCode.ERROR, description=rec.get("error_msg", ""))

        span.end(end_time=rec["end_ns"])
        provider.shutdown()


# ---------------------------------------------------------------------------
# Turn completion
# ---------------------------------------------------------------------------


def _emit_pending_llm_spans(
    state: dict,
    transcript_path: Optional[str],
    config: dict,
) -> None:
    """Emit any new LLM spans from the transcript that haven't been sent yet.

    Called from handle_post_tool (real-time, after each tool response) and
    handle_stop (catches the final LLM response after the last tool call).
    Updates state.current_trace in-place so the caller must save state afterwards.
    """
    current_trace = state.get("current_trace")
    if not current_trace or not transcript_path:
        return

    (_, _, _, _, _, _, SpanKind, _, _, _) = _otel_imports()

    all_llm_spans = _extract_llm_spans_for_turn(
        transcript_path,
        current_trace.get("human_count_at_start", 0),
        current_trace["trace_id"],
        current_trace["root_span_id"],
    )

    emitted = current_trace.get("emitted_llm_span_count", 0)
    new_spans = all_llm_spans[emitted:]

    # Always keep last_llm_output in sync with the last known assistant message,
    # even if there are no new spans to emit (e.g. all were emitted by post_tool
    # and the final text response was already the last one processed).
    if all_llm_spans:
        current_trace["last_llm_output"] = all_llm_spans[-1]["attributes"].get(
            "output.value",
            "",
        )

    if not new_spans:
        return

    for sp in new_spans:
        sp["kind"] = SpanKind.CLIENT

    _build_and_export_spans(
        config=config,
        session_id=state["session_id"],
        username=state.get("username", "unknown"),
        span_records=new_spans,
    )

    current_trace["emitted_llm_span_count"] = emitted + len(new_spans)


def _complete_turn(
    state: dict,
    config: dict,
    transcript_path: Optional[str],
    end_ns: int,
) -> None:
    """Send the CHAIN root span for the current turn's trace.

    LLM spans are emitted in real-time from _emit_pending_llm_spans (called
    from handle_post_tool and handle_stop), so they do not need to be re-sent
    here.  The root span's output.value is taken from the last LLM output
    already tracked in state.
    """
    current_trace = state.get("current_trace")
    if not current_trace:
        return

    (_, _, _, _, _, _, SpanKind, _, _, _) = _otel_imports()

    trace_id = current_trace["trace_id"]
    root_span_id = current_trace["root_span_id"]
    turn_start_ns = current_trace["turn_start_ns"]
    turn_number = current_trace.get("turn_number", 1)
    prompt_preview = current_trace.get("prompt_preview", "")
    final_output = current_trace.get("last_llm_output", "")
    # For the first turn of a subagent the CHAIN span is a child of the parent's
    # Agent tool span, making all subagent work visible inside the parent trace.
    parent_agent_span_id = current_trace.get("parent_agent_span_id")

    # Root CHAIN span
    username = state.get("username", "unknown")
    root_attrs: dict[str, Any] = {
        "openinference.span.kind": "CHAIN",
        "session.id": state["session_id"],
        "user.id": username,
        "arthur.turn_number": turn_number,
    }
    if prompt_preview:
        root_attrs["input.value"] = prompt_preview
        root_attrs["input.mime_type"] = "text/plain"
    if final_output:
        root_attrs["output.value"] = final_output
        root_attrs["output.mime_type"] = "text/plain"

    _build_and_export_spans(
        config=config,
        session_id=state["session_id"],
        username=username,
        span_records=[
            {
                "trace_id_hex": trace_id,
                "span_id_hex": root_span_id,
                "parent_span_id_hex": parent_agent_span_id,
                "name": "claude-code-turn",
                "kind": SpanKind.INTERNAL,
                "start_ns": turn_start_ns,
                "end_ns": end_ns,
                "attributes": root_attrs,
                "force_span_id": True,
            },
        ],
    )


# ---------------------------------------------------------------------------
# Tool span record builder (shared by post_tool and post_tool_failure)
# ---------------------------------------------------------------------------


def _build_tool_span_record(
    tool_name: str,
    tool_input: Any,
    tool_response: Any,
    start_ns: int,
    end_ns: int,
    trace_id: str,
    root_span_id: str,
    is_failure: bool = False,
    error_msg: str = "",
    span_id: Optional[str] = None,
) -> dict:
    """Build a span record dict for a tool call (success or failure).

    Returns a record suitable for passing to _build_and_export_spans.
    The kind field is left as None so _build_and_export_spans defaults to
    SpanKind.CLIENT.
    """
    if tool_name in RETRIEVER_TOOLS:
        span_kind_str = "RETRIEVER"
    elif tool_name in ("Agent", "Task"):
        span_kind_str = "AGENT"
    else:
        span_kind_str = "TOOL"

    input_value = (
        json.dumps(tool_input) if not isinstance(tool_input, str) else tool_input
    )
    output_str = (
        json.dumps(tool_response)
        if not isinstance(tool_response, str)
        else tool_response
    )

    if is_failure and error_msg:
        output_value = f"ERROR: {error_msg}\n{output_str}"
    else:
        output_value = output_str

    attrs: dict[str, Any] = {
        "openinference.span.kind": span_kind_str,
        "tool.name": tool_name,
        "input.value": input_value,
        "input.mime_type": "application/json",
        "output.value": output_value,
        "output.mime_type": "application/json",
    }

    # agent.name for AGENT spans: prefer subagent_type from input, fall back to tool_name
    if span_kind_str == "AGENT":
        agent_name = ""
        if isinstance(tool_input, dict):
            agent_name = tool_input.get("subagent_type", "")
        attrs["agent.name"] = agent_name or tool_name

    schema = TOOL_SCHEMAS.get(tool_name)
    if schema:
        attrs["tool.json_schema"] = json.dumps(schema)

    # RETRIEVER-specific attributes per OpenInference spec
    if span_kind_str == "RETRIEVER":
        if tool_name == "WebSearch":
            query = tool_input.get("query", "") if isinstance(tool_input, dict) else ""
            attrs["input.value"] = query
            attrs["input.mime_type"] = "text/plain"
        elif tool_name == "WebFetch":
            url = tool_input.get("url", "") if isinstance(tool_input, dict) else ""
            attrs["input.value"] = url
            attrs["input.mime_type"] = "text/plain"

        # First retrieval document content
        if isinstance(tool_response, str):
            doc_content = _truncate(tool_response)
        else:
            doc_content = _truncate(json.dumps(tool_response))
        attrs["retrieval.documents.0.document.content"] = doc_content

    rec: dict[str, Any] = {
        "trace_id_hex": trace_id,
        "span_id_hex": span_id or _new_span_id(),
        "parent_span_id_hex": root_span_id,
        "name": tool_name,
        "kind": None,  # defaults to SpanKind.CLIENT in _build_and_export_spans
        "start_ns": start_ns,
        "end_ns": end_ns,
        "attributes": attrs,
        "force_span_id": span_id is not None,
    }

    if is_failure:
        rec["error"] = True
        rec["error_msg"] = error_msg

    return rec


# ---------------------------------------------------------------------------
# Handlers
# ---------------------------------------------------------------------------


def handle_user_prompt_submit(data: dict, config: dict) -> None:
    """Handle UserPromptSubmit hook: complete previous trace and start a new one.

    Fires before Claude processes the prompt, giving an accurate turn_start_ns
    and the exact prompt text without transcript parsing.
    """
    session_id = data.get("session_id", "unknown")
    prompt = data.get("prompt", "")
    now_ns = time.time_ns()

    state = _load_state(session_id)

    # Initialize session on first event
    if not state.get("session_id"):
        state["session_id"] = session_id
        state["session_start_ns"] = now_ns
        state["username"] = os.environ.get(
            "USER",
            os.environ.get("USERNAME", "unknown"),
        )
        state["human_msg_count"] = 0
        state["turn_number"] = 0

        # If a parent session pre-registered an Agent invocation, inherit its
        # trace context so this subagent's spans are nested inside the parent trace.
        # Pass the prompt so parallel agents can each claim their own context.
        parent_ctx = _claim_pending_agent_context(prompt=prompt)
        if parent_ctx:
            state["inherited_trace_id"] = parent_ctx["parent_trace_id"]
            state["parent_agent_span_id"] = parent_ctx["parent_span_id"]

    # Complete the previous turn's trace if one is in progress.
    # Use the cached transcript path from state so that any mid-session
    # project-dir changes (e.g. gh pr create writing to a worktree dir) do
    # not silently redirect us to a shadow transcript file.
    _carry_trace_id = None
    _carry_parent_span = None
    if state.get("current_trace"):
        # When handle_pre_tool fires first for a new subagent session it claims
        # the pending parent context and stores it inside current_trace.  Save
        # those values before completing the old trace so the real first turn
        # (created below) inherits the same trace ID and parent span — otherwise
        # UPS would generate a fresh trace ID and sever the subagent connection.
        _old_ct = state["current_trace"]
        if _old_ct.get("_pretool_created") and _old_ct.get("parent_agent_span_id"):
            _carry_trace_id = _old_ct["trace_id"]
            _carry_parent_span = _old_ct["parent_agent_span_id"]
        transcript_path = _get_cached_transcript_path(data, state, session_id)
        _emit_pending_llm_spans(state, transcript_path, config)
        _complete_turn(state, config, transcript_path, now_ns)
        # Clear the cached path so the new turn resolves it fresh below
        state.pop("transcript_path", None)

    # Count human messages for LLM span extraction.
    # UserPromptSubmit fires before Claude processes the prompt, so the new
    # message is not yet in the transcript — human_count_at_start equals the
    # current count (the new prompt will become count+1 in the transcript).
    # Resolve and immediately cache the transcript path for this new turn.
    transcript_path = _get_cached_transcript_path(data, state, session_id)
    current_human_count = (
        _count_human_messages(transcript_path) if transcript_path else 0
    )

    turn_number = state.get("turn_number", 0) + 1
    state["turn_number"] = turn_number
    state["human_msg_count"] = current_human_count

    # Use the inherited trace ID if this session was spawned as a subagent, so
    # all spans land in the parent's trace.  Also consume the parent agent span
    # ID (used only for the first CHAIN span so it becomes a child of the Agent
    # tool span in the parent trace).
    # _carry_* values are non-None only when handle_pre_tool fired before UPS
    # and already placed the inherited context inside current_trace; in that
    # case state-level keys were already consumed by the pre_tool fallback.
    trace_id = (
        state.pop("inherited_trace_id", None) or _carry_trace_id or _new_trace_id()
    )
    parent_agent_span_id = state.pop("parent_agent_span_id", None) or _carry_parent_span

    state["current_trace"] = {
        "trace_id": trace_id,
        "root_span_id": _new_span_id(),
        "turn_start_ns": now_ns,
        "turn_number": turn_number,
        "human_count_at_start": current_human_count,
        "prompt_preview": _truncate(prompt) if prompt else "",
        # Non-None only for the first turn of a subagent; cleared after _complete_turn
        "parent_agent_span_id": parent_agent_span_id,
    }

    _save_state(session_id, state)


def handle_pre_tool(data: dict, config: dict) -> None:
    session_id = data.get("session_id", "unknown")
    tool_name = data.get("tool_name", "")
    tool_input = data.get("tool_input", {})
    now_ns = time.time_ns()

    state = _load_state(session_id)

    # Initialize session on first call
    if not state.get("session_id"):
        state["session_id"] = session_id
        state["session_start_ns"] = now_ns
        state["username"] = os.environ.get(
            "USER",
            os.environ.get("USERNAME", "unknown"),
        )
        state["human_msg_count"] = 0
        state["turn_number"] = 0

        # Claim any pending parent context in case UserPromptSubmit didn't fire
        # for this subagent (e.g. Explore-type agents that are spawned
        # programmatically bypass the UserPromptSubmit hook).  prompt="" triggers
        # the newest-first fallback since the prompt text is not available here.
        parent_ctx = _claim_pending_agent_context(prompt="")
        if parent_ctx:
            state["inherited_trace_id"] = parent_ctx["parent_trace_id"]
            state["parent_agent_span_id"] = parent_ctx["parent_span_id"]

    # Detect context continuation: Claude Code compresses context and continues
    # without firing UserPromptSubmit for the resumption message. The symptom is
    # current_trace still set from the previous turn, but the transcript has grown
    # by more than one human message since that trace started.
    if state.get("current_trace"):
        ct = state["current_trace"]
        _tp = _get_cached_transcript_path(data, state, session_id)
        if _tp:
            _cur_human = _count_human_messages(_tp)
            _start = ct.get("human_count_at_start", 0)
            if _cur_human > _start + 1:
                # A new turn started without UserPromptSubmit. Complete the old
                # trace and fall through to the fallback that creates a new one.
                _emit_pending_llm_spans(state, _tp, config)
                _complete_turn(state, config, _tp, now_ns)
                state.pop("current_trace", None)
                state.pop("pending_tools", None)

    # Fallback: create a trace if UserPromptSubmit has not set one yet.
    # This handles cases where the hook isn't registered or fires before the
    # UserPromptSubmit event is available.
    if not state.get("current_trace"):
        transcript_path = _get_cached_transcript_path(data, state, session_id)
        current_human_count = (
            _count_human_messages(transcript_path) if transcript_path else 0
        )
        prompt_preview = (
            _get_latest_human_message(transcript_path) if transcript_path else ""
        )

        turn_number = state.get("turn_number", 0) + 1
        state["turn_number"] = turn_number
        state["human_msg_count"] = current_human_count

        # Consume any inherited parent context stored by the session init block
        # (or by a UserPromptSubmit that ran in a prior hook invocation).
        trace_id = state.pop("inherited_trace_id", None) or _new_trace_id()
        parent_agent_span_id = state.pop("parent_agent_span_id", None)

        state["current_trace"] = {
            "trace_id": trace_id,
            "root_span_id": _new_span_id(),
            "turn_start_ns": now_ns,
            "turn_number": turn_number,
            # human_count_at_start is one less than the current count so that
            # _extract_llm_spans_for_turn finds entries where count > this value,
            # i.e. entries belonging to the NEW turn whose prompt is the
            # current_human_count-th human message.
            "human_count_at_start": max(0, current_human_count - 1),
            "prompt_preview": _truncate(prompt_preview) if prompt_preview else "",
            "parent_agent_span_id": parent_agent_span_id,
            # Sentinel so handle_user_prompt_submit can detect that this trace
            # was created by the pre_tool fallback path (before UPS fired) and
            # carry the inherited context forward to the real first turn.
            "_pretool_created": True,
        }

    pending_entry: dict[str, Any] = {
        "tool_name": tool_name,
        "tool_input": tool_input,
        "start_ns": now_ns,
    }

    # For Agent/Task tool calls, pre-allocate the span ID and write a side-channel
    # file so the spawned subagent can inherit this trace.
    current_trace = state.get("current_trace")
    if tool_name in ("Agent", "Task") and current_trace:
        agent_span_id = _new_span_id()
        pending_entry["pre_allocated_span_id"] = agent_span_id
        if isinstance(tool_input, dict):
            # "Agent" uses "prompt"; "Task" may use "description"
            agent_prompt = tool_input.get("prompt", tool_input.get("description", ""))
        else:
            agent_prompt = ""
        _write_pending_agent_context(
            parent_session_id=session_id,
            parent_trace_id=current_trace["trace_id"],
            parent_span_id=agent_span_id,
            agent_prompt=agent_prompt,
        )
        # Key by span_id so parallel Agent/Task calls in the same session don't
        # overwrite each other — post_tool finds the right entry by tool_input match.
        pending_key = f"{tool_name}#{agent_span_id}"
    else:
        pending_key = tool_name

    state.setdefault("pending_tools", {})[pending_key] = pending_entry
    _save_state(session_id, state)


def _get_latest_human_message(transcript_path: str) -> str:
    """Return the text of the most recent human message in the transcript."""
    try:
        last = ""
        for line in Path(transcript_path).read_text().splitlines():
            if not line.strip():
                continue
            try:
                entry = json.loads(line)
                if _is_human_message(entry):
                    content = entry.get("message", {}).get("content", "")
                    if isinstance(content, str):
                        last = content
            except json.JSONDecodeError:
                pass
        return last
    except Exception:
        return ""


def _find_pending_tool_entry(
    state: dict,
    tool_name: str,
    data_tool_input: Any,
) -> tuple[dict, str]:
    """Return (pending_entry, pending_key) for the given tool.

    For Agent/Task tools, parallel calls each get a compound key
    ``"{tool_name}#{span_id}"``.  We find the right entry by matching
    ``tool_input``; if no match, fall back to the first entry with the right
    prefix so single-agent sessions and legacy state files still work.

    For all other tools the key is just ``tool_name``.
    """
    pending = state.get("pending_tools", {})
    prefix = f"{tool_name}#"

    if tool_name in ("Agent", "Task"):
        # Pass 1: exact tool_input match
        if data_tool_input is not None:
            for k, v in pending.items():
                if k.startswith(prefix) and v.get("tool_input") == data_tool_input:
                    return v, k
        # Pass 2: any entry with this tool's prefix (single-agent / legacy key)
        for k, v in pending.items():
            if k.startswith(prefix) or k == tool_name:
                return v, k
        return {}, tool_name

    entry = pending.get(tool_name, {})
    return entry, tool_name


def _find_active_agent_span_id(
    state: dict,
    current_pending_key: str,
) -> Optional[str]:
    """Return the pre_allocated_span_id of the innermost pending Agent/Task call.

    When an inline subagent (e.g. Explore) makes tool calls within the parent
    session, those calls appear in pending_tools alongside the Agent entry.
    This detects that situation and returns the Agent's span_id so the sub-tool
    is correctly parented under the Agent span rather than the CHAIN root.
    For nested agents, the most-recently-started one wins (innermost parent).
    """
    pending = state.get("pending_tools", {})
    best_start_ns = -1
    best_span_id: Optional[str] = None
    for k, v in pending.items():
        if k == current_pending_key:
            continue  # don't self-reference
        base_tool = k.split("#")[0] if "#" in k else k
        if base_tool in ("Agent", "Task") and v.get("pre_allocated_span_id"):
            if v.get("start_ns", 0) > best_start_ns:
                best_start_ns = v["start_ns"]
                best_span_id = v["pre_allocated_span_id"]
    return best_span_id


def handle_post_tool(data: dict, config: dict) -> None:
    session_id = data.get("session_id", "unknown")
    end_ns = time.time_ns()

    # Read state once to build the tool span (no mutation needed).
    state = _load_state(session_id)
    if not state:
        log.warning("No state found for session %s in post_tool", session_id)
        return

    current_trace = state.get("current_trace")
    if not current_trace:
        log.debug("No current trace for session %s in post_tool", session_id)
        return

    tool_name = data.get("tool_name", "unknown")
    data_tool_input = data.get("tool_input")
    current_tool, pending_key = _find_pending_tool_entry(
        state,
        tool_name,
        data_tool_input,
    )
    tool_input = data_tool_input or current_tool.get("tool_input", {})
    tool_response = data.get("tool_response", {})
    start_ns = current_tool.get("start_ns", end_ns - 1_000_000)
    # For Agent tool calls the span ID was pre-allocated at PreToolUse time so
    # the subagent could reference it as its parent span.
    pre_allocated_span_id = current_tool.get("pre_allocated_span_id")

    # If an Agent/Task tool is still pending, this tool ran inside that agent
    # (inline subagent pattern). Parent it to the agent span, not the CHAIN root.
    active_agent_span_id = _find_active_agent_span_id(state, pending_key)
    span_record = _build_tool_span_record(
        tool_name=tool_name,
        tool_input=tool_input,
        tool_response=tool_response,
        start_ns=start_ns,
        end_ns=end_ns,
        trace_id=current_trace["trace_id"],
        root_span_id=active_agent_span_id or current_trace["root_span_id"],
        span_id=pre_allocated_span_id,
    )

    # Export the tool span before acquiring the lock — this is the slow network
    # I/O and does not mutate state, so it is safe to run outside the lock.
    _build_and_export_spans(
        config=config,
        session_id=session_id,
        username=state.get("username", "unknown"),
        span_records=[span_record],
    )

    # Acquire an exclusive per-session lock before emitting LLM spans.
    # Parallel PostToolUse processes race on emitted_llm_span_count; the lock
    # ensures each process re-reads the latest count and never double-emits.
    # Resolve the transcript path *inside* the lock so we always use the path
    # that was cached in state at UserPromptSubmit time — not the (potentially
    # stale or redirected) path carried in the current hook payload.
    with _session_lock(session_id):
        state = _load_state(session_id)
        # Use the compound key resolved above; fall back to tool_name for
        # any other entry that might have been written with the old plain key.
        pt = state.get("pending_tools", {})
        pt.pop(pending_key, None)
        pt.pop(tool_name, None)
        transcript_path = _get_cached_transcript_path(data, state, session_id)
        _emit_pending_llm_spans(state, transcript_path, config)
        _save_state(session_id, state)


def handle_post_tool_failure(data: dict, config: dict) -> None:
    """Handle PostToolUseFailure: emit an error TOOL/RETRIEVER/AGENT span."""
    session_id = data.get("session_id", "unknown")
    end_ns = time.time_ns()

    state = _load_state(session_id)
    if not state:
        log.warning("No state found for session %s in post_tool_failure", session_id)
        return

    current_trace = state.get("current_trace")
    if not current_trace:
        log.debug("No current trace for session %s in post_tool_failure", session_id)
        return

    tool_name = data.get("tool_name", "unknown")
    data_tool_input = data.get("tool_input")
    current_tool, pending_key = _find_pending_tool_entry(
        state,
        tool_name,
        data_tool_input,
    )
    tool_input = data_tool_input or current_tool.get("tool_input", {})
    tool_response = data.get("tool_response", {})
    start_ns = current_tool.get("start_ns", end_ns - 1_000_000)

    # Extract error message from various possible fields
    error_msg = ""
    if isinstance(tool_response, dict):
        error_msg = tool_response.get("error", tool_response.get("message", ""))
    elif isinstance(tool_response, str):
        error_msg = tool_response
    if not error_msg:
        error_msg = data.get("error", data.get("error_message", "Tool call failed"))

    active_agent_span_id = _find_active_agent_span_id(state, pending_key)
    span_record = _build_tool_span_record(
        tool_name=tool_name,
        tool_input=tool_input,
        tool_response=tool_response,
        start_ns=start_ns,
        end_ns=end_ns,
        trace_id=current_trace["trace_id"],
        root_span_id=active_agent_span_id or current_trace["root_span_id"],
        is_failure=True,
        error_msg=error_msg,
    )

    _build_and_export_spans(
        config=config,
        session_id=session_id,
        username=state.get("username", "unknown"),
        span_records=[span_record],
    )

    with _session_lock(session_id):
        state = _load_state(session_id)
        pt = state.get("pending_tools", {})
        pt.pop(pending_key, None)
        pt.pop(tool_name, None)
        transcript_path = _get_cached_transcript_path(data, state, session_id)
        _emit_pending_llm_spans(state, transcript_path, config)
        _save_state(session_id, state)


def handle_stop(data: dict, config: dict) -> None:
    session_id = data.get("session_id", "unknown")
    end_ns = time.time_ns()

    with _session_lock(session_id):
        state = _load_state(session_id)
        if not state.get("current_trace"):
            log.debug("No active trace for session %s at stop", session_id)
            return

        transcript_path = _get_cached_transcript_path(data, state, session_id)

        # Detect context continuation: same check as handle_pre_tool.
        # When no tool calls happen during a continuation turn, handle_pre_tool
        # never fires, so the stale trace would otherwise be completed with the
        # wrong LLM spans.  Complete the old trace here and start a new one so
        # the normal stop logic below runs against the correct turn.
        if transcript_path:
            ct = state["current_trace"]
            _cur_human = _count_human_messages(transcript_path)
            _start = ct.get("human_count_at_start", 0)
            if _cur_human > _start + 1:
                _emit_pending_llm_spans(state, transcript_path, config)
                _complete_turn(state, config, transcript_path, end_ns)
                state.pop("current_trace", None)
                state.pop("pending_tools", None)

                current_human_count = _cur_human
                prompt_preview = _get_latest_human_message(transcript_path)
                turn_number = state.get("turn_number", 0) + 1
                state["turn_number"] = turn_number
                state["human_msg_count"] = current_human_count
                state["current_trace"] = {
                    "trace_id": _new_trace_id(),
                    "root_span_id": _new_span_id(),
                    "turn_start_ns": end_ns,
                    "turn_number": turn_number,
                    "human_count_at_start": max(0, current_human_count - 1),
                    "prompt_preview": (
                        _truncate(prompt_preview) if prompt_preview else ""
                    ),
                }

        # Emit the final LLM span(s) — the last response has no tool call after it,
        # so handle_post_tool never got the chance to emit it.
        _emit_pending_llm_spans(state, transcript_path, config)

        _complete_turn(state, config, transcript_path, end_ns)

        _delete_state(session_id)

    _cleanup_stale_states()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    if len(sys.argv) < 2:
        print(
            "Usage: claude_code_tracer.py <pre_tool|post_tool|post_tool_failure|user_prompt_submit|stop>",
            file=sys.stderr,
        )
        sys.exit(0)

    event = sys.argv[1]

    try:
        raw = sys.stdin.read()
        data = json.loads(raw) if raw.strip() else {}
    except json.JSONDecodeError as e:
        log.warning("Failed to parse stdin JSON: %s", e)
        data = {}

    config = discover_config()
    if config is None:
        log.debug("Arthur Engine not configured, skipping tracing.")
        sys.exit(0)

    try:
        if event == "pre_tool":
            handle_pre_tool(data, config)
        elif event == "post_tool":
            handle_post_tool(data, config)
        elif event == "post_tool_failure":
            handle_post_tool_failure(data, config)
        elif event == "user_prompt_submit":
            handle_user_prompt_submit(data, config)
        elif event == "stop":
            handle_stop(data, config)
        else:
            log.warning("Unknown event: %s", event)
    except Exception as e:
        log.warning("Tracer error (%s): %s", event, e, exc_info=True)

    sys.exit(0)


if __name__ == "__main__":
    main()
