"""Claude Code hook handlers.

Entry point: python3 -m staso.claude_code

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 1 (kind=llm)   <- extracted from transcript in Stop
    |- Tool span   (kind=tool)  <- timed via PreToolUse/PostToolUse pair
    +- LLM span 2 (kind=llm)
"""
from __future__ import annotations

import json
import logging
import os
import sys
import time
import uuid
from datetime import datetime, timezone
from typing import Any

from .sender import ClaudeCodeSender, sender_from_env
from .state import HookState, acquire_lock, delete_state, load_state, release_lock, save_state
from .transcript import AssistantMessage, count_lines, parse_assistant_messages

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

_SOURCE = "claude-code"
_AGENT_ID = os.environ.get("STASO_AGENT_ID", "claude-code")


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


def _new_id() -> str:
    return str(uuid.uuid4())


def _now_iso() -> str:
    return datetime.now(tz=timezone.utc).isoformat()


def _ts_to_iso(unix_ts: float) -> str:
    return datetime.fromtimestamp(unix_ts, tz=timezone.utc).isoformat()


def _parse_iso_to_epoch(iso_str: str) -> float | None:
    """Parse ISO 8601 timestamp to unix epoch seconds. Returns None on failure."""
    if not iso_str:
        return None
    try:
        return datetime.fromisoformat(iso_str).timestamp()
    except (ValueError, TypeError):
        return None


def _safe_json(obj: Any) -> str:
    try:
        return json.dumps(obj, default=str)
    except (TypeError, ValueError):
        logger.debug("staso: failed to serialize object: %s", type(obj).__name__)
        return "{}"


def _base_span(
    *,
    trace_id: str,
    span_id: str,
    parent_span_id: str | None,
    name: str,
    kind: str,
    timestamp: str,
    duration_ms: float,
    session_id: str,
    status: str = "ok",
    model: str = "",
    input_tokens: int = 0,
    output_tokens: int = 0,
    total_tokens: int = 0,
    input_data: dict[str, Any] | None = None,
    output_data: dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    error_message: str | None = None,
) -> dict[str, Any]:
    return {
        "trace_id": trace_id,
        "span_id": span_id,
        "parent_span_id": parent_span_id,
        "name": name,
        "kind": kind,
        "timestamp": timestamp,
        "duration_ms": round(duration_ms, 2),
        "status": status,
        "model": model,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens,
        "input": _safe_json(input_data) if input_data else "",
        "output": _safe_json(output_data) if output_data else "",
        "attributes": _safe_json(metadata) if metadata else "{}",
        "error_message": error_message,
        "session_id": session_id,
        "agent_id": _AGENT_ID,
        "user_id": "",
    }


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 _aggregate_llm_messages(
    messages: list[AssistantMessage],
) -> 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


def _turn_span_name(prompt: str, turn_count: int) -> str:
    """Generate a human-readable turn span name from the user prompt."""
    if not prompt or not prompt.strip():
        return f"turn {turn_count}"
    truncated = prompt.strip().replace("\n", " ")
    if len(truncated) > 80:
        truncated = truncated[:77] + "..."
    return truncated


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


def _auto_repair_hooks(cwd: str) -> None:
    """Ensure all current hook events are registered in settings.

    Called on SessionStart so that upgrading the staso package is the only
    step needed -- the next session auto-adds any new hook events.
    """
    from pathlib import Path as _Path

    from ..cli import _build_hooks_config, _HOOK_EVENTS, _is_staso_hook, _merge_settings

    # Find settings file that contains staso hooks (project-level first)
    candidates = [
        _Path(cwd) / ".claude" / "settings.local.json",
        _Path.home() / ".claude" / "settings.json",
    ]
    settings_path = None
    existing: dict = {}
    for path in candidates:
        if not path.exists():
            continue
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        hooks = data.get("hooks", {})
        if any(_is_staso_hook(e) for entries in hooks.values() for e in entries):
            settings_path = path
            existing = data
            break

    if settings_path is None:
        return

    # Check if all events have staso hooks registered
    existing_hooks = existing.get("hooks", {})
    all_present = all(
        any(_is_staso_hook(e) for e in existing_hooks.get(event, []))
        for event in _HOOK_EVENTS
    )
    if all_present:
        return

    # Extract command from any existing staso hook entry
    command = None
    for event_entries in existing_hooks.values():
        for entry in event_entries:
            if not _is_staso_hook(entry):
                continue
            for h in entry.get("hooks", [entry]):
                cmd = h.get("command", "")
                if "staso.claude_code" in cmd:
                    command = cmd
                    break
            if command:
                break
    if not command:
        return

    merged = _merge_settings(existing, _build_hooks_config(command), {})
    try:
        settings_path.write_text(json.dumps(merged, indent=2), encoding="utf-8")
        logger.debug("staso: auto-repaired hooks config (%s)", settings_path)
    except OSError:
        pass


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

    try:
        _auto_repair_hooks(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: ClaudeCodeSender) -> 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", "")
        state.last_transcript_line = count_lines(payload.get("transcript_path", ""))
        save_state(state)
    finally:
        release_lock(session_id)


def handle_pre_tool_use(payload: dict[str, Any], _sender: ClaudeCodeSender) -> 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: ClaudeCodeSender) -> 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,
            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_post_tool_use_failure(payload: dict[str, Any], sender: ClaudeCodeSender) -> 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")
        error = payload.get("error", "")

        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,
            status="error",
            input_data={"tool_input": payload.get("tool_input") or {}},
            metadata={"tool_name": tool_name, "source": _SOURCE},
            error_message=error,
        )
        sender.send([span])
        save_state(state)
    finally:
        release_lock(session_id)


