from haiku.rag.agents.chat.state import (
    MAX_QA_HISTORY,
    QAResponse,
    build_document_filter,
    build_multi_document_filter,
    combine_filters,
)


def test_build_document_filter_simple():
    """Test build_document_filter with simple name."""
    result = build_document_filter("mytest")
    assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
    assert "LOWER(title) LIKE LOWER('%mytest%')" in result


def test_build_document_filter_with_spaces():
    """Test build_document_filter handles spaces correctly."""
    result = build_document_filter("TB MED 593")
    # Should include both the original (with spaces) and without spaces
    assert "LOWER(uri) LIKE LOWER('%TB MED 593%')" in result
    assert "LOWER(uri) LIKE LOWER('%TBMED593%')" in result
    assert "LOWER(title) LIKE LOWER('%TB MED 593%')" in result
    assert "LOWER(title) LIKE LOWER('%TBMED593%')" in result


def test_build_document_filter_escapes_quotes():
    """Test build_document_filter escapes single quotes."""
    result = build_document_filter("O'Reilly")
    # Single quotes should be doubled for SQL escaping
    assert "O''Reilly" in result


def test_build_multi_document_filter_empty():
    """Test build_multi_document_filter returns None for empty list."""
    result = build_multi_document_filter([])
    assert result is None


def test_build_multi_document_filter_single():
    """Test build_multi_document_filter with single document."""
    result = build_multi_document_filter(["mytest"])
    assert result is not None
    assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
    assert "LOWER(title) LIKE LOWER('%mytest%')" in result
    # Single document should not have extra wrapping parentheses
    assert " OR (" not in result


def test_build_multi_document_filter_multiple():
    """Test build_multi_document_filter with multiple documents."""
    result = build_multi_document_filter(["doc1", "doc2"])
    assert result is not None
    # Should have OR-combined filters
    assert "doc1" in result
    assert "doc2" in result
    assert " OR (" in result


def test_combine_filters_both_none():
    """Test combine_filters with both None."""
    result = combine_filters(None, None)
    assert result is None


def test_combine_filters_first_only():
    """Test combine_filters with only first filter."""
    result = combine_filters("uri = 'test'", None)
    assert result == "uri = 'test'"


def test_combine_filters_second_only():
    """Test combine_filters with only second filter."""
    result = combine_filters(None, "title = 'doc'")
    assert result == "title = 'doc'"


def test_combine_filters_both():
    """Test combine_filters combines with AND."""
    result = combine_filters("uri = 'test'", "title = 'doc'")
    assert result == "(uri = 'test') AND (title = 'doc')"


def test_max_qa_history_constant():
    """Test MAX_QA_HISTORY constant value."""
    assert MAX_QA_HISTORY == 50


def test_chat_deps_state_getter_returns_namespaced_state():
    """Test ChatDeps.state getter returns state under namespaced key."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(
        session_id="test-123",
        qa_history=[
            QAResponse(question="Q1", answer="A1", confidence=0.9),
        ],
    )

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    state = deps.state
    assert state is not None
    assert AGUI_STATE_KEY in state
    assert state[AGUI_STATE_KEY]["session_id"] == "test-123"
    assert len(state[AGUI_STATE_KEY]["qa_history"]) == 1
    assert state[AGUI_STATE_KEY]["qa_history"][0]["question"] == "Q1"


def test_chat_deps_state_getter_without_namespace():
    """Test ChatDeps.state getter returns flat state when no state_key."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="test-123")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=None,
    )

    state = deps.state
    assert state is not None
    assert "session_id" in state
    assert state["session_id"] == "test-123"


def test_chat_deps_state_getter_returns_none_without_session():
    """Test ChatDeps.state getter returns None when no session_state."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import ChatDeps

    mock_client = MagicMock()
    mock_config = MagicMock()

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=None,
    )

    assert deps.state is None


def test_chat_deps_state_setter_updates_from_namespaced_state():
    """Test ChatDeps.state setter updates session_state from namespaced incoming state."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="initial")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    # Simulate incoming AG-UI state with namespaced key
    incoming_state = {
        AGUI_STATE_KEY: {
            "session_id": "updated-123",
            "qa_history": [
                {"question": "Q1", "answer": "A1", "confidence": 0.9, "citations": []}
            ],
            "citations": [],
        }
    }

    deps.state = incoming_state

    assert deps.session_state is not None
    assert deps.session_state.session_id == "updated-123"
    assert len(deps.session_state.qa_history) == 1
    assert deps.session_state.qa_history[0].question == "Q1"


