"""
Tests for automatic training trigger functionality.

This module tests the automatic triggering of index training based on
vector count thresholds after upsert operations.
"""

import pytest
from unittest.mock import Mock, patch, MagicMock
import sys

# Mock cyborgdb_core before importing
sys.modules["cyborgdb_core"] = MagicMock()

# Now import after mocking
from cyborgdb_service.core.training_manager import TrainingManager, TrainingJob
from cyborgdb_service.api.schemas.vectors import UpsertRequest
from cyborgdb_service.api.routes.vectors import upsert_vectors
from cyborgdb_service.api.routes.indexes import get_training_status


class TestTrainingManager:
    """Test the TrainingManager class functionality."""

    def test_should_retrain_untrained_index(self):
        """Test retraining logic for untrained indexes (n_lists=1)."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # For untrained index, effective n_lists = 1
            # Threshold = 1 * 10000 = 10000
            assert not manager.should_retrain("test_index", 9999, 128, False, "ivfflat")
            assert manager.should_retrain("test_index", 10001, 128, False, "ivfflat")
            assert manager.should_retrain("test_index", 20000, 128, False, "ivfflat")

    def test_should_retrain_trained_index(self):
        """Test retraining logic for trained indexes."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # For trained ivfflat index with n_lists=128
            # Threshold = 128 * 10000 = 1,280,000
            assert not manager.should_retrain(
                "test_index", 1_000_000, 128, True, "ivfflat"
            )
            assert manager.should_retrain("test_index", 1_280_001, 128, True, "ivfflat")
            assert manager.should_retrain("test_index", 2_000_000, 128, True, "ivfflat")

    def test_should_retrain_various_n_lists(self):
        """Test retraining logic with various n_lists values."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # n_lists = 64, trained ivfflat
            # Threshold = 64 * 10000 = 640,000
            assert not manager.should_retrain(
                "test_index", 639_999, 64, True, "ivfflat"
            )
            assert manager.should_retrain("test_index", 640_001, 64, True, "ivfflat")

            # n_lists = 256, trained ivfflat
            # Threshold = 256 * 10000 = 2,560,000
            assert not manager.should_retrain(
                "test_index", 2_559_999, 256, True, "ivfflat"
            )
            assert manager.should_retrain("test_index", 2_560_001, 256, True, "ivfflat")

    def test_should_retrain_ivf_index_type_constraints(self):
        """Test index type constraints for IVF and IVFPQ indexes."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # IVF index - can train when untrained, cannot retrain when trained
            assert manager.should_retrain(
                "ivf_index", 15000, 128, False, "ivf"
            )  # Initial training allowed
            assert not manager.should_retrain(
                "ivf_index", 2_000_000, 128, True, "ivf"
            )  # Retraining blocked

            # IVFPQ index - can train when untrained, cannot retrain when trained
            assert manager.should_retrain(
                "ivfpq_index", 15000, 128, False, "ivfpq"
            )  # Initial training allowed
            assert not manager.should_retrain(
                "ivfpq_index", 2_000_000, 128, True, "ivfpq"
            )  # Retraining blocked

            # IVFFlat index - can train and retrain
            assert manager.should_retrain(
                "ivfflat_index", 15000, 128, False, "ivfflat"
            )  # Initial training allowed
            assert manager.should_retrain(
                "ivfflat_index", 2_000_000, 128, True, "ivfflat"
            )  # Retraining allowed

    def test_should_retrain_case_insensitive_index_types(self):
        """Test that index type checking is case insensitive."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # Test uppercase index types
            assert not manager.should_retrain("test_index", 2_000_000, 128, True, "IVF")
            assert not manager.should_retrain(
                "test_index", 2_000_000, 128, True, "IVFPQ"
            )
            assert manager.should_retrain("test_index", 2_000_000, 128, True, "IVFFLAT")

            # Test mixed case index types
            assert not manager.should_retrain("test_index", 2_000_000, 128, True, "IvF")
            assert not manager.should_retrain(
                "test_index", 2_000_000, 128, True, "IvfPq"
            )
            assert manager.should_retrain("test_index", 2_000_000, 128, True, "IvfFlat")

    def test_is_training_tracking(self):
        """Test tracking of indexes being trained."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # Initially no indexes are being trained
            assert not manager.is_training("index1")
            assert not manager.is_training("index2")

            # Set index as currently training
            with manager._lock:
                manager._currently_training = "index1"

            assert manager.is_training("index1")
            assert not manager.is_training("index2")

            # Clear currently training
            with manager._lock:
                manager._currently_training = None

            assert not manager.is_training("index1")

    def test_trigger_training_if_needed_below_threshold(self):
        """Test that training is not triggered when below threshold."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()
            manager._training_client = Mock()

            # Below threshold - should not trigger
            triggered = manager.trigger_training_if_needed(
                index_name="test_index",
                index_key_hex="0" * 64,
                num_vectors=5000,
                n_lists=128,
                is_trained=False,
                index_type="ivfflat",
            )

            assert not triggered
            assert not manager.is_training("test_index")

    def test_trigger_training_if_needed_above_threshold(self):
        """Test that training is triggered when above threshold."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            with patch.object(
                TrainingManager, "_ensure_worker_running"
            ):  # Don't start worker
                manager = TrainingManager()
                manager._training_client = Mock()

                # Above threshold - should trigger
                triggered = manager.trigger_training_if_needed(
                    index_name="test_index",
                    index_key_hex="0" * 64,
                    num_vectors=15000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                assert triggered
                assert manager.is_training("test_index")

    def test_trigger_training_already_training(self):
        """Test that training is not triggered if already in progress."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # Mark index as currently training
            with manager._lock:
                manager._currently_training = "test_index"

            # Should not trigger again
            triggered = manager.trigger_training_if_needed(
                index_name="test_index",
                index_key_hex="0" * 64,
                num_vectors=15000,
                n_lists=128,
                is_trained=False,
                index_type="ivfflat",
            )

            assert not triggered
            assert manager.is_training("test_index")

    def test_get_training_status(self):
        """Test getting the training status."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # Set currently training and queue an index
            with manager._lock:
                manager._currently_training = "index1"

            job = TrainingJob(index_name="index2", index_key_hex="0" * 64)
            manager._queue.put(job)

            status = manager.get_training_status()

            assert status["currently_training"] == "index1"
            assert status["queued_indexes"] == ["index2"]
            assert status["training_indexes"] == ["index1", "index2"]
            assert status["retrain_threshold"] == 10000

    def test_custom_retrain_threshold(self):
        """Test custom retrain threshold from environment variable."""
        with patch.dict("os.environ", {"RETRAIN_THRESHOLD": "5000"}):
            with patch.object(TrainingManager, "_initialize_training_client"):
                manager = TrainingManager()
                assert manager._retrain_threshold == 5000

                # Test with custom threshold
                # Threshold = 1 * 5000 = 5000 for untrained
                assert not manager.should_retrain(
                    "test_index", 4999, 128, False, "ivfflat"
                )
                assert manager.should_retrain("test_index", 5001, 128, False, "ivfflat")

    def test_trigger_training_blocked_for_trained_ivf_index(self):
        """Test that training is blocked for trained IVF/IVFPQ indexes."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()
            manager._training_client = Mock()

            # Trained IVF index - should not trigger even if above threshold
            triggered = manager.trigger_training_if_needed(
                index_name="test_index",
                index_key_hex="0" * 64,
                num_vectors=2_000_000,  # Well above threshold
                n_lists=128,
                is_trained=True,
                index_type="ivf",
            )

            assert not triggered
            assert not manager.is_training("test_index")

            # Trained IVFPQ index - should not trigger even if above threshold
            triggered = manager.trigger_training_if_needed(
                index_name="test_index2",
                index_key_hex="0" * 64,
                num_vectors=2_000_000,  # Well above threshold
                n_lists=128,
                is_trained=True,
                index_type="ivfpq",
            )

            assert not triggered
            assert not manager.is_training("test_index2")


class TestUpsertWithAutoTrigger:
    """Test the upsert endpoint with auto-trigger functionality."""

    @pytest.mark.asyncio
    async def test_upsert_triggers_training(self):
        """Test that upsert triggers training when threshold is exceeded."""
        # Create request
        request = UpsertRequest(
            index_name="test_index",
            index_key="0" * 64,
            items=[{"id": f"item_{i}", "vector": [0.1] * 128} for i in range(100)],
        )

        # Mock dependencies
        mock_client = Mock()
        mock_index = Mock()
        mock_index.upsert.return_value = None
        mock_index.get_num_vectors.return_value = 15000  # Above threshold
        mock_index.is_trained.return_value = False
        mock_index.index_config.return_value = {"n_lists": 128, "type": "ivfflat"}

        with patch(
            "cyborgdb_service.api.routes.vectors.load_index", return_value=mock_index
        ):
            with patch(
                "cyborgdb_service.api.routes.vectors.get_training_manager"
            ) as mock_get_manager:
                mock_manager = Mock()
                mock_manager._retrain_threshold = 10000
                # Now sync, not AsyncMock
                mock_manager.trigger_training_if_needed = Mock(return_value=True)
                mock_get_manager.return_value = mock_manager

                # Call upsert
                response = await upsert_vectors(request, mock_client)

                # Verify upsert was called
                mock_index.upsert.assert_called_once()

                # Verify training was triggered
                mock_manager.trigger_training_if_needed.assert_called_once_with(
                    index_name="test_index",
                    index_key_hex="0" * 64,
                    num_vectors=15000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                # Check response
                assert response["status"] == "success"
                assert "Upserted 100 vectors" in response["message"]
                assert response.get("training_triggered")
                assert "training_message" in response

    @pytest.mark.asyncio
    async def test_upsert_no_training_below_threshold(self):
        """Test that upsert doesn't trigger training when below threshold."""
        # Create request
        request = UpsertRequest(
            index_name="test_index",
            index_key="0" * 64,
            items=[{"id": f"item_{i}", "vector": [0.1] * 128} for i in range(10)],
        )

        # Mock dependencies
        mock_client = Mock()
        mock_index = Mock()
        mock_index.upsert.return_value = None
        mock_index.get_num_vectors.return_value = 5000  # Below threshold
        mock_index.is_trained.return_value = False
        mock_index.index_config.return_value = {"n_lists": 128, "type": "ivfflat"}

        with patch(
            "cyborgdb_service.api.routes.vectors.load_index", return_value=mock_index
        ):
            with patch(
                "cyborgdb_service.api.routes.vectors.get_training_manager"
            ) as mock_get_manager:
                mock_manager = Mock()
                mock_manager._retrain_threshold = 10000
                mock_manager.trigger_training_if_needed = Mock(return_value=False)
                mock_get_manager.return_value = mock_manager

                # Call upsert
                response = await upsert_vectors(request, mock_client)

                # Verify upsert was called
                mock_index.upsert.assert_called_once()

                # Verify training check was called
                mock_manager.trigger_training_if_needed.assert_called_once()

                # Check response - no training triggered
                assert response["status"] == "success"
                assert "Upserted 10 vectors" in response["message"]
                assert "training_triggered" not in response

    @pytest.mark.asyncio
    async def test_upsert_training_check_failure(self):
        """Test that upsert succeeds even if training check fails."""
        # Create request
        request = UpsertRequest(
            index_name="test_index",
            index_key="0" * 64,
            items=[{"id": "item_1", "vector": [0.1] * 128}],
        )

        # Mock dependencies
        mock_client = Mock()
        mock_index = Mock()
        mock_index.upsert.return_value = None
        mock_index.get_num_vectors.side_effect = Exception("Failed to get num vectors")

        with patch(
            "cyborgdb_service.api.routes.vectors.load_index", return_value=mock_index
        ):
            with patch("cyborgdb_service.api.routes.vectors.logger") as mock_logger:
                # Call upsert
                response = await upsert_vectors(request, mock_client)

                # Verify upsert was called
                mock_index.upsert.assert_called_once()

                # Check that warning was logged
                mock_logger.warning.assert_called_once()

                # Check response - upsert succeeded with warning
                assert response["status"] == "success"
                assert "Upserted 1 vectors" in response["message"]
                assert (
                    response.get("warning")
                    == "Training check failed but upsert succeeded"
                )


class TestTrainingStatusEndpoint:
    """Test the training status endpoint."""

    @pytest.mark.asyncio
    async def test_get_training_status_endpoint(self):
        """Test the /indexes/training-status endpoint."""
        with patch(
            "cyborgdb_service.api.routes.indexes.get_training_manager"
        ) as mock_get_manager:
            mock_manager = Mock()
            mock_manager.get_training_status.return_value = {
                "currently_training": "index1",
                "queued_indexes": ["index2"],
                "queue_size": 1,
                "retrain_threshold": 10000,
                "worker_running": True,
            }
            mock_get_manager.return_value = mock_manager

            # Call endpoint
            response = await get_training_status()

            # Verify response
            assert response["currently_training"] == "index1"
            assert response["queued_indexes"] == ["index2"]
            assert response["retrain_threshold"] == 10000
            mock_manager.get_training_status.assert_called_once()


class TestIntegrationScenarios:
    """Integration test scenarios for auto-trigger functionality."""

    def test_auto_training_trigger_for_untrained_index(self):
        """Test that auto training is triggered when threshold is exceeded for untrained index."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            with patch.object(
                TrainingManager, "_ensure_worker_running"
            ):  # Don't start worker
                manager = TrainingManager()
                manager._training_client = Mock()

                # Simulate an untrained index reaching the threshold
                # For untrained index, threshold = 1 * 10000 = 10000
                triggered = manager.trigger_training_if_needed(
                    index_name="new_index",
                    index_key_hex="a" * 64,
                    num_vectors=10001,  # Just above threshold
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                assert triggered
                assert manager.is_training("new_index")

    def test_auto_retraining_only_for_ivfflat(self):
        """Test that auto retraining is triggered only for ivfflat indexes when already trained."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            with patch.object(
                TrainingManager, "_ensure_worker_running"
            ):  # Don't start worker
                manager = TrainingManager()
                manager._training_client = Mock()

                # Test 1: IVFFlat - should allow retraining when trained
                triggered_ivfflat = manager.trigger_training_if_needed(
                    index_name="ivfflat_trained",
                    index_key_hex="b" * 64,
                    num_vectors=1_280_001,  # Above threshold for n_lists=128
                    n_lists=128,
                    is_trained=True,  # Already trained
                    index_type="ivfflat",
                )

                assert triggered_ivfflat
                assert manager.is_training("ivfflat_trained")

                # Test 2: IVF - should NOT allow retraining when trained
                triggered_ivf = manager.trigger_training_if_needed(
                    index_name="ivf_trained",
                    index_key_hex="c" * 64,
                    num_vectors=2_000_000,  # Well above threshold
                    n_lists=128,
                    is_trained=True,  # Already trained
                    index_type="ivf",
                )

                assert not triggered_ivf
                assert not manager.is_training("ivf_trained")

                # Test 3: IVFPQ - should NOT allow retraining when trained
                triggered_ivfpq = manager.trigger_training_if_needed(
                    index_name="ivfpq_trained",
                    index_key_hex="d" * 64,
                    num_vectors=2_000_000,  # Well above threshold
                    n_lists=128,
                    is_trained=True,  # Already trained
                    index_type="ivfpq",
                )

                assert not triggered_ivfpq
                assert not manager.is_training("ivfpq_trained")

                # Test 4: IVFSQ - should allow retraining when trained (like IVFFLAT)
                triggered_ivfsq = manager.trigger_training_if_needed(
                    index_name="ivfsq_trained",
                    index_key_hex="e" * 64,
                    num_vectors=1_280_001,  # Above threshold for n_lists=128
                    n_lists=128,
                    is_trained=True,  # Already trained
                    index_type="ivfsq",
                )

                assert triggered_ivfsq
                assert manager.is_training("ivfsq_trained")

    def test_multiple_upserts_trigger_once(self):
        """Test that multiple rapid upserts only trigger training once."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            with patch.object(
                TrainingManager, "_ensure_worker_running"
            ):  # Don't start worker
                manager = TrainingManager()
                manager._training_client = Mock()

                # First upsert triggers training
                triggered1 = manager.trigger_training_if_needed(
                    index_name="test_index",
                    index_key_hex="0" * 64,
                    num_vectors=15000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                # Second upsert should not trigger (already in queue)
                triggered2 = manager.trigger_training_if_needed(
                    index_name="test_index",
                    index_key_hex="0" * 64,
                    num_vectors=16000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                assert triggered1
                assert not triggered2
                assert manager.is_training("test_index")

    def test_multiple_indexes_get_queued(self):
        """Test that multiple indexes get queued for training."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            with patch.object(
                TrainingManager, "_ensure_worker_running"
            ):  # Don't start worker
                manager = TrainingManager()
                manager._training_client = Mock()

                # Trigger training for first index
                triggered1 = manager.trigger_training_if_needed(
                    index_name="index1",
                    index_key_hex="0" * 64,
                    num_vectors=15000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                # Trigger training for second index - should get queued
                triggered2 = manager.trigger_training_if_needed(
                    index_name="index2",
                    index_key_hex="1" * 64,
                    num_vectors=15000,
                    n_lists=128,
                    is_trained=False,
                    index_type="ivfflat",
                )

                assert triggered1
                assert triggered2  # Second index gets queued
                assert manager.is_training("index1")
                assert manager.is_training("index2")

                # Check status
                status = manager.get_training_status()
                # Both should be in training_indexes since worker hasn't started
                assert len(status["training_indexes"]) == 2
                assert "index1" in status["training_indexes"]
                assert "index2" in status["training_indexes"]

    def test_queue_position_tracking(self):
        """Test that queue positions are tracked correctly."""
        with patch.object(TrainingManager, "_initialize_training_client"):
            manager = TrainingManager()

            # Manually set up state for testing
            with manager._lock:
                manager._currently_training = "index0"

            # Queue some indexes
            for i in range(1, 4):
                job = TrainingJob(index_name=f"index{i}", index_key_hex=f"{i}" * 64)
                manager._queue.put(job)

            # Check positions
            assert manager.get_queue_position("index0") == 0  # Currently training
            assert manager.get_queue_position("index1") == 1
            assert manager.get_queue_position("index2") == 2
            assert manager.get_queue_position("index3") == 3
            assert manager.get_queue_position("index99") is None


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
