import uuid

from ag_ui.core import StateDeltaEvent

from haiku.rag.agents.chat.state import (
    MAX_QA_HISTORY,
    ChatSessionState,
    QAResponse,
    SessionContext,
    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_default_state():
    """Test ChatDeps.state getter returns default state when not explicitly set."""
    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,
    )

    state = deps.state
    assert state is not None
    assert "session_id" in state
    assert state["qa_history"] == []
    assert state["citations"] == []


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_updates_default_state():
    """Test ChatDeps.state setter updates the default 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,
    )

    original_session_id = deps.session_state.session_id

    # Update with incoming state
    deps.state = {"session_id": "updated-123", "qa_history": [], "citations": []}

    assert deps.session_state.session_id == "updated-123"
    assert deps.session_state.session_id != original_session_id


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_index_assignment():
    """Test get_or_assign_index basic index assignment behavior.

    Verifies:
    - First chunk gets index 1
    - Second unique chunk gets index 2
    - Same chunk_id always returns same index
    """
    from haiku.rag.agents.chat.state import ChatSessionState

    session_state = ChatSessionState(session_id="test")

    # First chunk gets index 1
    index1 = session_state.get_or_assign_index("chunk-abc")
    assert index1 == 1

    # Second unique chunk gets index 2
    index2 = session_state.get_or_assign_index("chunk-def")
    assert index2 == 2

    # Same chunk_id returns same index (not incremented)
    index1_again = session_state.get_or_assign_index("chunk-abc")
    assert index1_again == 1


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

    session_state = ChatSessionState(session_id="test")

    # First round 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 regardless of order
    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_roundtrip():
    """Test citation_registry serializes and deserializes correctly for AG-UI state."""
    from haiku.rag.agents.chat.state import ChatSessionState

    # Create state and assign indices
    original = ChatSessionState(session_id="test")
    original.get_or_assign_index("chunk-a")
    original.get_or_assign_index("chunk-b")

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

    # Deserialize (simulating AG-UI state restoration)
    restored = ChatSessionState.model_validate(state_dict)

    # Existing chunks should return their persisted indices
    assert restored.get_or_assign_index("chunk-a") == 1
    assert restored.get_or_assign_index("chunk-b") == 2
    # New chunk should get next index
    assert restored.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"]


def test_chat_session_state_auto_generates_session_id():
    """New ChatSessionState should have a valid UUID session_id."""
    state = ChatSessionState()
    assert state.session_id
    assert len(state.session_id) == 36  # UUID format
    # Verify it's a valid UUID
    uuid.UUID(state.session_id)


def test_chat_session_state_preserves_explicit_session_id():
    """Explicit session_id should be preserved."""
    state = ChatSessionState(session_id="my-custom-id")
    assert state.session_id == "my-custom-id"


def test_chat_session_state_each_instance_gets_unique_id():
    """Each new instance should get a unique session_id."""
    state1 = ChatSessionState()
    state2 = ChatSessionState()
    assert state1.session_id != state2.session_id


def test_chat_session_state_initial_context_default_none():
    """Initial context should default to None."""
    state = ChatSessionState()
    assert state.initial_context is None


def test_chat_session_state_initial_context_preserved():
    """Explicit initial_context should be preserved."""
    state = ChatSessionState(initial_context="Background info about the project")
    assert state.initial_context == "Background info about the project"


def test_chat_session_state_initial_context_serialization():
    """initial_context should serialize and deserialize correctly."""
    state = ChatSessionState(
        session_id="test-123",
        initial_context="User is working on authentication",
    )
    state_dict = state.model_dump()
    assert state_dict["initial_context"] == "User is working on authentication"

    restored = ChatSessionState.model_validate(state_dict)
    assert restored.initial_context == "User is working on authentication"


def test_chat_session_state_model_dump_json_serializes_datetime():
    """model_dump(mode='json') should serialize datetime to ISO string.

    Agent tools use model_dump(mode='json') when creating StateSnapshotEvent
    to ensure datetime fields are JSON-serializable for external clients
    persisting AG-UI state to database JSON columns.
    """
    from datetime import datetime

    session_state = ChatSessionState(
        session_id="test",
        session_context=SessionContext(
            summary="Test summary",
            last_updated=datetime(2025, 1, 27, 12, 0, 0),
        ),
    )

    # This is how agent.py creates snapshots for StateSnapshotEvent
    snapshot = session_state.model_dump(mode="json")

    # datetime should be serialized as ISO string, not datetime object
    assert isinstance(snapshot["session_context"]["last_updated"], str)
    assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"


def test_emit_state_event_returns_none_when_no_changes():
    """emit_state_event returns None when states are identical."""
    from haiku.rag.agents.chat.state import emit_state_event

    state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])

    event = emit_state_event(state, state)

    assert event is None


def test_emit_state_event_returns_delta_with_changes():
    """emit_state_event returns StateDeltaEvent with JSON Patch ops for changes."""
    from ag_ui.core import EventType, StateDeltaEvent

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

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

    event = emit_state_event(current_state, new_state)

    assert isinstance(event, StateDeltaEvent)
    assert event.type == EventType.STATE_DELTA
    assert len(event.delta) > 0
    # Delta should contain an "add" operation for the new qa_history entry
    ops = event.delta
    qa_history_op = next((op for op in ops if "/qa_history" in op["path"]), None)
    assert qa_history_op is not None


def test_emit_state_event_delta_with_state_key():
    """emit_state_event wraps delta paths with state_key namespace."""
    from ag_ui.core import StateDeltaEvent

    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event

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

    event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY)

    assert isinstance(event, StateDeltaEvent)
    # Paths should be namespaced under state_key
    for op in event.delta:
        assert op["path"].startswith(f"/{AGUI_STATE_KEY}")


def test_emit_state_event_delta_produces_valid_patch():
    """emit_state_event delta can be applied to reproduce new state."""
    import jsonpatch

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

    current_state = ChatSessionState(
        session_id="test-123",
        qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
        citations=[],
    )
    new_state = ChatSessionState(
        session_id="test-123",
        qa_history=[
            QAResponse(question="Q1", answer="A1", confidence=0.9),
            QAResponse(question="Q2", answer="A2", confidence=0.8),
        ],
        citations=[],
    )

    event = emit_state_event(current_state, new_state)
    assert isinstance(event, StateDeltaEvent)

    # Apply patch to current state and verify it produces new state
    current_snapshot = current_state.model_dump(mode="json")
    patched = jsonpatch.apply_patch(current_snapshot, event.delta)
    new_snapshot = new_state.model_dump(mode="json")
    assert patched == new_snapshot