"""Test streaming edit handling to prevent duplicate responses."""

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, MagicMock, patch

import nio
import pytest

from mindroom.bot import AgentBot
from mindroom.config.agent import AgentConfig
from mindroom.config.main import Config
from mindroom.config.models import ModelConfig, RouterConfig
from mindroom.handled_turns import TurnRecord
from mindroom.matrix.cache.thread_history_result import thread_history_result
from mindroom.matrix.users import AgentMatrixUser
from tests.conftest import (
    TEST_PASSWORD,
    bind_runtime_paths,
    drain_coalescing,
    install_runtime_cache_support,
    make_matrix_client_mock,
    replace_edit_regenerator_deps,
    runtime_paths_for,
    test_runtime_paths,
)
from tests.identity_helpers import persist_entity_accounts

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator
    from pathlib import Path


async def _empty_event_iterator() -> AsyncGenerator[object, None]:
    if False:
        yield None


def _make_matrix_client_mock() -> AsyncMock:
    client = make_matrix_client_mock(user_id="@mindroom_calculator:localhost")
    client.room_get_event_relations = MagicMock(side_effect=lambda *_args, **_kwargs: _empty_event_iterator())
    return client


def _delivery_resolution(response_event_id: str | None) -> str | None:
    return response_event_id


def setup_test_bot(
    agent: AgentMatrixUser,
    storage_path: Path,
    room_id: str,
    enable_streaming: bool = False,
    config: Config | None = None,
) -> AgentBot:
    """Set up a test bot with all required mocks."""
    runtime_paths = test_runtime_paths(storage_path)
    if config is None:
        config = bind_runtime_paths(Config(router=RouterConfig(model="default")), runtime_paths)
    else:
        try:
            runtime_paths_for(config)
        except KeyError:
            config = bind_runtime_paths(config, runtime_paths)
    persist_entity_accounts(config, runtime_paths_for(config))

    bot = AgentBot(
        agent,
        storage_path,
        config,
        runtime_paths_for(config),
        rooms=[room_id],
        enable_streaming=enable_streaming,
    )
    bot.client = _make_matrix_client_mock()
    install_runtime_cache_support(bot)
    bot._conversation_cache.get_thread_history = AsyncMock(return_value=thread_history_result([], is_full_history=True))
    bot._conversation_cache.get_dispatch_thread_history = AsyncMock(
        return_value=thread_history_result([], is_full_history=True),
    )
    bot._conversation_cache.get_dispatch_thread_snapshot = AsyncMock(
        return_value=thread_history_result([], is_full_history=False),
    )
    bot._conversation_resolver.fetch_thread_history = AsyncMock(
        return_value=thread_history_result([], is_full_history=True),
    )

    # Mock orchestrator
    mock_orchestrator = MagicMock()
    mock_orchestrator.current_config = config
    bot.orchestrator = mock_orchestrator

    return bot


@pytest.fixture
def mock_agent_user() -> AgentMatrixUser:
    """Create a mock agent user."""
    return AgentMatrixUser(
        agent_name="calculator",
        password=TEST_PASSWORD,
        display_name="CalculatorAgent",
        user_id="@mindroom_calculator:localhost",
    )


