"""Codex hook handlers.

Entry point: python3 -m staso.codex

Reads a JSON event from stdin, dispatches to the right handler.
Each hook invocation is a separate process; state is persisted to disk.

Trace model: one trace per conversation turn. All turns share session_id.
  Turn span (root, kind=agent)
    |- LLM span   (kind=llm)   <- extracted from transcript in Stop
    |- Tool span   (kind=tool)  <- timed via PreToolUse/PostToolUse pair (shell/exec only)
    +- Tool span   (kind=tool)  <- extracted from transcript in Stop (all other tools)

Known limitation: When embed_env_in_command is used (Codex default), env
vars including STASO_API_KEY are visible in ``ps aux`` output.  Codex does
not support an ``env`` section in its hooks config, so this is the only
mechanism available.
"""
from __future__ import annotations

import json
import logging
import os
import sys
import time
from typing import Any

from staso._common.helpers import new_id, now_iso, ts_to_iso, turn_span_name
from staso._common.repair import auto_sync_enabled, repair_hooks
from staso._common.sender import SpanSender, sender_from_env
from staso._common.span import base_span
from .state import (
    HookState,
    acquire_lock,
    gc_stale_states,
    load_state,
    release_lock,
    save_state,
)
from .transcript import (
    LLMMessage,
    ToolCallRecord,
    find_transcript,
    parse_turn_llm_data,
    parse_turn_tool_calls,
)

logger = logging.getLogger("staso.codex")

_SOURCE = "codex"
_AGENT_ID = os.environ.get("STASO_AGENT_ID", "codex")


# -- helpers -------------------------------------------------------------------


def _load_active_state(session_id: str) -> HookState | None:
    """Load state only when there is an active turn. Returns None otherwise."""
    state = load_state(session_id)
    if state is None or state.current_trace_id is None or state.current_turn_span_id is None:
        return None
    return state


def _resolve_tool_name(
    payload_name: str,
    tool_use_id: str,
    transcript_path: str,
) -> str:
    """Resolve the real tool name from the transcript.

    Codex hardcodes ``tool_name: "Bash"`` in all hook payloads.  The actual
    tool name (exec_command, apply_patch, update_plan, etc.) only appears in
    the JSONL transcript as ``function_call`` response items.  We scan the
    transcript backwards for a matching ``call_id`` to recover the real name.
    """
    if payload_name != "Bash" or not transcript_path or not tool_use_id:
        return payload_name

    try:
        from pathlib import Path

        path = Path(transcript_path)
        if not path.exists():
            return payload_name

        # Read lines in reverse — the matching call_id is usually near the end
        lines = path.read_text(encoding="utf-8").splitlines()
        for line in reversed(lines):
            if tool_use_id not in line:
                continue
            try:
                entry = json.loads(line)
                # transcript stores entries at top-level or under "payload"
                p = entry if "call_id" in entry else entry.get("payload", {})
                if p.get("call_id") == tool_use_id:
                    ptype = p.get("type", "")
                    # function_call (exec_command, update_plan, etc.)
                    if ptype == "function_call":
                        real_name = p.get("name", "")
                        if real_name:
                            return real_name
                    # custom_tool_call (apply_patch freeform)
                    if ptype == "custom_tool_call":
                        real_name = p.get("name", "")
                        if real_name:
                            return real_name
            except (json.JSONDecodeError, KeyError):
                continue

        return payload_name
    except Exception:
        logger.debug("staso: transcript tool-name lookup failed", exc_info=True)
        return payload_name


def _aggregate_llm_messages(
    messages: list[LLMMessage],
) -> tuple[int, int, str]:
    """Return (total_input_tokens, total_output_tokens, last_model)."""
    total_in = sum(m.input_tokens for m in messages)
    total_out = sum(m.output_tokens for m in messages)
    last_model = messages[-1].model if messages else ""
    return total_in, total_out, last_model


# -- handlers ------------------------------------------------------------------


def handle_session_start(payload: dict[str, Any], _sender: SpanSender) -> None:
    session_id = payload["session_id"]
    cwd = payload.get("cwd", "")

    try:
        gc_stale_states()
    except Exception:
        logger.debug("staso: stale state GC failed", exc_info=True)

    if auto_sync_enabled():
        try:
            repair_hooks("codex", cwd)
        except Exception:
            logger.debug("staso: auto-repair failed", exc_info=True)

    if not acquire_lock(session_id):
        return
    try:
        # Find and cache the transcript path for this session
        transcript = find_transcript(session_id)
        transcript_path = str(transcript) if transcript else ""

        state = HookState(
            session_id=session_id,
            cwd=cwd,
            transcript_path=transcript_path,
        )
        save_state(state)
    finally:
        release_lock(session_id)