def handle_stop(payload: dict[str, Any], sender: ClaudeCodeSender | 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

        transcript_path = payload.get("transcript_path", "")
        llm_messages, new_line = parse_assistant_messages(
            transcript_path, state.last_transcript_line
        )

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

            for i, msg in enumerate(llm_messages):
                ts = msg.timestamp if msg.timestamp else _ts_to_iso(turn_start)
                ts_epoch = _parse_iso_to_epoch(ts)
                if i + 1 < len(llm_messages):
                    next_ts = llm_messages[i + 1].timestamp
                    next_epoch = _parse_iso_to_epoch(next_ts)
                    if ts_epoch and next_epoch:
                        llm_duration_ms = max((next_epoch - ts_epoch) * 1000, 0.0)
                    else:
                        llm_duration_ms = 0.0
                else:
                    llm_duration_ms = max((turn_end - ts_epoch) * 1000, 0.0) if ts_epoch else 0.0

                is_synthetic = msg.model.startswith("<")
                if is_synthetic:
                    content_text = " ".join(
                        b.get("text", "") for b in msg.content if b.get("type") == "text"
                    ).strip()
                    llm_span = _base_span(
                        trace_id=state.current_trace_id,
                        span_id=_new_id(),
                        parent_span_id=state.current_turn_span_id,
                        name="internal",
                        kind="custom",
                        timestamp=ts,
                        duration_ms=llm_duration_ms,
                        session_id=session_id,
                        output_data={"content": msg.content} if msg.content else None,
                        metadata={
                            "source": _SOURCE,
                            "synthetic": True,
                            "synthetic_content": content_text,
                        },
                    )
                else:
                    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:claude",
                        kind="llm",
                        timestamp=ts,
                        duration_ms=llm_duration_ms,
                        session_id=session_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,
                            "response_id": msg.message_id,
                            "cache_read_input_tokens": msg.cache_read_tokens,
                            "cache_creation_input_tokens": msg.cache_creation_tokens,
                        },
                    )
                spans.append(llm_span)

            total_in, total_out, last_model = _aggregate_llm_messages(llm_messages)

            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,
                model=last_model,
                input_tokens=total_in,
                output_tokens=total_out,
                total_tokens=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.last_transcript_line = new_line
        state.pending_tool_starts.clear()
        state.pending_subagent_starts.clear()
        save_state(state)
    finally:
        release_lock(session_id)


def handle_stop_failure(payload: dict[str, Any], sender: ClaudeCodeSender | 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
            error_type = payload.get("error", "unknown")
            error_details = payload.get("error_details", "")

            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,
                status="error",
                input_data={"prompt": state.current_turn_prompt} if state.current_turn_prompt else None,
                metadata={
                    "source": _SOURCE,
                    "project": project_name,
                    "cwd": state.cwd,
                    "turn_number": state.turn_count,
                    "error_type": error_type,
                },
                error_message=f"{error_type}: {error_details}" if error_details else error_type,
            )
            sender.send([turn_span])

        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()
        state.pending_subagent_starts.clear()
        save_state(state)
    finally:
        release_lock(session_id)


def handle_subagent_start(payload: dict[str, Any], _sender: ClaudeCodeSender) -> 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
        agent_id = payload.get("agent_id", "")
        if agent_id:
            state.pending_subagent_starts[agent_id] = time.time()
            save_state(state)
    finally:
        release_lock(session_id)


def handle_subagent_stop(payload: dict[str, Any], sender: ClaudeCodeSender) -> 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()
        agent_id = payload.get("agent_id", "")
        start_time = state.pending_subagent_starts.pop(agent_id, None)
        duration_ms = (now - start_time) * 1000 if start_time else 0.0

        transcript_path = payload.get("agent_transcript_path") or payload.get("transcript_path", "")
        llm_messages, _ = parse_assistant_messages(transcript_path, 0)

        total_in, total_out, last_model = _aggregate_llm_messages(llm_messages)

        agent_type = payload.get("agent_type", "")

        span = _base_span(
            trace_id=state.current_trace_id,
            span_id=_new_id(),
            parent_span_id=state.current_turn_span_id,
            name=f"subagent:{agent_type}" if agent_type else "subagent",
            kind="agent",
            timestamp=_ts_to_iso(start_time) if start_time else _now_iso(),
            duration_ms=duration_ms,
            session_id=session_id,
            model=last_model,
            input_tokens=total_in,
            output_tokens=total_out,
            total_tokens=total_in + total_out,
            metadata={
                "source": _SOURCE,
                "subagent": True,
                "agent_id": agent_id,
                "agent_type": agent_type,
            },
        )
        sender.send([span])
        save_state(state)
    finally:
        release_lock(session_id)


def handle_session_end(payload: dict[str, Any], _sender: ClaudeCodeSender) -> None:
    delete_state(payload["session_id"])


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

_HANDLERS = {
    "SessionStart": handle_session_start,
    "UserPromptSubmit": handle_user_prompt_submit,
    "PreToolUse": handle_pre_tool_use,
    "PostToolUse": handle_post_tool_use,
    "PostToolUseFailure": handle_post_tool_use_failure,
    "Stop": handle_stop,
    "StopFailure": handle_stop_failure,
    "SubagentStart": handle_subagent_start,
    "SubagentStop": handle_subagent_stop,
    "SessionEnd": handle_session_end,
}


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 ("SessionStart", "UserPromptSubmit", "Stop", "StopFailure", "SessionEnd"):
        return

    try:
        handler(payload, sender)  # type: ignore[arg-type]
    except Exception:
        logger.debug("staso: hook handler error", 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)
