# -*- coding: utf-8 -*-
"""Location: ./tests/integration/test_llmchat_endpoints.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Keval Mahajan

Integration Tests for LLM Chat API Router

This test suite covers:
1. Multi-Worker Coordination Tests
2. Redis Failure Scenarios
3. Chat History Persistence Tests

All tests use mocked dependencies (Redis, LLM providers, MCP clients) to ensure
isolation from external services.
"""

import json
import pytest
from typing import Any, Dict
from unittest.mock import AsyncMock, Mock, patch
from uuid import uuid4

# ==================== FIXTURES ====================


@pytest.fixture
def mock_redis():
    """Mock Redis client"""
    redis_mock = AsyncMock()
    redis_mock.get = AsyncMock(return_value=None)
    redis_mock.set = AsyncMock(return_value=True)
    redis_mock.delete = AsyncMock(return_value=True)
    redis_mock.expire = AsyncMock(return_value=True)
    redis_mock.rpush = AsyncMock(return_value=1)
    redis_mock.lrange = AsyncMock(return_value=[])
    redis_mock.ltrim = AsyncMock(return_value=True)
    redis_mock.llen = AsyncMock(return_value=0)
    redis_mock.close = AsyncMock()
    return redis_mock


@pytest.fixture
def mock_mcp_client():
    """Mock MCP Client"""
    client = AsyncMock()
    client.connect = AsyncMock()
    client.disconnect = AsyncMock()
    client.get_tools = AsyncMock(return_value=[])
    client.call_tool = AsyncMock(return_value={"result": "test"})
    client.is_connected = True
    return client


@pytest.fixture
def mock_agent():
    """Mock LangGraph Agent"""
    agent = AsyncMock()

    async def mock_astream_events(messages, version):
        yield {"event": "on_chat_model_stream", "data": {"chunk": Mock(content="Test ")}}
        yield {"event": "on_chat_model_stream", "data": {"chunk": Mock(content="response")}}

    async def mock_ainvoke(messages):
        return {"messages": [Mock(content="Test response")]}

    agent.astream_events = mock_astream_events
    agent.ainvoke = mock_ainvoke
    return agent


# ==================== HELPER FUNCTIONS ====================


def create_connect_payload(user_id: str, model: str = "test-model") -> Dict[str, Any]:
    """Helper to create connect request payload.

    Note: Models must be configured via Admin UI -> LLM Settings.
    The model parameter should match a model_id from the gateway.
    """
    return {
        "user_id": user_id,
        "server": {"url": "http://test-mcp-server.com/mcp", "transport": "streamable_http", "auth_token": "test-token"},
        "llm": {
            "model": model,
            "temperature": 0.7,
        },
        "streaming": False,
    }


def create_chat_payload(user_id: str, message: str, streaming: bool = False) -> Dict[str, Any]:
    """Helper to create chat request payload"""
    return {"user_id": user_id, "message": message, "streaming": streaming}


# ==================== MULTI-WORKER COORDINATION TESTS ====================


class TestMultiWorkerCoordination:
    """Tests for multi-worker coordination scenarios"""

    @pytest.mark.asyncio
    async def test_session_handoff_between_workers(self, mock_redis):
        """Test session handoff when different workers handle requests"""
        from mcpgateway.routers.llmchat_router import set_active_session, _active_key

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "user123"
            mock_service = AsyncMock()

            await set_active_session(user_id, mock_service)

            mock_redis.set.assert_called()
            call_args = mock_redis.set.call_args[0]
            assert _active_key(user_id) in str(call_args)

            with patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-2"):
                mock_redis.get = AsyncMock(return_value="worker-1")
                owner_result = await mock_redis.get(_active_key(user_id))
                assert owner_result == "worker-1"

    @pytest.mark.asyncio
    async def test_distributed_lock_acquisition_and_release(self, mock_redis):
        """Test distributed lock mechanism for session initialization"""
        from mcpgateway.routers.llmchat_router import _try_acquire_lock, _release_lock_safe

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "user123"

            mock_redis.set = AsyncMock(return_value=True)
            acquired = await _try_acquire_lock(user_id)
            assert acquired is True

            mock_redis.set.assert_called()

            mock_redis.get = AsyncMock(return_value="worker-1")
            await _release_lock_safe(user_id)
            mock_redis.delete.assert_called()

    @pytest.mark.asyncio
    async def test_session_ttl_expiration_and_renewal(self, mock_redis):
        """Test session TTL expiration and automatic renewal"""
        from mcpgateway.routers.llmchat_router import get_active_session

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            # Import locally to avoid affecting other tests
            from mcpgateway.routers import llmchat_router as router_module

            user_id = "user123"
            mock_service = AsyncMock()

            # Temporarily add to sessions
            original_sessions = router_module.active_sessions.copy()
            router_module.active_sessions[user_id] = mock_service

            try:
                mock_redis.get = AsyncMock(return_value="worker-1")
                mock_redis.expire = AsyncMock(return_value=True)

                session = await get_active_session(user_id)

                assert session is not None
                assert session == mock_service
                mock_redis.expire.assert_called()
            finally:
                # Restore original state
                router_module.active_sessions.clear()
                router_module.active_sessions.update(original_sessions)