def handle_user_prompt_submit(payload: dict[str, Any], _sender: SpanSender) -> None:
    session_id = payload["session_id"]
    if not acquire_lock(session_id):
        return
    try:
        state = load_state(session_id) or HookState(
            session_id=session_id,
            cwd=payload.get("cwd", ""),
        )

        # Discover transcript if not found during SessionStart
        if not state.transcript_path:
            transcript = find_transcript(session_id)
            if transcript:
                state.transcript_path = str(transcript)

        state.turn_count += 1
        state.current_trace_id = new_id()
        state.current_turn_span_id = new_id()
        state.current_turn_start = time.time()
        state.current_turn_prompt = payload.get("prompt", "")
        save_state(state)
    finally:
        release_lock(session_id)


def handle_pre_tool_use(payload: dict[str, Any], _sender: SpanSender) -> None:
    session_id = payload["session_id"]
    if not acquire_lock(session_id):
        return
    try:
        state = _load_active_state(session_id)
        if state is None:
            return
        tool_use_id = payload.get("tool_use_id", "")
        if tool_use_id:
            state.pending_tool_starts[tool_use_id] = time.time()
            save_state(state)
    finally:
        release_lock(session_id)


def handle_post_tool_use(payload: dict[str, Any], sender: SpanSender) -> None:
    session_id = payload["session_id"]
    if not acquire_lock(session_id):
        return
    try:
        state = _load_active_state(session_id)
        if state is None:
            return

        now = time.time()
        tool_use_id = payload.get("tool_use_id", "")
        start_time = state.pending_tool_starts.pop(tool_use_id, None)
        duration_ms = (now - start_time) * 1000 if start_time else 0.0

        raw_tool_name = payload.get("tool_name", "unknown")
        tool_input = payload.get("tool_input") or {}
        tool_response = payload.get("tool_response") or {}
        cwd = state.cwd or payload.get("cwd", "")

        # Codex hardcodes tool_name="Bash" for all hooks — resolve the real
        # name (exec_command, update_plan, apply_patch, etc.) from the transcript.
        tool_name = _resolve_tool_name(
            raw_tool_name, tool_use_id, state.transcript_path
        )

        span = base_span(
            trace_id=state.current_trace_id,
            span_id=new_id(),
            parent_span_id=state.current_turn_span_id,
            name=f"tool:{tool_name}",
            kind="tool",
            timestamp=ts_to_iso(start_time) if start_time else now_iso(),
            duration_ms=duration_ms,
            session_id=session_id,
            agent_id=_AGENT_ID,
            input_data={"tool_input": tool_input},
            output_data={"tool_response": tool_response},
            metadata={"tool_name": tool_name, "source": _SOURCE, "cwd": cwd},
        )
        sender.send([span])
        save_state(state)
    finally:
        release_lock(session_id)


