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

import json
from unittest.mock import MagicMock, patch

import pytest

from staso.claude_code.hook import (
    _turn_span_name,
    handle_pre_tool_use,
    handle_post_tool_use,
    handle_post_tool_use_failure,
    handle_session_end,
    handle_session_start,
    handle_stop,
    handle_stop_failure,
    handle_subagent_start,
    handle_subagent_stop,
    handle_user_prompt_submit,
    run,
)
from staso.claude_code.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.claude_code.state._state_dir", lambda: tmp_path)

        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


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

        # Pre-create state (as SessionStart would)
        save_state(HookState(session_id="s1", cwd="/proj"))

        handle_user_prompt_submit(
            {
                "hook_event_name": "UserPromptSubmit",
                "session_id": "s1",
                "prompt": "Fix the bug",
                "transcript_path": str(tmp_path / "nonexistent.jsonl"),
                "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.claude_code.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",
             "transcript_path": str(tmp_path / "x.jsonl"), "cwd": "/"},
            _make_sender(),
        )
        state = load_state("s2")
        assert state.turn_count == 3


class TestPreToolUse:
    def test_stores_start_time(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.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": "Read",
                "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.claude_code.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": "Read",
                "tool_input": {"file_path": "/src/auth.py"},
                "tool_response": {"content": "file contents"},
            },
            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:Read"
        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.claude_code.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.claude_code.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 TestPostToolUseFailure:
    def test_sends_error_span(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

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

        sender = _make_sender()
        handle_post_tool_use_failure(
            {
                "hook_event_name": "PostToolUseFailure",
                "session_id": "s1",
                "tool_name": "Bash",
                "error": "command not found",
            },
            sender,
        )

        spans = sender.send.call_args[0][0]
        assert spans[0]["status"] == "error"
        assert spans[0]["error_message"] == "command not found"


class TestStop:
    def test_sends_turn_and_llm_spans(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        transcript = tmp_path / "session.jsonl"
        transcript.write_text(
            json.dumps({
                "type": "assistant",
                "message": {
                    "id": "msg_1",
                    "model": "claude-sonnet-4-6",
                    "content": [{"type": "text", "text": "done"}],
                    "usage": {"input_tokens": 100, "output_tokens": 50},
                },
                "timestamp": "2026-01-01T00:00:01Z",
            }),
            encoding="utf-8",
        )

        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"
        state.last_transcript_line = 0
        save_state(state)

        sender = _make_sender()
        handle_stop(
            {
                "hook_event_name": "Stop",
                "session_id": "s1",
                "transcript_path": str(transcript),
                "last_assistant_message": "done",
            },
            sender,
        )

        sender.send.assert_called_once()
        spans = sender.send.call_args[0][0]
        kinds = {s["kind"] for s in spans}
        assert "agent" in kinds
        assert "llm" in kinds

        turn_span = next(s for s in spans if s["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["input_tokens"] == 100
        assert turn_span["output_tokens"] == 50

        llm_span = next(s for s in spans if s["kind"] == "llm")
        assert llm_span["model"] == "claude-sonnet-4-6"
        assert llm_span["parent_span_id"] == "turn-span-1"

    def test_resets_turn_state_after_stop(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        transcript = tmp_path / "t.jsonl"
        transcript.write_text("", encoding="utf-8")

        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",
             "transcript_path": str(transcript), "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.claude_code.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",
             "transcript_path": "", "last_assistant_message": ""},
            sender,
        )
        sender.send.assert_not_called()

    def test_synthetic_message_produces_custom_span(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        transcript = tmp_path / "session.jsonl"
        transcript.write_text(
            json.dumps({
                "type": "assistant",
                "message": {
                    "id": "msg_syn",
                    "model": "<synthetic>",
                    "content": [{"type": "text", "text": "context compacted"}],
                    "usage": {"input_tokens": 0, "output_tokens": 0},
                },
                "timestamp": "2026-01-01T00:00:01Z",
            }),
            encoding="utf-8",
        )

        state = HookState(session_id="s3", cwd="/proj", turn_count=1)
        state.current_trace_id = "trace-syn"
        state.current_turn_span_id = "turn-span-syn"
        state.current_turn_start = 1000.0
        state.last_transcript_line = 0
        save_state(state)

        sender = _make_sender()
        handle_stop(
            {
                "hook_event_name": "Stop",
                "session_id": "s3",
                "transcript_path": str(transcript),
                "last_assistant_message": "",
            },
            sender,
        )

        spans = sender.send.call_args[0][0]
        custom_span = next(s for s in spans if s["kind"] == "custom")
        assert custom_span["name"] == "internal"
        assert custom_span["model"] == ""
        assert custom_span["input_tokens"] == 0
        assert custom_span["output_tokens"] == 0
        metadata = json.loads(custom_span["attributes"])
        assert metadata["synthetic"] is True
        assert metadata["synthetic_content"] == "context compacted"
        assert not any(s["kind"] == "llm" for s in spans)


class TestStopWithoutSender:
    def test_resets_state_when_sender_is_none(self, tmp_path, monkeypatch):
        """State must reset even if sender is unavailable (e.g. API key removed mid-session)."""
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        transcript = tmp_path / "t.jsonl"
        transcript.write_text("", encoding="utf-8")

        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",
             "transcript_path": str(transcript), "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


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"

    def test_turn_span_name_in_stop(self, tmp_path, monkeypatch):
        """The turn span name should use the user prompt, not turn number + project."""
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        transcript = tmp_path / "t.jsonl"
        transcript.write_text("", encoding="utf-8")

        state = HookState(session_id="s1", cwd="/proj", turn_count=1)
        state.current_trace_id = "t"
        state.current_turn_span_id = "sp"
        state.current_turn_start = 0.0
        state.current_turn_prompt = "Fix the login bug"
        save_state(state)

        sender = _make_sender()
        handle_stop(
            {"hook_event_name": "Stop", "session_id": "s1",
             "transcript_path": str(transcript), "last_assistant_message": "done"},
            sender,
        )

        spans = sender.send.call_args[0][0]
        turn_span = next(s for s in spans if s["kind"] == "agent")
        assert turn_span["name"] == "Fix the login bug"
        metadata = json.loads(turn_span["attributes"])
        assert metadata["turn_number"] == 1
        assert metadata["project"] == "proj"


class TestStopFailure:
    def test_sends_error_turn_span(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        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_failure(
            {
                "hook_event_name": "StopFailure",
                "session_id": "s1",
                "error": "rate_limit",
                "error_details": "Too many requests",
                "last_assistant_message": "",
            },
            sender,
        )

        sender.send.assert_called_once()
        spans = sender.send.call_args[0][0]
        assert len(spans) == 1
        span = spans[0]
        assert span["status"] == "error"
        assert "rate_limit" in span["error_message"]
        metadata = json.loads(span["attributes"])
        assert metadata["error_type"] == "rate_limit"

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


class TestSubagentStart:
    def test_stores_start_time(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.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_subagent_start(
            {
                "hook_event_name": "SubagentStart",
                "session_id": "s1",
                "agent_id": "agent-abc",
                "agent_type": "Explore",
            },
            _make_sender(),
        )

        updated = load_state("s1")
        assert "agent-abc" in updated.pending_subagent_starts
        assert updated.pending_subagent_starts["agent-abc"] > 0


class TestSubagentStop:
    def test_sends_subagent_span_with_duration(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)
        import time

        state = HookState(session_id="s1", cwd="/proj")
        state.current_trace_id = "trace-1"
        state.current_turn_span_id = "turn-span-1"
        state.pending_subagent_starts["agent-abc"] = time.time() - 5.0
        save_state(state)

        transcript = tmp_path / "subagent.jsonl"
        transcript.write_text(
            json.dumps({
                "type": "assistant",
                "message": {
                    "id": "msg_sub",
                    "model": "claude-opus-4-6",
                    "content": [{"type": "text", "text": "found it"}],
                    "usage": {"input_tokens": 500, "output_tokens": 200},
                },
                "timestamp": "2026-01-01T00:00:01Z",
            }),
            encoding="utf-8",
        )

        sender = _make_sender()
        handle_subagent_stop(
            {
                "hook_event_name": "SubagentStop",
                "session_id": "s1",
                "agent_id": "agent-abc",
                "agent_type": "Explore",
                "agent_transcript_path": str(transcript),
            },
            sender,
        )

        sender.send.assert_called_once()
        span = sender.send.call_args[0][0][0]
        assert span["kind"] == "agent"
        assert span["name"] == "subagent:Explore"
        assert span["duration_ms"] >= 4900
        assert span["input_tokens"] == 500
        assert span["output_tokens"] == 200
        metadata = json.loads(span["attributes"])
        assert metadata["subagent"] is True
        assert metadata["agent_type"] == "Explore"

        updated = load_state("s1")
        assert "agent-abc" not in updated.pending_subagent_starts


class TestSessionEnd:
    def test_deletes_state(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.claude_code.state._state_dir", lambda: tmp_path)

        save_state(HookState(session_id="s1", cwd="/"))
        handle_session_end({"session_id": "s1"}, _make_sender())
        assert load_state("s1") is None


class TestDispatcher:
    def test_unknown_event_is_noop(self, tmp_path, monkeypatch, capsys):
        monkeypatch.setenv("STASO_API_KEY", "ak_test")
        with patch("staso.claude_code.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_missing_session_id_handled_gracefully(self, monkeypatch):
        monkeypatch.setenv("STASO_API_KEY", "ak_test")
        run({"hook_event_name": "SessionEnd"})