# ==================== REDIS FAILURE SCENARIOS ====================


class TestRedisFailureScenarios:
    """Tests for Redis failure and fallback scenarios"""

    @pytest.mark.asyncio
    async def test_graceful_degradation_redis_unavailable(self):
        """Test graceful degradation when Redis is unavailable"""
        from mcpgateway.routers.llmchat_router import set_user_config, get_user_config, set_active_session, get_active_session
        from mcpgateway.services.mcp_client_chat_service import MCPClientConfig, MCPServerConfig, LLMConfig, OllamaConfig

        with patch("mcpgateway.routers.llmchat_router.redis_client", None):

            user_id = "test_user_" + str(uuid4())
            config = MCPClientConfig(mcp_server=MCPServerConfig(url="http://test.com/mcp"), llm=LLMConfig(provider="ollama", config=OllamaConfig(model="llama2")))

            await set_user_config(user_id, config)
            retrieved_config = await get_user_config(user_id)
            assert retrieved_config == config

            mock_service = AsyncMock()
            await set_active_session(user_id, mock_service)
            session = await get_active_session(user_id)
            assert session == mock_service

            # Cleanup
            from mcpgateway.routers import llmchat_router as router_module

            if user_id in router_module.user_configs:
                del router_module.user_configs[user_id]
            if user_id in router_module.active_sessions:
                del router_module.active_sessions[user_id]

    @pytest.mark.asyncio
    async def test_redis_connection_loss_during_active_session(self, mock_redis):
        """Test handling of Redis connection loss during an active session"""
        from mcpgateway.routers.llmchat_router import get_active_session
        from mcpgateway.routers import llmchat_router as router_module

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "test_user_" + str(uuid4())
            mock_service = AsyncMock()

            original_sessions = router_module.active_sessions.copy()
            router_module.active_sessions[user_id] = mock_service

            try:
                mock_redis.get = AsyncMock(return_value="worker-1")
                mock_redis.expire = AsyncMock(side_effect=ConnectionError("Redis connection lost"))

                session = await get_active_session(user_id)
                assert session == mock_service
            finally:
                router_module.active_sessions.clear()
                router_module.active_sessions.update(original_sessions)


# ==================== CHAT HISTORY PERSISTENCE TESTS ====================


class TestChatHistoryPersistence:
    """Tests for chat history persistence and management"""

    @pytest.mark.asyncio
    async def test_history_with_redis(self, mock_redis):
        """Test chat history with Redis backend"""
        from mcpgateway.services.mcp_client_chat_service import ChatHistoryManager

        user_id = "test_user_" + str(uuid4())

        messages = [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]

        # Mock Redis to return JSON data
        mock_redis.get = AsyncMock(return_value=json.dumps(messages))

        history_manager = ChatHistoryManager(redis_client=mock_redis, max_messages=100)

        history = await history_manager.get_history(user_id)

        mock_redis.get.assert_called()
        assert len(history) == 2
        assert history[0]["content"] == "Hello"
        assert history[1]["content"] == "Hi there!"

    @pytest.mark.asyncio
    async def test_message_ordering_in_memory(self):
        """Test message ordering with in-memory storage"""
        from mcpgateway.services.mcp_client_chat_service import ChatHistoryManager

        user_id = "test_user_" + str(uuid4())
        history_manager = ChatHistoryManager(redis_client=None)

        for i in range(5):
            await history_manager.append_message(user_id, "user", f"Message {i}")

        history = await history_manager.get_history(user_id)
        assert len(history) == 5
        for i, msg in enumerate(history):
            assert msg["content"] == f"Message {i}"

        # Cleanup
        await history_manager.clear_history(user_id)

    @pytest.mark.asyncio
    async def test_history_size_limits(self):
        """Test chat history size limits and automatic trimming"""
        from mcpgateway.services.mcp_client_chat_service import ChatHistoryManager

        user_id = "test_user_" + str(uuid4())
        max_messages = 10
        history_manager = ChatHistoryManager(redis_client=None, max_messages=max_messages)

        for i in range(15):
            await history_manager.append_message(user_id, "user", f"Message {i}")

        history = await history_manager.get_history(user_id)
        assert len(history) <= max_messages

        # Cleanup
        await history_manager.clear_history(user_id)

    @pytest.mark.asyncio
    async def test_clear_history(self, mock_redis):
        """Test clear history operation"""
        from mcpgateway.services.mcp_client_chat_service import ChatHistoryManager

        user_id = "test_user_" + str(uuid4())

        history_manager = ChatHistoryManager(redis_client=mock_redis)

        mock_redis.delete = AsyncMock(return_value=True)
        await history_manager.clear_history(user_id)

        mock_redis.delete.assert_called()