def test_chat_deps_state_setter_handles_none():
    """Test ChatDeps.state setter handles None gracefully."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="original")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
    )

    # Setting None should not raise and should not change state
    deps.state = None

    assert deps.session_state is not None
    assert deps.session_state.session_id == "original"


def test_chat_deps_state_setter_without_session_state():
    """Test ChatDeps.state setter does nothing when session_state is None."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import ChatDeps

    mock_client = MagicMock()
    mock_config = MagicMock()

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=None,
    )

    # Should not raise even with valid incoming state
    deps.state = {"session_id": "test", "qa_history": [], "citations": []}

    assert deps.session_state is None


def test_chat_deps_state_setter_with_citation_dicts():
    """Test ChatDeps.state setter converts citation dicts to Citation."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="test")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    incoming_state = {
        AGUI_STATE_KEY: {
            "session_id": "test",
            "qa_history": [],
            "citations": [
                {
                    "index": 1,
                    "document_id": "doc-1",
                    "chunk_id": "chunk-1",
                    "document_uri": "test.md",
                    "document_title": "Test Doc",
                    "page_numbers": [1, 2],
                    "headings": ["Intro"],
                    "content": "Test content",
                }
            ],
        }
    }

    deps.state = incoming_state

    assert deps.session_state is not None
    assert len(deps.session_state.citations) == 1
    citation = deps.session_state.citations[0]
    assert citation.document_id == "doc-1"
    assert citation.chunk_id == "chunk-1"
    assert citation.page_numbers == [1, 2]


def test_chat_deps_state_getter_includes_session_context():
    """Test ChatDeps.state getter includes session_context when present."""
    from datetime import datetime
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import (
        AGUI_STATE_KEY,
        ChatDeps,
        ChatSessionState,
        SessionContext,
    )

    mock_client = MagicMock()
    mock_config = MagicMock()

    now = datetime.now()
    session_state = ChatSessionState(
        session_id="test-123",
        session_context=SessionContext(
            summary="User discussed authentication.",
            last_updated=now,
        ),
    )

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    state = deps.state
    assert state is not None
    assert AGUI_STATE_KEY in state
    assert state[AGUI_STATE_KEY]["session_context"] is not None
    assert (
        state[AGUI_STATE_KEY]["session_context"]["summary"]
        == "User discussed authentication."
    )


def test_chat_deps_state_setter_ignores_session_context():
    """Test ChatDeps.state setter ignores session_context from client.

    The agent owns session_context via server-side cache, so client-provided
    session_context should be ignored to prevent stale state overwriting.
    """
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import (
        AGUI_STATE_KEY,
        ChatDeps,
        ChatSessionState,
        SessionContext,
    )

    mock_client = MagicMock()
    mock_config = MagicMock()

    # Start with a session_context (e.g., from cache)
    session_state = ChatSessionState(
        session_id="test",
        session_context=SessionContext(summary="Server-side context"),
    )
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    # Client sends different session_context (stale)
    incoming_state = {
        AGUI_STATE_KEY: {
            "session_id": "test",
            "qa_history": [],
            "citations": [],
            "session_context": {
                "summary": "Client-provided stale context",
                "last_updated": "2025-01-15T10:30:00",
            },
        }
    }

    deps.state = incoming_state

    # session_context should NOT be overwritten
    assert deps.session_state is not None
    assert deps.session_state.session_context is not None
    assert deps.session_state.session_context.summary == "Server-side context"


def test_citation_registry_get_or_assign_index_first_chunk():
    """Test get_or_assign_index assigns index 1 to first chunk."""
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")
    index = session_state.get_or_assign_index("chunk-abc")
    assert index == 1


def test_citation_registry_get_or_assign_index_second_chunk():
    """Test get_or_assign_index assigns incremental indices to new chunks."""
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")
    index1 = session_state.get_or_assign_index("chunk-abc")
    index2 = session_state.get_or_assign_index("chunk-def")
    assert index1 == 1
    assert index2 == 2


def test_citation_registry_get_or_assign_index_same_chunk():
    """Test get_or_assign_index returns same index for same chunk_id."""
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")
    index1 = session_state.get_or_assign_index("chunk-abc")
    index2 = session_state.get_or_assign_index("chunk-abc")
    assert index1 == index2 == 1


def test_citation_registry_get_or_assign_index_stability():
    """Test citation indices are stable across multiple calls."""
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")

    # First call assigns indices 1, 2, 3
    idx_a = session_state.get_or_assign_index("chunk-a")
    idx_b = session_state.get_or_assign_index("chunk-b")
    idx_c = session_state.get_or_assign_index("chunk-c")

    # Second round - existing chunks keep their indices
    assert session_state.get_or_assign_index("chunk-b") == idx_b
    assert session_state.get_or_assign_index("chunk-a") == idx_a
    assert session_state.get_or_assign_index("chunk-c") == idx_c

    # New chunk gets next index
    idx_d = session_state.get_or_assign_index("chunk-d")
    assert idx_d == 4


def test_citation_registry_serialization():
    """Test citation_registry is included in model_dump for AG-UI state."""
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")
    session_state.get_or_assign_index("chunk-a")
    session_state.get_or_assign_index("chunk-b")

    state_dict = session_state.model_dump()
    assert "citation_registry" in state_dict
    assert state_dict["citation_registry"] == {"chunk-a": 1, "chunk-b": 2}


def test_citation_registry_deserialization():
    """Test citation_registry is restored from dict."""
    from haiku.rag.agents.chat.state import ChatSessionState

    # Simulate state from AG-UI using model_validate (proper Pydantic deserialization)
    session_state = ChatSessionState.model_validate(
        {
            "session_id": "test",
            "citations": [],
            "qa_history": [],
            "citation_registry": {"chunk-a": 1, "chunk-b": 2},
        }
    )

    # Existing chunks should return their persisted indices
    assert session_state.get_or_assign_index("chunk-a") == 1
    assert session_state.get_or_assign_index("chunk-b") == 2
    # New chunk should get next index
    assert session_state.get_or_assign_index("chunk-c") == 3


def test_chat_deps_state_getter_includes_citation_registry():
    """Test ChatDeps.state getter includes citation_registry."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="test")
    session_state.get_or_assign_index("chunk-a")

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    state = deps.state
    assert state is not None
    assert AGUI_STATE_KEY in state
    assert state[AGUI_STATE_KEY]["citation_registry"] == {"chunk-a": 1}


