"""
Test for smart auth profile rotation fix
Tests that rotation properly considers model-specific cooldowns
"""

import json
import os
import tempfile
import unittest
from datetime import datetime, timedelta
from unittest.mock import patch

# Import the functions we want to test
from browser_utils.auth_rotation import _find_best_profile_in_dirs, _normalize_model_id


class TestSmartRotationFix(unittest.TestCase):
    """Test smart rotation logic with model-specific cooldowns"""

    def setUp(self):
        """Set up test environment"""
        self.test_dir = tempfile.mkdtemp()
        self.saved_dir = os.path.join(self.test_dir, "saved")
        self.emergency_dir = os.path.join(self.test_dir, "emergency")
        os.makedirs(self.saved_dir, exist_ok=True)
        os.makedirs(self.emergency_dir, exist_ok=True)

        # Create test profile files
        self.profile1 = os.path.join(self.saved_dir, "profile1.json")
        self.profile2 = os.path.join(self.saved_dir, "profile2.json")
        self.profile3 = os.path.join(self.saved_dir, "profile3.json")
        self.emergency_profile = os.path.join(self.emergency_dir, "emergency.json")

        # Create dummy profile content
        profile_content = {"cookies": [{"name": "test", "value": "test"}]}

        for profile_path in [
            self.profile1,
            self.profile2,
            self.profile3,
            self.emergency_profile,
        ]:
            with open(profile_path, "w") as f:
                json.dump(profile_content, f)

    def tearDown(self):
        """Clean up test environment"""
        import shutil

        shutil.rmtree(self.test_dir, ignore_errors=True)

    def test_normalize_model_id(self):
        """Test model ID normalization"""
        test_cases = [
            ("gemini 3 pro preview", "gemini-3-pro-preview"),
            ("gemini-2.5-pro", "gemini-2.5-pro"),
            (
                "gemini 2.5 pro",
                "gemini-2.5-pro",
            ),  # Should preserve dots for known models
            ("default", "default"),
            ("", "default"),
            (None, "default"),
        ]

        for input_id, expected in test_cases:
            result = _normalize_model_id(input_id)
            self.assertEqual(result, expected, f"Failed for input: {input_id}")

    def test_find_best_profile_gemini_3_pro_preview(self):
        """Test that profiles with gemini-3-pro-preview cooldown are excluded"""
        # Mock cooldown data similar to user's config
        cooldown_data = {
            self.profile1: {
                "gemini-3-pro-preview": (
                    datetime.now() + timedelta(hours=1)
                ).timestamp(),
                "default": (datetime.now() + timedelta(hours=1)).timestamp(),
            },
            self.profile2: {
                "default": (datetime.now() + timedelta(hours=1)).timestamp()
            },
            # profile3 has no cooldown - should be selected
            # emergency_profile has no cooldown - should be selected
        }

        with (
            patch("browser_utils.auth_rotation._COOLDOWN_PROFILES", cooldown_data),
            patch("browser_utils.auth_rotation.get_profile_usage", return_value=1),
        ):
            # Test with gemini-3-pro-preview - should exclude profile1 but include others
            best_profile = _find_best_profile_in_dirs(
                [self.saved_dir, self.emergency_dir],
                target_model_id="gemini 3 pro preview",
            )

            # Should not be profile1 (has gemini-3-pro-preview cooldown)
            self.assertNotEqual(best_profile, self.profile1)
            # Should be either profile2, profile3, or emergency_profile
            self.assertIn(
                best_profile, [self.profile2, self.profile3, self.emergency_profile]
            )

    def test_find_best_profile_no_model_specific_cooldown(self):
        """Test that profiles without specific model cooldown are included"""
        # Mock cooldown data - only default cooldown
        cooldown_data = {
            self.profile1: {
                "default": (datetime.now() + timedelta(hours=1)).timestamp()
            },
            # Other profiles have no cooldown
        }

        with (
            patch("browser_utils.auth_rotation._COOLDOWN_PROFILES", cooldown_data),
            patch("browser_utils.auth_rotation.get_profile_usage", return_value=1),
        ):
            # Test with gemini-3-pro-preview - should include all profiles since none have specific cooldown
            best_profile = _find_best_profile_in_dirs(
                [self.saved_dir, self.emergency_dir],
                target_model_id="gemini 3 pro preview",
            )

            # Should be able to select any profile (including profile1)
            self.assertIsNotNone(best_profile)

    def test_find_best_profile_global_cooldown_excludes_all(self):
        """Test that global cooldown excludes all profiles"""
        # Mock cooldown data with global cooldown
        cooldown_data = {
            self.profile1: {
                "global": (datetime.now() + timedelta(hours=1)).timestamp()
            },
            self.profile2: {
                "global": (datetime.now() + timedelta(hours=1)).timestamp()
            },
            self.profile3: {
                "global": (datetime.now() + timedelta(hours=1)).timestamp()
            },
            self.emergency_profile: {
                "global": (datetime.now() + timedelta(hours=1)).timestamp()
            },
        }

        with (
            patch("browser_utils.auth_rotation._COOLDOWN_PROFILES", cooldown_data),
            patch("browser_utils.auth_rotation.get_profile_usage", return_value=1),
        ):
            # Test with any model - should return None since all have global cooldown
            best_profile = _find_best_profile_in_dirs(
                [self.saved_dir, self.emergency_dir],
                target_model_id="gemini 3 pro preview",
            )

            # Should return None since all profiles are in global cooldown
            self.assertIsNone(best_profile)

    def test_find_best_profile_expired_cooldown(self):
        """Test that expired cooldowns don't exclude profiles"""
        # Mock cooldown data with expired cooldown
        cooldown_data = {
            self.profile1: {
                "gemini-3-pro-preview": (
                    datetime.now() - timedelta(hours=1)
                ).timestamp(),  # Expired
                "default": (datetime.now() - timedelta(hours=1)).timestamp(),  # Expired
            },
        }

        with (
            patch("browser_utils.auth_rotation._COOLDOWN_PROFILES", cooldown_data),
            patch("browser_utils.auth_rotation.get_profile_usage", return_value=1),
        ):
            # Test with gemini-3-pro-preview - should include profile1 since cooldown expired
            best_profile = _find_best_profile_in_dirs(
                [self.saved_dir, self.emergency_dir],
                target_model_id="gemini 3 pro preview",
            )

            # Should be able to select profile1 since cooldown expired
            self.assertIsNotNone(best_profile)
            # Profile1 should be selectable since its cooldown expired, but we can't guarantee
            # it will be selected over other profiles due to usage-based sorting
            valid_profiles = [
                self.profile1,
                self.profile2,
                self.profile3,
                self.emergency_profile,
            ]
            self.assertIn(best_profile, valid_profiles)


if __name__ == "__main__":
    unittest.main()