def handle_stop(payload: dict[str, Any], sender: SpanSender | None) -> None:
    session_id = payload["session_id"]
    if not acquire_lock(session_id):
        return
    try:
        state = load_state(session_id)
        if state is None or state.current_trace_id is None or state.current_turn_span_id is None:
            return

        # Last chance to discover transcript if earlier attempts failed
        if not state.transcript_path:
            transcript = find_transcript(session_id)
            if transcript:
                state.transcript_path = str(transcript)
                logger.debug("staso: transcript discovered in Stop: %s", transcript)

        # Extract LLM data from transcript (by turn count, not line offset)
        llm_messages, total_turns = parse_turn_llm_data(
            state.transcript_path, state.turns_processed
        )

        if sender is not None:
            spans: list[dict[str, Any]] = []
            turn_end = time.time()
            turn_start = state.current_turn_start or turn_end
            turn_duration_ms = (turn_end - turn_start) * 1000
            project_name = os.path.basename(state.cwd) if state.cwd else _SOURCE

            # Create LLM spans from transcript data
            for msg in llm_messages:
                llm_span = base_span(
                    trace_id=state.current_trace_id,
                    span_id=new_id(),
                    parent_span_id=state.current_turn_span_id,
                    name=f"llm:{msg.model}" if msg.model else "llm:openai",
                    kind="llm",
                    timestamp=msg.timestamp or ts_to_iso(turn_start),
                    duration_ms=turn_duration_ms,  # single LLM call per turn
                    session_id=session_id,
                    agent_id=_AGENT_ID,
                    model=msg.model,
                    input_tokens=msg.input_tokens,
                    output_tokens=msg.output_tokens,
                    total_tokens=msg.total_tokens,
                    output_data={"content": msg.content} if msg.content else None,
                    metadata={
                        "source": _SOURCE,
                        "cached_input_tokens": msg.cached_input_tokens,
                        "reasoning_output_tokens": msg.reasoning_output_tokens,
                    },
                )
                spans.append(llm_span)

            # Extract non-shell tool calls from transcript (apply_patch,
            # update_plan, view_image, spawn_agent, etc.).  These tools
            # don't fire hooks — they only appear in the transcript.
            try:
                hook_call_ids = set(state.pending_tool_starts.keys())
                tool_calls = parse_turn_tool_calls(
                    state.transcript_path,
                    state.turns_processed,
                    hook_call_ids,
                )
                for tc in tool_calls:
                    input_data = (
                        {"tool_input": tc.input_data}
                        if isinstance(tc.input_data, dict)
                        else {"tool_input": tc.input_data}
                    )
                    output_data = (
                        {"tool_response": tc.output_data}
                        if isinstance(tc.output_data, dict)
                        else {"tool_response": tc.output_data}
                    )
                    tool_span = base_span(
                        trace_id=state.current_trace_id,
                        span_id=new_id(),
                        parent_span_id=state.current_turn_span_id,
                        name=f"tool:{tc.name}",
                        kind="tool",
                        timestamp=tc.timestamp or ts_to_iso(turn_start),
                        duration_ms=0.0,  # no timing info for transcript-only tools
                        session_id=session_id,
                        agent_id=_AGENT_ID,
                        input_data=input_data,
                        output_data=output_data,
                        metadata={
                            "tool_name": tc.name,
                            "source": _SOURCE,
                            "cwd": state.cwd or "",
                        },
                    )
                    spans.append(tool_span)
            except Exception:
                logger.debug("staso: transcript tool-call extraction failed", exc_info=True)

            # Only put tokens on the turn span when there are NO LLM child spans
            # (to avoid double-counting in the dashboard trace total).
            total_in, total_out, last_model = _aggregate_llm_messages(llm_messages)
            model = last_model or payload.get("model", "")
            has_llm_children = len(llm_messages) > 0

            turn_span = base_span(
                trace_id=state.current_trace_id,
                span_id=state.current_turn_span_id,
                parent_span_id=None,
                name=turn_span_name(state.current_turn_prompt, state.turn_count),
                kind="agent",
                timestamp=ts_to_iso(turn_start),
                duration_ms=turn_duration_ms,
                session_id=session_id,
                agent_id=_AGENT_ID,
                model=model,
                input_tokens=0 if has_llm_children else total_in,
                output_tokens=0 if has_llm_children else total_out,
                total_tokens=0 if has_llm_children else (total_in + total_out),
                input_data=(
                    {"prompt": state.current_turn_prompt} if state.current_turn_prompt else None
                ),
                output_data=(
                    {"response": payload.get("last_assistant_message", "")}
                    if payload.get("last_assistant_message")
                    else None
                ),
                metadata={
                    "source": _SOURCE,
                    "project": project_name,
                    "cwd": state.cwd,
                    "turn_number": state.turn_count,
                },
            )
            spans.append(turn_span)

            sender.send(spans)

        # Always reset turn state, even if sender is unavailable
        state.current_trace_id = None
        state.current_turn_span_id = None
        state.current_turn_start = None
        state.current_turn_prompt = ""
        state.turns_processed = total_turns
        state.pending_tool_starts.clear()
        save_state(state)
    finally:
        release_lock(session_id)


# -- dispatcher ----------------------------------------------------------------

_HANDLERS = {
    "SessionStart": handle_session_start,
    "UserPromptSubmit": handle_user_prompt_submit,
    "PreToolUse": handle_pre_tool_use,
    "PostToolUse": handle_post_tool_use,
    "Stop": handle_stop,
}

_SENDERLESS_EVENTS = frozenset(("SessionStart", "UserPromptSubmit", "Stop"))


def run(payload: dict[str, Any]) -> None:
    """Dispatch a hook payload to the appropriate handler."""
    event = payload.get("hook_event_name", "")
    handler = _HANDLERS.get(event)
    if handler is None:
        return

    sender = sender_from_env()
    if sender is None and event not in _SENDERLESS_EVENTS:
        return

    try:
        handler(payload, sender)  # type: ignore[arg-type]
    except Exception:
        logger.warning("staso: %s handler error", event, exc_info=True)


def main() -> None:
    """Entry point: read JSON from stdin, run the hook."""
    debug = os.environ.get("STASO_DEBUG", "").lower() in ("1", "true")
    logging.basicConfig(
        level=logging.DEBUG if debug else logging.WARNING,
        format="%(levelname)s staso: %(message)s",
    )
    try:
        payload = json.loads(sys.stdin.read())
    except (json.JSONDecodeError, OSError) as exc:
        logger.warning("staso: failed to read hook payload: %s", exc)
        sys.exit(0)

    run(payload)
