from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from typer.testing import CliRunner

from haiku.rag.cli import _cli as cli

runner = CliRunner()


def test_chat_command():
    """Test chat command launches chat TUI."""
    with patch("haiku.rag.chat.run_chat") as mock_chat:
        mock_chat.return_value = None

        result = runner.invoke(cli, ["chat"])

        assert result.exit_code == 0
        mock_chat.assert_called_once()


@pytest.mark.asyncio
async def test_chat_app_has_required_widgets(temp_db_path: Path):
    """Test that ChatApp has the required widgets: ChatHistory, Input."""
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            chat_history = app.query_one(ChatHistory)
            assert chat_history is not None

            from textual.widgets import Input

            chat_input = app.query_one(Input)
            assert chat_input is not None


@pytest.mark.asyncio
async def test_chat_app_quit_binding(temp_db_path: Path):
    """Test that pressing ctrl+q quits the app."""
    from haiku.rag.chat.app import ChatApp

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test() as pilot:
            # App should be running
            assert app.is_running

            # Press ctrl+q to quit
            await pilot.press("ctrl+q")

            # App should have exited
            assert not app.is_running


@pytest.mark.asyncio
async def test_chat_history_can_add_message(temp_db_path: Path):
    """Test that ChatHistory can display messages."""
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            chat_history = app.query_one(ChatHistory)

            # Add a user message
            await chat_history.add_message("user", "Hello, how are you?")
            assert len(chat_history.messages) == 1
            assert chat_history.messages[0] == ("user", "Hello, how are you?")

            # Add an assistant message
            await chat_history.add_message("assistant", "I'm doing well, thank you!")
            assert len(chat_history.messages) == 2


@pytest.mark.asyncio
async def test_chat_app_calls_agent_on_submit(temp_db_path: Path):
    """Test that submitting a message triggers agent invocation."""
    from contextlib import asynccontextmanager

    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    # Create a mock stream result
    mock_result = MagicMock()
    mock_result.output = "This is the agent's response."

    # Mock stream that returns the result
    mock_stream = MagicMock()
    mock_stream.get_result = AsyncMock(return_value=mock_result)

    async def mock_stream_text():
        yield "This is the agent's response."

    mock_stream.stream_text = mock_stream_text

    @asynccontextmanager
    async def mock_run_stream(*args, **kwargs):
        yield mock_stream

    # Create a mock agent that returns a fixed response
    mock_agent = MagicMock()
    mock_agent.run_stream = mock_run_stream

    with (
        patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client),
        patch("haiku.rag.chat.app.create_chat_agent", return_value=mock_agent),
    ):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test() as pilot:
            # Type a message in the input
            await pilot.press("H", "e", "l", "l", "o")
            await pilot.press("enter")

            # Give the app time to process
            await pilot.pause()

            # Verify the response was added to chat history
            chat_history = app.query_one(ChatHistory)
            assert len(chat_history.messages) >= 2  # User message + agent response


@pytest.mark.asyncio
async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
    """Test that ChatHistory can display inline tool calls."""
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory, ToolCallWidget

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            chat_history = app.query_one(ChatHistory)

            # Add a tool call
            tool_widget = await chat_history.add_tool_call("search", {"query": "test"})
            assert isinstance(tool_widget, ToolCallWidget)
            assert tool_widget._complete is False

            # Mark it complete
            chat_history.mark_tool_complete(tool_widget)
            assert tool_widget._complete is True


@pytest.mark.asyncio
async def test_chat_history_can_add_citations(temp_db_path: Path):
    """Test that ChatHistory can display inline citations."""
    from haiku.rag.agents.research.models import Citation
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            chat_history = app.query_one(ChatHistory)

            test_citations = [
                Citation(
                    index=1,
                    document_id="doc1",
                    chunk_id="chunk1",
                    document_uri="file:///test/doc1.pdf",
                    document_title="Test Document 1",
                    page_numbers=[1, 2],
                    headings=["Section 1"],
                    content="This is some test content from doc 1",
                ),
                Citation(
                    index=2,
                    document_id="doc2",
                    chunk_id="chunk2",
                    document_uri="file:///test/doc2.pdf",
                    document_title="Test Document 2",
                    page_numbers=[5],
                    headings=["Section 2", "Subsection"],
                    content="This is test content from doc 2",
                ),
            ]

            await chat_history.add_citations(test_citations)

            # Verify citation widgets were added
            citation_widgets = chat_history.query(CitationWidget)
            assert len(list(citation_widgets)) == 2


@pytest.mark.asyncio
async def test_chat_history_thinking_indicator(temp_db_path: Path):
    """Test that ChatHistory can show and hide thinking indicator."""
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory, ThinkingWidget

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test() as pilot:
            chat_history = app.query_one(ChatHistory)

            # Show thinking indicator
            await chat_history.show_thinking()
            thinking = chat_history.query(ThinkingWidget)
            assert len(list(thinking)) == 1

            # Hide thinking indicator (remove is deferred, so pause)
            chat_history.hide_thinking()
            await pilot.pause()
            thinking = chat_history.query(ThinkingWidget)
            assert len(list(thinking)) == 0