def test_chat_deps_state_setter_restores_citation_registry():
    """Test ChatDeps.state setter restores citation_registry from incoming state."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="test")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    # Simulate incoming AG-UI state with citation_registry
    incoming_state = {
        AGUI_STATE_KEY: {
            "session_id": "test",
            "qa_history": [],
            "citations": [],
            "citation_registry": {"chunk-x": 1, "chunk-y": 2},
        }
    }

    deps.state = incoming_state

    assert deps.session_state is not None
    # Registry should be restored
    assert deps.session_state.get_or_assign_index("chunk-x") == 1
    assert deps.session_state.get_or_assign_index("chunk-y") == 2
    # New chunk gets next index
    assert deps.session_state.get_or_assign_index("chunk-z") == 3


def test_chat_deps_state_setter_restores_document_filter():
    """Test ChatDeps.state setter restores document_filter from incoming state."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(session_id="test")
    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    incoming_state = {
        AGUI_STATE_KEY: {
            "session_id": "test",
            "qa_history": [],
            "citations": [],
            "document_filter": ["doc1.pdf", "doc2.pdf"],
        }
    }

    deps.state = incoming_state

    assert deps.session_state is not None
    assert deps.session_state.document_filter == ["doc1.pdf", "doc2.pdf"]


def test_chat_deps_state_getter_includes_document_filter():
    """Test ChatDeps.state getter includes document_filter."""
    from unittest.mock import MagicMock

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState

    mock_client = MagicMock()
    mock_config = MagicMock()

    session_state = ChatSessionState(
        session_id="test-123",
        document_filter=["doc1.pdf", "doc2.pdf"],
    )

    deps = ChatDeps(
        client=mock_client,
        config=mock_config,
        session_state=session_state,
        state_key=AGUI_STATE_KEY,
    )

    state = deps.state
    assert state is not None
    assert AGUI_STATE_KEY in state
    assert state[AGUI_STATE_KEY]["document_filter"] == ["doc1.pdf", "doc2.pdf"]