"""Tests for MongoDB client wrapper."""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from pymongo.errors import PyMongoError, ConnectionFailure, DuplicateKeyError

from src.geaicontext.clients.mongo_client import MongoClient, get_mongo_client


class TestMongoClient:
    """Test cases for MongoClient class."""

    @pytest.fixture
    def mongo_client(self):
        """Create a fresh MongoClient instance."""
        return MongoClient()

    @pytest.mark.asyncio
    async def test_init(self, mongo_client):
        """Test MongoClient initialization."""
        assert mongo_client._client is None
        assert mongo_client._db is None
        assert mongo_client._collection is None
        assert mongo_client._connected is False

    @pytest.mark.asyncio
    async def test_connect_success(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test successful MongoDB connection."""
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            
            assert mongo_client._connected is True
            mock_motor_client.admin.command.assert_called_once_with('ping')
            mock_mongo_collection.create_index.assert_called_once_with("customer_id", unique=True)

    @pytest.mark.asyncio
    async def test_connect_connection_failure(self, mongo_client):
        """Test MongoDB connection failure."""
        mock_client = MagicMock()
        mock_client.admin.command = AsyncMock(side_effect=ConnectionFailure("Connection refused"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_client):
            with pytest.raises(ConnectionFailure):
                await mongo_client.connect()
            
            assert mongo_client._connected is False

    @pytest.mark.asyncio
    async def test_connect_unexpected_error(self, mongo_client):
        """Test MongoDB connection with unexpected error."""
        mock_client = MagicMock()
        mock_client.admin.command = AsyncMock(side_effect=Exception("Unexpected error"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_client):
            with pytest.raises(Exception):
                await mongo_client.connect()

    @pytest.mark.asyncio
    async def test_disconnect(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test MongoDB disconnection."""
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            await mongo_client.disconnect()
            
            assert mongo_client._connected is False
            mock_motor_client.close.assert_called_once()

    @pytest.mark.asyncio
    async def test_disconnect_when_not_connected(self, mongo_client):
        """Test disconnect when not connected."""
        await mongo_client.disconnect()
        assert mongo_client._connected is False

    @pytest.mark.asyncio
    async def test_health_check_healthy(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test health check when connected."""
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.health_check()
            
            assert result is True

    @pytest.mark.asyncio
    async def test_health_check_not_connected(self, mongo_client):
        """Test health check when not connected."""
        result = await mongo_client.health_check()
        assert result is False

    @pytest.mark.asyncio
    async def test_health_check_ping_fails(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test health check when ping fails."""
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            # Make ping fail after connection
            mock_motor_client.admin.command = AsyncMock(side_effect=PyMongoError("Ping failed"))
            
            result = await mongo_client.health_check()
            assert result is False

    @pytest.mark.asyncio
    async def test_get_context_found(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_context, sample_customer_id):
        """Test retrieving existing context."""
        document = {**sample_context, "_id": "some_object_id", "customer_id": sample_customer_id}
        mock_mongo_collection.find_one = AsyncMock(return_value=document)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.get_context(sample_customer_id)
            
            assert result is not None
            assert "_id" not in result
            assert result["name"] == "John Doe"
            mock_mongo_collection.find_one.assert_called_once_with({"customer_id": sample_customer_id})

    @pytest.mark.asyncio
    async def test_get_context_not_found(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_customer_id):
        """Test retrieving non-existent context."""
        mock_mongo_collection.find_one = AsyncMock(return_value=None)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.get_context(sample_customer_id)
            
            assert result is None

    @pytest.mark.asyncio
    async def test_get_context_error(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_customer_id):
        """Test get_context with MongoDB error."""
        mock_mongo_collection.find_one = AsyncMock(side_effect=PyMongoError("Database error"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            
            with pytest.raises(PyMongoError):
                await mongo_client.get_context(sample_customer_id)

    @pytest.mark.asyncio
    async def test_save_context_insert(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_context, sample_customer_id):
        """Test saving new context (insert)."""
        mock_result = MagicMock()
        mock_result.upserted_id = "new_id"
        mock_result.matched_count = 0
        mock_mongo_collection.replace_one = AsyncMock(return_value=mock_result)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.save_context(sample_customer_id, sample_context)
            
            assert result is True
            mock_mongo_collection.replace_one.assert_called_once()

    @pytest.mark.asyncio
    async def test_save_context_update(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_context, sample_customer_id):
        """Test saving existing context (update)."""
        mock_result = MagicMock()
        mock_result.upserted_id = None
        mock_result.matched_count = 1
        mock_mongo_collection.replace_one = AsyncMock(return_value=mock_result)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.save_context(sample_customer_id, sample_context)
            
            assert result is True

    @pytest.mark.asyncio
    async def test_save_context_duplicate_key_error(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_context, sample_customer_id):
        """Test save_context with duplicate key error."""
        mock_mongo_collection.replace_one = AsyncMock(side_effect=DuplicateKeyError("Duplicate key"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.save_context(sample_customer_id, sample_context)
            
            assert result is False

    @pytest.mark.asyncio
    async def test_save_context_general_error(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_context, sample_customer_id):
        """Test save_context with general error."""
        mock_mongo_collection.replace_one = AsyncMock(side_effect=Exception("General error"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            
            with pytest.raises(Exception):
                await mongo_client.save_context(sample_customer_id, sample_context)

    @pytest.mark.asyncio
    async def test_delete_context_success(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_customer_id):
        """Test deleting existing context."""
        mock_result = MagicMock()
        mock_result.deleted_count = 1
        mock_mongo_collection.delete_one = AsyncMock(return_value=mock_result)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.delete_context(sample_customer_id)
            
            assert result is True

    @pytest.mark.asyncio
    async def test_delete_context_not_found(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_customer_id):
        """Test deleting non-existent context."""
        mock_result = MagicMock()
        mock_result.deleted_count = 0
        mock_mongo_collection.delete_one = AsyncMock(return_value=mock_result)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.delete_context(sample_customer_id)
            
            assert result is False

    @pytest.mark.asyncio
    async def test_delete_context_not_connected(self, mongo_client, sample_customer_id):
        """Test delete_context when not connected."""
        with pytest.raises(PyMongoError, match="MongoDB client not connected"):
            await mongo_client.delete_context(sample_customer_id)

    @pytest.mark.asyncio
    async def test_delete_context_error(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection, sample_customer_id):
        """Test delete_context with MongoDB error."""
        mock_mongo_collection.delete_one = AsyncMock(side_effect=PyMongoError("Delete error"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            
            with pytest.raises(PyMongoError):
                await mongo_client.delete_context(sample_customer_id)

    @pytest.mark.asyncio
    async def test_count_contexts_success(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test counting contexts."""
        mock_mongo_collection.count_documents = AsyncMock(return_value=42)
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            result = await mongo_client.count_contexts()
            
            assert result == 42

    @pytest.mark.asyncio
    async def test_count_contexts_not_connected(self, mongo_client):
        """Test count_contexts when not connected."""
        with pytest.raises(PyMongoError, match="MongoDB client not connected"):
            await mongo_client.count_contexts()

    @pytest.mark.asyncio
    async def test_count_contexts_error(self, mongo_client, mock_motor_client, mock_mongo_db, mock_mongo_collection):
        """Test count_contexts with MongoDB error."""
        mock_mongo_collection.count_documents = AsyncMock(side_effect=PyMongoError("Count error"))
        
        with patch("src.geaicontext.clients.mongo_client.AsyncIOMotorClient", return_value=mock_motor_client):
            await mongo_client.connect()
            
            with pytest.raises(PyMongoError):
                await mongo_client.count_contexts()


class TestGetMongoClient:
    """Test cases for get_mongo_client function."""

    @pytest.mark.asyncio
    async def test_get_mongo_client_creates_new(self, reset_global_clients):
        """Test that get_mongo_client creates a new client."""
        client = await get_mongo_client()
        assert client is not None
        assert isinstance(client, MongoClient)

    @pytest.mark.asyncio
    async def test_get_mongo_client_returns_same_instance(self, reset_global_clients):
        """Test that get_mongo_client returns the same instance."""
        client1 = await get_mongo_client()
        client2 = await get_mongo_client()
        assert client1 is client2
