import pytest
from unittest.mock import Mock, patch
from comfy_resolve.sources.huggingface import HuggingFaceSource
from comfy_resolve.sources.base import ResolutionResult, ResolvedSource


class TestHuggingFaceSource:
    """Tests for HuggingFace source."""
    
    def test_name_property(self):
        """Test that the name property returns 'huggingface'."""
        source = HuggingFaceSource()
        assert source.name == "huggingface"
    
    def test_is_available_always_true(self):
        """Test that is_available always returns True."""
        source = HuggingFaceSource()
        assert source.is_available({}) is True
        assert source.is_available({"api_keys": {"huggingface": "token"}}) is True
    
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_search_exact_match_high_confidence(self, mock_get):
        """Test that exact filename matches return high confidence."""
        # Mock the search API response
        mock_search_response = Mock()
        mock_search_response.status_code = 200
        mock_search_response.json.return_value = [
            {"id": "test-user/test-model", "modelId": "test-user/test-model"}
        ]
        
        # Mock the file list API response
        mock_files_response = Mock()
        mock_files_response.status_code = 200
        mock_files_response.json.return_value = [
            {"type": "file", "path": "model.safetensors"}
        ]
        
        # Mock the metadata API response
        mock_metadata_response = Mock()
        mock_metadata_response.status_code = 200
        mock_metadata_response.json.return_value = {
            "lfs": {"sha256": "abc123"},
            "size": 1000000
        }
        
        # Set up the mock to return different responses for different URLs
        def get_side_effect(url, *args, **kwargs):
            if "/models?" in url or "/models" == url.split("?")[0]:
                return mock_search_response
            elif "/tree/main" in url and "/tree/main/" not in url:
                return mock_files_response
            else:
                return mock_metadata_response
        
        mock_get.side_effect = get_side_effect
        
        source = HuggingFaceSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) > 0
        assert results[0].confidence == "high"
        assert results[0].input_name == "model.safetensors"
        assert results[0].resolved_name == "model.safetensors"
        assert results[0].source is not None
        assert results[0].source.repo == "test-user/test-model"
    
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_search_fuzzy_match_medium_confidence(self, mock_get):
        """Test that fuzzy matches above 85% return medium confidence."""
        mock_search_response = Mock()
        mock_search_response.status_code = 200
        mock_search_response.json.return_value = [
            {"id": "test-user/test-model"}
        ]
        
        mock_files_response = Mock()
        mock_files_response.status_code = 200
        mock_files_response.json.return_value = [
            {"type": "file", "path": "model_v2.safetensors"}  # Similar but not exact
        ]
        
        mock_metadata_response = Mock()
        mock_metadata_response.status_code = 200
        mock_metadata_response.json.return_value = {}
        
        def get_side_effect(url, *args, **kwargs):
            if "/models?" in url or "/models" == url.split("?")[0]:
                return mock_search_response
            elif "/tree/main" in url and "/tree/main/" not in url:
                return mock_files_response
            else:
                return mock_metadata_response
        
        mock_get.side_effect = get_side_effect
        
        source = HuggingFaceSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) > 0
        # The similarity should be high enough for medium confidence
        assert results[0].confidence in ["medium", "low"]
    
    @patch("comfy_resolve.sources.huggingface.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 = HuggingFaceSource()
        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("comfy_resolve.sources.huggingface.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 = HuggingFaceSource()
        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("comfy_resolve.sources.huggingface.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 = HuggingFaceSource()
        results = source.search("model.safetensors", "checkpoints")
        
        assert len(results) == 1
        assert results[0].confidence == "unknown"
        assert results[0].error is not None
    
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_list_repo_files_403_raises_gated_error(self, mock_get):
        """Test that 403 errors raise helpful gated model message."""
        mock_response = Mock()
        mock_response.status_code = 403
        mock_get.return_value = mock_response
        
        source = HuggingFaceSource()
        
        with pytest.raises(Exception) as exc_info:
            source._list_repo_files("test-user/gated-model")
        
        assert "gated" in str(exc_info.value).lower()
        assert "huggingface.co" in str(exc_info.value)
    
    @patch("comfy_resolve.sources.huggingface.requests.head")
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_get_file_metadata_returns_sha256_and_size(self, mock_get, mock_head):
        """Test that get_file_metadata extracts SHA-256 and file size."""
        # Mock HEAD request with headers
        mock_head_response = Mock()
        mock_head_response.status_code = 200
        mock_head_response.headers = {
            "X-Linked-Etag": '"' + "a" * 64 + '"',  # Valid SHA-256
            "Content-Length": "1234567890"
        }
        mock_head.return_value = mock_head_response
        
        source = HuggingFaceSource()
        metadata = source.get_file_metadata("test-user/test-model", "model.safetensors")
        
        assert "sha256" in metadata
        assert metadata["sha256"] == "a" * 64
        assert "size" in metadata
        assert metadata["size"] == 1234567890
    
    @patch("comfy_resolve.sources.huggingface.requests.head")
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_get_file_metadata_fallback_to_api(self, mock_get, mock_head):
        """Test that get_file_metadata falls back to API if headers don't have SHA."""
        # Mock HEAD request without SHA in headers
        mock_head_response = Mock()
        mock_head_response.status_code = 200
        mock_head_response.headers = {
            "Content-Length": "1234567890"
        }
        mock_head.return_value = mock_head_response
        
        # Mock API request with SHA
        mock_api_response = Mock()
        mock_api_response.status_code = 200
        mock_api_response.json.return_value = {
            "lfs": {"sha256": "b" * 64},
            "size": 1234567890
        }
        mock_get.return_value = mock_api_response
        
        source = HuggingFaceSource()
        metadata = source.get_file_metadata("test-user/test-model", "model.safetensors")
        
        assert "sha256" in metadata
        assert metadata["sha256"] == "b" * 64
    
    @patch("comfy_resolve.sources.huggingface.requests.head")
    def test_get_file_metadata_401_raises_token_error(self, mock_head):
        """Test that 401 errors raise token invalid message."""
        mock_response = Mock()
        mock_response.status_code = 401
        mock_head.return_value = mock_response
        
        source = HuggingFaceSource()
        
        with pytest.raises(Exception) as exc_info:
            source.get_file_metadata("test-user/test-model", "model.safetensors")
        
        assert "token invalid" in str(exc_info.value).lower()
    
    @patch("comfy_resolve.sources.huggingface.requests.head")
    def test_get_file_metadata_403_raises_gated_error(self, mock_head):
        """Test that 403 errors raise gated model message."""
        mock_response = Mock()
        mock_response.status_code = 403
        mock_head.return_value = mock_response
        
        source = HuggingFaceSource()
        
        with pytest.raises(Exception) as exc_info:
            source.get_file_metadata("test-user/gated-model", "model.safetensors")
        
        assert "gated" in str(exc_info.value).lower()
    
    def test_get_file_metadata_safe_returns_empty_on_error(self):
        """Test that _get_file_metadata_safe returns empty dict on error."""
        source = HuggingFaceSource()
        
        # This should not raise an exception
        metadata = source._get_file_metadata_safe("invalid/repo", "nonexistent.file")
        
        assert isinstance(metadata, dict)
        assert len(metadata) == 0
    
    @patch("comfy_resolve.sources.huggingface.requests.get")
    def test_search_sorts_by_confidence(self, mock_get):
        """Test that search results are sorted by confidence."""
        mock_search_response = Mock()
        mock_search_response.status_code = 200
        mock_search_response.json.return_value = [
            {"id": "user1/model1"},
            {"id": "user2/model2"}
        ]
        
        # First repo has exact match, second has fuzzy match
        mock_files_response_1 = Mock()
        mock_files_response_1.status_code = 200
        mock_files_response_1.json.return_value = [
            {"type": "file", "path": "test.safetensors"}
        ]
        
        mock_files_response_2 = Mock()
        mock_files_response_2.status_code = 200
        mock_files_response_2.json.return_value = [
            {"type": "file", "path": "test_v2.safetensors"}
        ]
        
        mock_metadata_response = Mock()
        mock_metadata_response.status_code = 200
        mock_metadata_response.json.return_value = {}
        
        call_count = [0]
        
        def get_side_effect(url, *args, **kwargs):
            if "/models?" in url or "/models" == url.split("?")[0]:
                return mock_search_response
            elif "/tree/main" in url and "/tree/main/" not in url:
                call_count[0] += 1
                if call_count[0] == 1:
                    return mock_files_response_1
                else:
                    return mock_files_response_2
            else:
                return mock_metadata_response
        
        mock_get.side_effect = get_side_effect
        
        source = HuggingFaceSource()
        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"