# ==================== CONFIGURATION TESTS ====================


class TestConfigurationManagement:
    """Tests for configuration management"""

    @pytest.mark.asyncio
    async def test_config_storage_and_retrieval(self):
        """Test storing and retrieving user configuration"""
        from mcpgateway.routers.llmchat_router import set_user_config, get_user_config
        from mcpgateway.services.mcp_client_chat_service import MCPClientConfig, MCPServerConfig, LLMConfig, OllamaConfig

        with patch("mcpgateway.routers.llmchat_router.redis_client", None):
            user_id = "test_user_" + str(uuid4())
            config = MCPClientConfig(mcp_server=MCPServerConfig(url="http://test.com/mcp"), llm=LLMConfig(provider="ollama", config=OllamaConfig(model="llama2")))

            await set_user_config(user_id, config)
            retrieved = await get_user_config(user_id)

            assert retrieved == config

            # Cleanup
            from mcpgateway.routers import llmchat_router as router_module

            if user_id in router_module.user_configs:
                del router_module.user_configs[user_id]

    @pytest.mark.asyncio
    async def test_config_with_redis(self, mock_redis):
        """Test configuration storage with Redis"""
        from mcpgateway.routers.llmchat_router import set_user_config
        from mcpgateway.services.mcp_client_chat_service import MCPClientConfig, MCPServerConfig, LLMConfig, OllamaConfig

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis):
            user_id = "test_user_" + str(uuid4())
            config = MCPClientConfig(mcp_server=MCPServerConfig(url="http://test.com/mcp"), llm=LLMConfig(provider="ollama", config=OllamaConfig(model="llama2")))

            mock_redis.set = AsyncMock(return_value=True)
            await set_user_config(user_id, config)

            mock_redis.set.assert_called()


# ==================== LOCK MECHANISM TESTS ====================


class TestLockMechanism:
    """Tests for distributed lock mechanism"""

    @pytest.mark.asyncio
    async def test_lock_acquisition(self, mock_redis):
        """Test successful lock acquisition"""
        from mcpgateway.routers.llmchat_router import _try_acquire_lock

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "test_user_" + str(uuid4())

            mock_redis.set = AsyncMock(return_value=True)
            acquired = await _try_acquire_lock(user_id)

            assert acquired is True
            mock_redis.set.assert_called()

    @pytest.mark.asyncio
    async def test_lock_failure(self, mock_redis):
        """Test failed lock acquisition"""
        from mcpgateway.routers.llmchat_router import _try_acquire_lock

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "test_user_" + str(uuid4())

            mock_redis.set = AsyncMock(return_value=False)
            acquired = await _try_acquire_lock(user_id)

            assert acquired is False

    @pytest.mark.asyncio
    async def test_lock_release(self, mock_redis):
        """Test lock release"""
        from mcpgateway.routers.llmchat_router import _release_lock_safe

        with patch("mcpgateway.routers.llmchat_router.redis_client", mock_redis), patch("mcpgateway.routers.llmchat_router.WORKER_ID", "worker-1"):

            user_id = "test_user_" + str(uuid4())

            mock_redis.get = AsyncMock(return_value="worker-1")
            mock_redis.delete = AsyncMock(return_value=True)

            await _release_lock_safe(user_id)

            mock_redis.delete.assert_called()
