"""Tests for codex.hook -- event handlers and dispatcher."""
from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pytest

from staso._common.helpers import turn_span_name
from staso.codex.hook import (
    handle_pre_tool_use,
    handle_post_tool_use,
    handle_session_start,
    handle_stop,
    handle_user_prompt_submit,
    run,
)
from staso.codex.state import HookState, load_state, save_state


def _make_sender() -> MagicMock:
    sender = MagicMock()
    sender.send = MagicMock()
    return sender


class TestSessionStart:
    def test_creates_state_file(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        monkeypatch.setattr("staso.codex.hook.find_transcript", lambda sid: None)

        handle_session_start(
            {"hook_event_name": "SessionStart", "session_id": "s1", "cwd": "/home/user/proj"},
            _make_sender(),
        )
        state = load_state("s1")
        assert state is not None
        assert state.session_id == "s1"
        assert state.cwd == "/home/user/proj"
        assert state.turn_count == 0

    def test_caches_transcript_path(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        from pathlib import Path
        fake_path = Path("/fake/sessions/transcript.jsonl")
        monkeypatch.setattr("staso.codex.hook.find_transcript", lambda sid: fake_path)

        handle_session_start(
            {"hook_event_name": "SessionStart", "session_id": "s2", "cwd": "/proj"},
            _make_sender(),
        )
        state = load_state("s2")
        assert state.transcript_path == str(fake_path)


class TestUserPromptSubmit:
    def test_initialises_turn(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        save_state(HookState(session_id="s1", cwd="/proj"))

        handle_user_prompt_submit(
            {
                "hook_event_name": "UserPromptSubmit",
                "session_id": "s1",
                "prompt": "Fix the bug",
                "cwd": "/proj",
            },
            _make_sender(),
        )

        state = load_state("s1")
        assert state is not None
        assert state.turn_count == 1
        assert state.current_trace_id is not None
        assert state.current_turn_span_id is not None
        assert state.current_turn_start is not None
        assert state.current_turn_prompt == "Fix the bug"

    def test_increments_turn_count(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        save_state(HookState(session_id="s2", cwd="/", turn_count=2))
        handle_user_prompt_submit(
            {"hook_event_name": "UserPromptSubmit", "session_id": "s2", "prompt": "p", "cwd": "/"},
            _make_sender(),
        )
        state = load_state("s2")
        assert state.turn_count == 3

    def test_discovers_transcript_if_missing(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        save_state(HookState(session_id="s3", cwd="/proj", transcript_path=""))

        from pathlib import Path
        fake_path = Path("/found/transcript.jsonl")
        monkeypatch.setattr("staso.codex.hook.find_transcript", lambda sid: fake_path)

        handle_user_prompt_submit(
            {"hook_event_name": "UserPromptSubmit", "session_id": "s3", "prompt": "hi"},
            _make_sender(),
        )
        state = load_state("s3")
        assert state.transcript_path == str(fake_path)


class TestPreToolUse:
    def test_stores_start_time(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        state = HookState(session_id="s1", cwd="/proj")
        state.current_trace_id = "trace-1"
        state.current_turn_span_id = "turn-span-1"
        save_state(state)

        handle_pre_tool_use(
            {
                "hook_event_name": "PreToolUse",
                "session_id": "s1",
                "tool_name": "Bash",
                "tool_use_id": "tu-1",
            },
            _make_sender(),
        )

        updated = load_state("s1")
        assert "tu-1" in updated.pending_tool_starts
        assert updated.pending_tool_starts["tu-1"] > 0


class TestPostToolUse:
    def _state_with_turn(self, session_id: str) -> HookState:
        state = HookState(session_id=session_id, cwd="/proj")
        state.current_trace_id = "trace-1"
        state.current_turn_span_id = "turn-span-1"
        return state

    def test_sends_tool_span(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        save_state(self._state_with_turn("s1"))

        sender = _make_sender()
        handle_post_tool_use(
            {
                "hook_event_name": "PostToolUse",
                "session_id": "s1",
                "tool_name": "Bash",
                "tool_input": {"command": "ls -la"},
                "tool_response": {"stdout": "file1 file2"},
            },
            sender,
        )

        sender.send.assert_called_once()
        spans = sender.send.call_args[0][0]
        assert len(spans) == 1
        span = spans[0]
        assert span["kind"] == "tool"
        assert span["name"] == "tool:Bash"
        assert span["parent_span_id"] == "turn-span-1"
        assert span["trace_id"] == "trace-1"
        assert span["status"] == "ok"

    def test_precise_duration_from_pre_tool_use(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        import time

        state = self._state_with_turn("s1")
        state.pending_tool_starts["tu-1"] = time.time() - 2.5
        save_state(state)

        sender = _make_sender()
        handle_post_tool_use(
            {
                "hook_event_name": "PostToolUse",
                "session_id": "s1",
                "tool_name": "Bash",
                "tool_use_id": "tu-1",
                "tool_input": {"command": "npm test"},
                "tool_response": {"stdout": "ok"},
            },
            sender,
        )

        span = sender.send.call_args[0][0][0]
        assert span["duration_ms"] >= 2400

        updated = load_state("s1")
        assert "tu-1" not in updated.pending_tool_starts

    def test_no_op_without_active_turn(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        save_state(HookState(session_id="s2", cwd="/"))

        sender = _make_sender()
        handle_post_tool_use(
            {"hook_event_name": "PostToolUse", "session_id": "s2", "tool_name": "Bash"},
            sender,
        )
        sender.send.assert_not_called()


class TestStop:
    def test_sends_turn_span(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        monkeypatch.setattr(
            "staso.codex.hook.parse_turn_llm_data", lambda path, start: ([], start)
        )

        state = HookState(session_id="s1", cwd="/proj", turn_count=1)
        state.current_trace_id = "trace-1"
        state.current_turn_span_id = "turn-span-1"
        state.current_turn_start = 1000.0
        state.current_turn_prompt = "Do something"
        save_state(state)

        sender = _make_sender()
        handle_stop(
            {
                "hook_event_name": "Stop",
                "session_id": "s1",
                "model": "gpt-5.4",
                "last_assistant_message": "done",
            },
            sender,
        )

        sender.send.assert_called_once()
        spans = sender.send.call_args[0][0]
        assert len(spans) == 1  # just turn span, no LLM spans
        turn_span = spans[0]
        assert turn_span["kind"] == "agent"
        assert turn_span["trace_id"] == "trace-1"
        assert turn_span["span_id"] == "turn-span-1"
        assert turn_span["parent_span_id"] is None
        assert turn_span["model"] == "gpt-5.4"
        assert turn_span["name"] == "Do something"

        metadata = json.loads(turn_span["attributes"])
        assert metadata["source"] == "codex"
        assert metadata["turn_number"] == 1
        assert metadata["project"] == "proj"

    def test_emits_llm_span_from_transcript(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        from staso.codex.transcript import LLMMessage

        fake_msg = LLMMessage(
            model="gpt-5.4",
            input_tokens=12000,
            output_tokens=50,
            cached_input_tokens=9000,
            reasoning_output_tokens=10,
            content=[{"type": "output_text", "text": "Hello!"}],
            timestamp="2026-04-02T15:42:41.000Z",
        )
        monkeypatch.setattr(
            "staso.codex.hook.parse_turn_llm_data", lambda path, start: ([fake_msg], start + 10)
        )

        state = HookState(session_id="s2", cwd="/proj", turn_count=1)
        state.current_trace_id = "trace-2"
        state.current_turn_span_id = "turn-span-2"
        state.current_turn_start = 1000.0
        state.current_turn_prompt = "Hello"
        save_state(state)

        sender = _make_sender()
        handle_stop(
            {
                "hook_event_name": "Stop",
                "session_id": "s2",
                "last_assistant_message": "Hello!",
            },
            sender,
        )

        sender.send.assert_called_once()
        spans = sender.send.call_args[0][0]
        assert len(spans) == 2  # LLM span + turn span

        llm_span = spans[0]
        assert llm_span["kind"] == "llm"
        assert llm_span["name"] == "llm:gpt-5.4"
        assert llm_span["model"] == "gpt-5.4"
        assert llm_span["input_tokens"] == 12000
        assert llm_span["output_tokens"] == 50
        assert llm_span["total_tokens"] == 12050
        assert llm_span["parent_span_id"] == "turn-span-2"

        llm_meta = json.loads(llm_span["attributes"])
        assert llm_meta["cached_input_tokens"] == 9000
        assert llm_meta["reasoning_output_tokens"] == 10

        turn_span = spans[1]
        assert turn_span["kind"] == "agent"
        assert turn_span["model"] == "gpt-5.4"
        assert turn_span["input_tokens"] == 12000
        assert turn_span["output_tokens"] == 50

    def test_resets_turn_state_after_stop(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        monkeypatch.setattr(
            "staso.codex.hook.parse_turn_llm_data", lambda path, start: ([], start)
        )

        state = HookState(session_id="s1", cwd="/", turn_count=1)
        state.current_trace_id = "t"
        state.current_turn_span_id = "sp"
        state.current_turn_start = 0.0
        save_state(state)

        handle_stop(
            {"hook_event_name": "Stop", "session_id": "s1", "last_assistant_message": ""},
            _make_sender(),
        )

        updated = load_state("s1")
        assert updated.current_trace_id is None
        assert updated.current_turn_span_id is None

    def test_no_op_without_active_turn(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)

        save_state(HookState(session_id="s2", cwd="/"))
        sender = _make_sender()
        handle_stop(
            {"hook_event_name": "Stop", "session_id": "s2", "last_assistant_message": ""},
            sender,
        )
        sender.send.assert_not_called()

    def test_resets_state_when_sender_is_none(self, tmp_path, monkeypatch):
        """State must reset even if sender is unavailable."""
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        monkeypatch.setattr(
            "staso.codex.hook.parse_turn_llm_data", lambda path, start: ([], start)
        )

        state = HookState(session_id="s1", cwd="/proj", turn_count=1)
        state.current_trace_id = "trace-orphan"
        state.current_turn_span_id = "span-orphan"
        state.current_turn_start = 1000.0
        save_state(state)

        handle_stop(
            {"hook_event_name": "Stop", "session_id": "s1", "last_assistant_message": ""},
            None,
        )

        updated = load_state("s1")
        assert updated.current_trace_id is None
        assert updated.current_turn_span_id is None
        assert updated.current_turn_start is None

    def test_updates_last_transcript_line(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.state._state_dir", lambda: tmp_path)
        monkeypatch.setattr(
            "staso.codex.hook.parse_turn_llm_data", lambda path, start: ([], 42)
        )

        state = HookState(session_id="s1", cwd="/", turn_count=1, last_transcript_line=10)
        state.current_trace_id = "t"
        state.current_turn_span_id = "sp"
        state.current_turn_start = 0.0
        save_state(state)

        handle_stop(
            {"hook_event_name": "Stop", "session_id": "s1", "last_assistant_message": ""},
            _make_sender(),
        )

        updated = load_state("s1")
        assert updated.last_transcript_line == 42


class TestTurnSpanName:
    def test_uses_prompt_as_name(self):
        assert turn_span_name("Fix the bug in auth.py", 1) == "Fix the bug in auth.py"

    def test_truncates_long_prompt(self):
        long_prompt = "a" * 100
        result = turn_span_name(long_prompt, 1)
        assert len(result) == 80
        assert result.endswith("...")

    def test_falls_back_to_turn_number(self):
        assert turn_span_name("", 3) == "turn 3"
        assert turn_span_name("   ", 5) == "turn 5"
        assert turn_span_name(None, 2) == "turn 2"

    def test_collapses_newlines(self):
        assert turn_span_name("line one\nline two", 1) == "line one line two"


class TestDispatcher:
    def test_unknown_event_is_noop(self, tmp_path, monkeypatch):
        monkeypatch.setenv("STASO_API_KEY", "ak_test")
        with patch("staso.codex.hook.sender_from_env") as mock_sender:
            run({"hook_event_name": "UnknownEvent", "session_id": "s1"})
            mock_sender.assert_not_called()

    def test_no_api_key_skips_data_events(self, monkeypatch):
        monkeypatch.delenv("STASO_API_KEY", raising=False)
        run({"hook_event_name": "PostToolUse", "session_id": "s1", "tool_name": "Bash"})

    def test_no_session_end_event(self):
        """Codex has no SessionEnd event."""
        from staso.codex.hook import _HANDLERS
        assert "SessionEnd" not in _HANDLERS

    def test_no_subagent_events(self):
        """Codex has no subagent events."""
        from staso.codex.hook import _HANDLERS
        assert "SubagentStart" not in _HANDLERS
        assert "SubagentStop" not in _HANDLERS

    def test_no_failure_events(self):
        """Codex has no StopFailure or PostToolUseFailure events."""
        from staso.codex.hook import _HANDLERS
        assert "StopFailure" not in _HANDLERS
        assert "PostToolUseFailure" not in _HANDLERS