@pytest.mark.asyncio
async def test_clear_chat_resets_session(temp_db_path: Path):
    """Test that clearing chat resets the session state."""
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test() as pilot:
            chat_history = app.query_one(ChatHistory)

            # Add some messages
            await chat_history.add_message("user", "Hello")
            await chat_history.add_message("assistant", "Hi there")
            assert len(chat_history.messages) == 2

            # Store original session ID
            assert app.session_state is not None
            original_session_id = app.session_state.session_id

            # Clear chat via action (available through command palette)
            await app.action_clear_chat()
            await pilot.pause()

            # Verify messages cleared
            assert len(chat_history.messages) == 0

            # Verify session state reset (new session ID)
            assert app.session_state is not None
            assert app.session_state.session_id != original_session_id


@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_snapshot(
    temp_db_path: Path,
):
    """Test that _handle_stream_event extracts citations from STATE_SNAPSHOT events."""
    from ag_ui.core import EventType, StateSnapshotEvent
    from pydantic_ai import FunctionToolResultEvent
    from pydantic_ai.messages import ToolReturnPart

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
    from haiku.rag.chat.app import ChatApp

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            # Create a STATE_SNAPSHOT event with citations
            snapshot_event = StateSnapshotEvent(
                type=EventType.STATE_SNAPSHOT,
                snapshot={
                    AGUI_STATE_KEY: {
                        "session_id": "test",
                        "citations": [
                            {
                                "index": 1,
                                "document_id": "doc1",
                                "chunk_id": "chunk1",
                                "document_uri": "test.pdf",
                                "document_title": "Test Doc",
                                "content": "Test content",
                            }
                        ],
                        "qa_history": [],
                        "citation_registry": {"chunk1": 1},
                    }
                },
            )

            # Create a tool result with the snapshot in metadata
            tool_return = ToolReturnPart(
                tool_name="search",
                content="Found results",
                tool_call_id="test-call-1",
                metadata=[snapshot_event],
            )

            event = FunctionToolResultEvent(result=tool_return)

            # Handle the event
            await app._handle_stream_event(event)

            # Session state should be synced
            assert len(app.session_state.citations) == 1
            assert app.session_state.citations[0].chunk_id == "chunk1"
            assert app.session_state.citation_registry == {"chunk1": 1}


@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_delta(
    temp_db_path: Path,
):
    """Test that _handle_stream_event extracts citations from STATE_DELTA events.

    This test demonstrates that state deltas need to be applied to extract citations.
    After the first STATE_SNAPSHOT, subsequent tool calls emit STATE_DELTA events
    containing JSON Patch operations.
    """
    from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
    from pydantic_ai import FunctionToolResultEvent
    from pydantic_ai.messages import ToolReturnPart

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
    from haiku.rag.chat.app import ChatApp

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            # First, handle a STATE_SNAPSHOT to establish initial state
            initial_snapshot = StateSnapshotEvent(
                type=EventType.STATE_SNAPSHOT,
                snapshot={
                    AGUI_STATE_KEY: {
                        "session_id": "test",
                        "citations": [],
                        "qa_history": [],
                        "citation_registry": {},
                    }
                },
            )
            tool_return1 = ToolReturnPart(
                tool_name="search",
                content="Initial search",
                tool_call_id="test-call-1",
                metadata=[initial_snapshot],
            )
            event1 = FunctionToolResultEvent(result=tool_return1)
            await app._handle_stream_event(event1)
            assert len(app.session_state.citations) == 0

            # Now handle a STATE_DELTA event that adds citations
            delta_event = StateDeltaEvent(
                type=EventType.STATE_DELTA,
                delta=[
                    {
                        "op": "replace",
                        "path": f"/{AGUI_STATE_KEY}/citations",
                        "value": [
                            {
                                "index": 1,
                                "document_id": "doc1",
                                "chunk_id": "chunk1",
                                "document_uri": "test.pdf",
                                "document_title": "Test Doc",
                                "content": "Test content from delta",
                            }
                        ],
                    },
                    {
                        "op": "add",
                        "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
                        "value": 1,
                    },
                ],
            )
            tool_return2 = ToolReturnPart(
                tool_name="ask",
                content="Answer with citations",
                tool_call_id="test-call-2",
                metadata=[delta_event],
            )
            event2 = FunctionToolResultEvent(result=tool_return2)

            # Handle the delta event
            await app._handle_stream_event(event2)

            # Session state should be synced
            assert len(app.session_state.citations) == 1
            assert app.session_state.citations[0].chunk_id == "chunk1"
            assert app.session_state.citations[0].content == "Test content from delta"


