import pytest
from unittest.mock import Mock, patch
from comfy_resolve.sources.civitai import CivitaiSource
from comfy_resolve.sources.base import ResolutionResult, ResolvedSource


class TestCivitaiSource:
    """Tests for Civitai source."""
    
    def test_name_property(self):
        """Test that the name property returns 'civitai'."""
        source = CivitaiSource()
        assert source.name == "civitai"
    
    def test_is_available_with_env_var(self):
        """Test that is_available returns True when env var is set."""
        with patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"}):
            source = CivitaiSource()
            assert source.is_available({}) is True
    
    def test_is_available_with_config(self):
        """Test that is_available returns True when config has token."""
        source = CivitaiSource()
        config = {"api_keys": {"civitai": "test_token"}}
        assert source.is_available(config) is True
    
    def test_is_available_without_token(self):
        """Test that is_available returns False without token."""
        source = CivitaiSource()
        assert source.is_available({}) is False
    
    def test_search_without_token_returns_error(self):
        """Test that search without token returns error result."""
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
        assert "token required" in results[0].error.lower()
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_exact_match_high_confidence(self, mock_get):
        """Test that exact filename matches return high confidence."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "items": [
                {
                    "id": 12345,
                    "name": "Test Model",
                    "modelVersions": [
                        {
                            "files": [
                                {
                                    "name": "model.safetensors",
                                    "downloadUrl": "https://civitai.com/download/12345",
                                    "sizeKB": 1000,
                                    "hashes": {"SHA256": "abc123"}
                                }
                            ]
                        }
                    ]
                }
            ]
        }
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) > 0
        assert results[0].confidence == "high"
        assert results[0].resolved_name == "model.safetensors"
        assert results[0].source is not None
        assert results[0].source.repo == "civitai:12345"
        assert results[0].source.sha256 == "abc123"
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_401_returns_token_error(self, mock_get):
        """Test that 401 errors return token invalid message."""
        mock_response = Mock()
        mock_response.status_code = 401
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
        assert "token invalid" in results[0].error.lower()
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_429_returns_rate_limit_error(self, mock_get):
        """Test that 429 errors return rate limit message."""
        mock_response = Mock()
        mock_response.status_code = 429
        mock_response.headers = {"Retry-After": "60"}
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
        assert "rate limit" in results[0].error.lower()
        assert "60" in results[0].error
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_api_error_returns_error_result(self, mock_get):
        """Test that API errors return a result with error message."""
        mock_response = Mock()
        mock_response.status_code = 500
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
        assert "API error" in results[0].error
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_network_error_returns_error_result(self, mock_get):
        """Test that network errors return a result with error message."""
        mock_get.side_effect = Exception("Network timeout")
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_fuzzy_match_medium_confidence(self, mock_get):
        """Test that fuzzy matches above 85% return medium confidence."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "items": [
                {
                    "id": 12345,
                    "name": "Test Model",
                    "modelVersions": [
                        {
                            "files": [
                                {
                                    "name": "model_v2.safetensors",
                                    "downloadUrl": "https://civitai.com/download/12345",
                                    "sizeKB": 1000,
                                    "hashes": {}
                                }
                            ]
                        }
                    ]
                }
            ]
        }
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) > 0
        # The similarity should be high enough for medium or low confidence
        assert results[0].confidence in ["medium", "low"]
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_converts_size_kb_to_bytes(self, mock_get):
        """Test that file size is correctly converted from KB to bytes."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "items": [
                {
                    "id": 12345,
                    "name": "Test Model",
                    "modelVersions": [
                        {
                            "files": [
                                {
                                    "name": "model.safetensors",
                                    "downloadUrl": "https://civitai.com/download/12345",
                                    "sizeKB": 1024,  # 1024 KB = 1048576 bytes
                                    "hashes": {}
                                }
                            ]
                        }
                    ]
                }
            ]
        }
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) > 0
        assert results[0].source is not None
        assert results[0].source.file_size_bytes == 1048576
    
    @patch.dict("os.environ", {"COMFY_CIVITAI_TOKEN": "test_token"})
    @patch("comfy_resolve.sources.civitai.requests.get")
    def test_search_sorts_by_confidence(self, mock_get):
        """Test that search results are sorted by confidence."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "items": [
                {
                    "id": 12345,
                    "name": "Test Model",
                    "modelVersions": [
                        {
                            "files": [
                                {
                                    "name": "test.safetensors",
                                    "downloadUrl": "https://civitai.com/download/1",
                                    "sizeKB": 1000,
                                    "hashes": {}
                                },
                                {
                                    "name": "test_v2.safetensors",
                                    "downloadUrl": "https://civitai.com/download/2",
                                    "sizeKB": 1000,
                                    "hashes": {}
                                }
                            ]
                        }
                    ]
                }
            ]
        }
        mock_get.return_value = mock_response
        
        source = CivitaiSource()
        results = source.search("test.safetensors", "checkpoints")
        
        # First result should be the exact match with high confidence
        if len(results) > 0:
            assert results[0].confidence == "high"
            assert results[0].resolved_name == "test.safetensors"
