"""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)
    +- Tool span   (kind=tool)  <- timed via PreToolUse/PostToolUse pair

Note: LLM span extraction from Codex transcripts is not yet implemented.
Codex has no SessionEnd event; stale state is GC'd on SessionStart.

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,
)

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


# -- 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:
        state = HookState(session_id=session_id, cwd=cwd)
        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", ""),
        )
        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

        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", "")

        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

        if sender is not None:
            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
            model = payload.get("model", "")

            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_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,
                },
            )
            sender.send([turn_span])

        # 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.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)