@pytest.mark.asyncio
async def test_handle_stream_event_delta_with_preinitialized_state(
    temp_db_path: Path,
):
    """Test delta handling when _agui_state_snapshot is pre-initialized.

    This is the actual TUI scenario: session_state exists from the start,
    so the agent emits deltas (not snapshots) even on the first tool call.
    The TUI pre-initializes _agui_state_snapshot from session_state.
    """
    from ag_ui.core import EventType, StateDeltaEvent
    from pydantic_ai import FunctionToolResultEvent
    from pydantic_ai.messages import ToolReturnPart

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
    from haiku.rag.chat.app import ChatApp

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            # Pre-initialize _agui_state_snapshot (simulating what _run_agent does)
            app._agui_state_snapshot = {
                AGUI_STATE_KEY: {
                    "session_id": "test",
                    "citations": [],
                    "qa_history": [],
                    "citation_registry": {},
                    "document_filter": [],
                    "initial_context": None,
                    "session_context": None,
                }
            }

            # Now handle a STATE_DELTA event directly (no prior snapshot event)
            delta_event = StateDeltaEvent(
                type=EventType.STATE_DELTA,
                delta=[
                    {
                        "op": "replace",
                        "path": f"/{AGUI_STATE_KEY}/citations",
                        "value": [
                            {
                                "index": 1,
                                "document_id": "doc1",
                                "chunk_id": "chunk1",
                                "document_uri": "test.pdf",
                                "document_title": "Test Doc",
                                "content": "Content from first delta",
                            }
                        ],
                    },
                    {
                        "op": "add",
                        "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
                        "value": 1,
                    },
                ],
            )
            tool_return = ToolReturnPart(
                tool_name="ask",
                content="Answer with citations",
                tool_call_id="test-call-1",
                metadata=[delta_event],
            )
            event = FunctionToolResultEvent(result=tool_return)

            # Handle the delta event
            await app._handle_stream_event(event)

            # Session state should be synced
            assert len(app.session_state.citations) == 1
            assert app.session_state.citations[0].chunk_id == "chunk1"
            assert app.session_state.citations[0].content == "Content from first delta"


@pytest.mark.asyncio
async def test_handle_stream_event_syncs_session_context(temp_db_path: Path):
    """Test that _handle_stream_event syncs session_context to session_state."""
    from ag_ui.core import EventType, StateDeltaEvent
    from pydantic_ai import FunctionToolResultEvent
    from pydantic_ai.messages import ToolReturnPart

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
    from haiku.rag.chat.app import ChatApp

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test():
            # Pre-initialize state
            app._agui_state_snapshot = {
                AGUI_STATE_KEY: {
                    "session_id": "test",
                    "citations": [],
                    "qa_history": [],
                    "citation_registry": {},
                    "document_filter": [],
                    "initial_context": None,
                    "session_context": None,
                }
            }

            # Verify session_context starts as None
            assert app.session_state.session_context is None

            # Handle a delta that adds session_context
            delta_event = StateDeltaEvent(
                type=EventType.STATE_DELTA,
                delta=[
                    {
                        "op": "replace",
                        "path": f"/{AGUI_STATE_KEY}/session_context",
                        "value": {
                            "summary": "User asked about Python async patterns.",
                            "last_updated": "2025-01-15T10:30:00",
                        },
                    },
                ],
            )
            tool_return = ToolReturnPart(
                tool_name="ask",
                content="Answer about async",
                tool_call_id="test-call-1",
                metadata=[delta_event],
            )
            event = FunctionToolResultEvent(result=tool_return)

            await app._handle_stream_event(event)

            # Session context should be synced to session_state
            assert app.session_state.session_context is not None
            assert (
                app.session_state.session_context.summary
                == "User asked about Python async patterns."
            )


@pytest.mark.asyncio
async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
    """Test that pressing Enter on a focused citation toggles expand/collapse."""
    from haiku.rag.agents.research.models import Citation
    from haiku.rag.chat.app import ChatApp
    from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
        app = ChatApp(temp_db_path, read_only=True)

        async with app.run_test() as pilot:
            chat_history = app.query_one(ChatHistory)

            # Add a citation
            test_citation = Citation(
                index=1,
                document_id="doc1",
                chunk_id="chunk1",
                document_uri="file:///test/doc1.pdf",
                document_title="Test Document",
                page_numbers=[1],
                content="Test content",
            )
            await chat_history.add_citations([test_citation])

            # Get the citation widget
            citation_widget = chat_history.query_one(CitationWidget)
            assert citation_widget.collapsed is True

            # Focus the citation
            citation_widget.focus()
            await pilot.pause()

            # Press Enter to expand
            await pilot.press("enter")
            await pilot.pause()
            assert citation_widget.collapsed is False

            # Press Enter again to collapse
            await pilot.press("enter")
            await pilot.pause()
            assert citation_widget.collapsed is True