class TestStreamingEdits:
    """Test that streaming edits don't trigger duplicate responses."""

    def setup_method(self) -> None:
        """Set up test config."""
        self.config = Config(
            agents={
                "calculator": AgentConfig(display_name="CalculatorAgent", rooms=["!test:localhost"]),
                "helper": AgentConfig(display_name="HelperAgent", rooms=["!test:localhost"]),
            },
            teams={},
            room_models={},
            models={"default": ModelConfig(provider="ollama", id="test-model")},
            router=RouterConfig(model="default"),
        )

    @pytest.mark.asyncio
    @patch("mindroom.response_runner.ai_response")
    @patch("mindroom.response_runner.stream_agent_response")
    async def test_agent_regenerates_on_user_edits(
        self,
        mock_stream_agent_response: AsyncMock,  # noqa: ARG002
        mock_ai_response: AsyncMock,
        mock_agent_user: AgentMatrixUser,
        tmp_path: Path,
    ) -> None:
        """Test that agents regenerate their response when users edit messages."""
        # Set up bot
        bot = setup_test_bot(mock_agent_user, tmp_path, "!test:localhost", config=self.config)

        # Mock successful room_send response
        mock_send_response = MagicMock()
        mock_send_response.__class__ = nio.RoomSendResponse
        bot.client.room_send.return_value = mock_send_response

        # Mock AI response
        mock_ai_response.return_value = "I can help with that!"

        # Set up room
        mock_room = MagicMock()
        mock_room.room_id = "!test:localhost"

        # Initial message mentioning the agent
        initial_event = MagicMock()
        initial_event.sender = "@user:localhost"
        initial_event.body = "@mindroom_calculator:localhost: What's 2+2?"
        initial_event.event_id = "$initial123"
        initial_event.source = {
            "content": {
                "body": "@mindroom_calculator:localhost: What's 2+2?",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
            },
        }

        # Process initial message - bot should respond
        await bot._on_message(mock_room, initial_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 2  # thinking + final
        assert mock_ai_response.call_count == 1
        bot._turn_store.record_turn(
            TurnRecord.create(["$initial123"], response_event_id="$response123"),
        )

        # Reset mocks
        bot.client.room_send.reset_mock()
        mock_ai_response.reset_mock()

        # Edit event 1 - simulating streaming update
        edit_event1 = MagicMock()
        edit_event1.sender = "@user:localhost"
        edit_event1.body = "* @mindroom_calculator:localhost: What's 2+2? Can you show the work?"
        edit_event1.event_id = "$edit1"
        edit_event1.source = {
            "content": {
                "body": "* @mindroom_calculator:localhost: What's 2+2? Can you show the work?",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                "m.relates_to": {
                    "rel_type": "m.replace",
                    "event_id": "$initial123",  # References the original message
                },
                "m.new_content": {
                    "body": "@mindroom_calculator:localhost: What's 2+2? Can you show the work?",
                    "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                },
            },
        }

        # Process edit - bot SHOULD regenerate its response.
        mock_generate_response = AsyncMock(return_value=_delivery_resolution("$response123"))
        replace_edit_regenerator_deps(bot, generate_response=mock_generate_response)
        await bot._on_message(mock_room, edit_event1)
        await drain_coalescing(bot)
        assert mock_generate_response.await_count == 1

        # Edit event 2 - another streaming update
        edit_event2 = MagicMock()
        edit_event2.sender = "@user:localhost"
        edit_event2.body = "* @mindroom_calculator:localhost: What's 2+2? Can you show the work step by step?"
        edit_event2.event_id = "$edit2"
        edit_event2.source = {
            "content": {
                "body": "* @mindroom_calculator:localhost: What's 2+2? Can you show the work step by step?",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                "m.relates_to": {
                    "rel_type": "m.replace",
                    "event_id": "$initial123",  # Still references the original
                },
                "m.new_content": {
                    "body": "@mindroom_calculator:localhost: What's 2+2? Can you show the work step by step?",
                    "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                },
            },
        }

        # Reset mocks again
        bot.client.room_send.reset_mock()
        mock_ai_response.reset_mock()

        # Process second edit - bot should regenerate again.
        mock_generate_response = AsyncMock(return_value=_delivery_resolution("$response123"))
        replace_edit_regenerator_deps(bot, generate_response=mock_generate_response)
        await bot._on_message(mock_room, edit_event2)
        await drain_coalescing(bot)
        assert mock_generate_response.await_count == 1

    @pytest.mark.asyncio
    @patch("mindroom.response_runner.ai_response")
    async def test_agent_responds_to_new_messages_after_edits(
        self,
        mock_ai_response: AsyncMock,
        mock_agent_user: AgentMatrixUser,
        tmp_path: Path,
    ) -> None:
        """Test that agents still respond to new messages after seeing edits."""
        # Set up bot
        bot = setup_test_bot(mock_agent_user, tmp_path, "!test:localhost", config=self.config)

        # Mock successful room_send response
        mock_send_response = MagicMock()
        mock_send_response.__class__ = nio.RoomSendResponse
        bot.client.room_send.return_value = mock_send_response

        # Mock AI response
        mock_ai_response.return_value = "Here's the answer!"

        # Set up room
        mock_room = MagicMock()
        mock_room.room_id = "!test:localhost"

        # Mark that we already responded to some original message
        bot._turn_store.record_turn(TurnRecord.create(["$original123"]))

        # New message (NOT an edit) mentioning the agent
        new_event = MagicMock()
        new_event.sender = "@user:localhost"
        new_event.body = "@mindroom_calculator:localhost: What's 5+5?"
        new_event.event_id = "$new456"
        new_event.source = {
            "content": {
                "body": "@mindroom_calculator:localhost: What's 5+5?",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
            },
        }

        # Process new message - bot SHOULD respond
        await bot._on_message(mock_room, new_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 2  # thinking + final
        assert mock_ai_response.call_count == 1

    @pytest.mark.asyncio
    @patch("mindroom.response_runner.ai_response")
    async def test_agent_ignores_all_edits_from_agents(
        self,
        mock_ai_response: AsyncMock,
        mock_agent_user: AgentMatrixUser,
        tmp_path: Path,
    ) -> None:
        """Test that agents ignore ALL edits from other agents, even first-time mentions."""
        # Set up bot
        bot = setup_test_bot(mock_agent_user, tmp_path, "!test:localhost", config=self.config)

        # Mock successful room_send response
        mock_send_response = MagicMock()
        mock_send_response.__class__ = nio.RoomSendResponse
        bot.client.room_send.return_value = mock_send_response

        # Mock AI response
        mock_ai_response.return_value = "I can help with that!"

        # Set up room
        mock_room = MagicMock()
        mock_room.room_id = "!test:localhost"

        # Initial message from another agent WITHOUT mentioning calculator
        initial_event = MagicMock()
        initial_event.sender = "@mindroom_helper:localhost"  # Another agent
        initial_event.body = "Let me calculate something..."
        initial_event.event_id = "$initial123"
        initial_event.source = {
            "content": {
                "body": "Let me calculate something...",
            },
        }

        # Process initial message - calculator should NOT respond (not mentioned)
        await bot._on_message(mock_room, initial_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 0
        assert mock_ai_response.call_count == 0

        # Edit from agent that NOW mentions calculator (with in-progress marker)
        edit_event = MagicMock()
        edit_event.sender = "@mindroom_helper:localhost"  # Same agent
        edit_event.body = "* Let me calculate something... @mindroom_calculator:localhost can you help? ⋯"
        edit_event.event_id = "$edit1"
        edit_event.source = {
            "content": {
                "body": "* Let me calculate something... @mindroom_calculator:localhost can you help? ⋯",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                "m.relates_to": {
                    "rel_type": "m.replace",
                    "event_id": "$initial123",
                },
                "m.new_content": {
                    "body": "Let me calculate something... @mindroom_calculator:localhost can you help? ⋯",
                    "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                },
            },
        }

        # Process edit - calculator should STILL NOT respond (it's an edit from an agent)
        await bot._on_message(mock_room, edit_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 0
        assert mock_ai_response.call_count == 0

    @pytest.mark.asyncio
    @patch("mindroom.response_runner.ai_response")
    async def test_agent_responds_to_user_edits_with_new_mentions(
        self,
        mock_ai_response: AsyncMock,
        mock_agent_user: AgentMatrixUser,
        tmp_path: Path,
    ) -> None:
        """Test that agents DO NOT respond to user edits (even those that add mentions).

        This is by design - the bot ignores all edits to prevent confusion.
        Users should send a new message if they want a new response after editing.
        """
        # Set up bot
        bot = setup_test_bot(mock_agent_user, tmp_path, "!test:localhost", config=self.config)

        # Mock successful room_send response
        mock_send_response = MagicMock()
        mock_send_response.__class__ = nio.RoomSendResponse
        bot.client.room_send.return_value = mock_send_response

        # Mock AI response
        mock_ai_response.return_value = "I can help with that!"

        # Set up room
        mock_room = MagicMock()
        mock_room.room_id = "!test:localhost"

        # Initial message from user WITHOUT mentioning calculator
        initial_event = MagicMock()
        initial_event.sender = "@user:localhost"  # Regular user
        initial_event.body = "I need some help..."
        initial_event.event_id = "$initial123"
        initial_event.source = {
            "content": {
                "body": "I need some help...",
            },
        }

        # Process initial message - calculator should NOT respond (not mentioned)
        await bot._on_message(mock_room, initial_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 0
        assert mock_ai_response.call_count == 0

        # Edit from user that NOW mentions calculator
        edit_event = MagicMock()
        edit_event.sender = "@user:localhost"  # Same user
        edit_event.body = "* I need some help... @mindroom_calculator:localhost what's 2+2?"
        edit_event.event_id = "$edit1"
        edit_event.source = {
            "content": {
                "body": "* I need some help... @mindroom_calculator:localhost what's 2+2?",
                "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                "m.relates_to": {
                    "rel_type": "m.replace",
                    "event_id": "$initial123",
                },
                "m.new_content": {
                    "body": "I need some help... @mindroom_calculator:localhost what's 2+2?",
                    "m.mentions": {"user_ids": ["@mindroom_calculator:localhost"]},
                },
            },
        }

        # Process edit - calculator should NOT respond (bot ignores all edits)
        await bot._on_message(mock_room, edit_event)
        await drain_coalescing(bot)
        assert bot.client.room_send.call_count == 0
        assert mock_ai_response.call_count == 0
