import pytest
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock
from comfy_resolve.resolver import Resolver
from comfy_resolve.parser import AssetReference
from comfy_resolve.sources.base import ResolutionResult, ResolvedSource


class TestResolver:
    """Tests for the Resolver class."""
    
    def test_resolver_initialization(self):
        """Test that resolver initializes with config."""
        config = {"comfyui_dir": "C:/test"}
        resolver = Resolver(config)
        assert resolver.config == config
        assert resolver.hf_source is not None
        assert resolver.civitai_source is not None
    
    @patch("comfy_resolve.resolver.get_model_path")
    def test_local_check_exact_match(self, mock_get_model_path, tmp_path):
        """Test that local check finds exact filename matches."""
        # Create a test file
        test_file = tmp_path / "test_model.safetensors"
        test_file.write_text("test content")
        
        mock_get_model_path.return_value = tmp_path
        
        config = {"comfyui_dir": str(tmp_path)}
        resolver = Resolver(config)
        
        asset = AssetReference(
            raw_name="test_model.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver._check_local(asset)
        
        assert result is not None
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "local"
        assert result.resolved_name == "test_model.safetensors"
    
    @patch("comfy_resolve.resolver.get_model_path")
    def test_local_check_fuzzy_match(self, mock_get_model_path, tmp_path):
        """Test that local check finds fuzzy matches."""
        # Create a test file with slightly different name
        test_file = tmp_path / "test_model_v2.safetensors"
        test_file.write_text("test content")
        
        mock_get_model_path.return_value = tmp_path
        
        config = {"comfyui_dir": str(tmp_path)}
        resolver = Resolver(config)
        
        asset = AssetReference(
            raw_name="test_model.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver._check_local(asset)
        
        assert result is not None
        assert result.confidence in ("high", "medium")
        assert result.source is not None
        assert result.source.source_type == "local"
    
    @patch("comfy_resolve.resolver.storage.get_mapping")
    def test_cache_check_returns_cached_result(self, mock_get_mapping):
        """Test that cache check returns cached mappings."""
        mock_get_mapping.return_value = {
            "resolved_name": "cached_model.safetensors",
            "confidence": "high",
            "source": {
                "type": "huggingface",
                "repo": "test/repo",
                "file": "cached_model.safetensors",
                "download_url": "https://example.com/model",
                "sha256": "abc123",
                "file_size_bytes": 1000000
            }
        }
        
        config = {"comfyui_dir": "C:/test"}
        resolver = Resolver(config)
        
        asset = AssetReference(
            raw_name="cached_model.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver._check_cache(asset)
        
        assert result is not None
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "huggingface"
        assert result.source.repo == "test/repo"
    
    def test_config_mappings_direct_url(self):
        """Test that config mappings work for direct URLs."""
        config = {
            "comfyui_dir": "C:/test",
            "mappings": {
                "custom_model.safetensors": {
                    "type": "direct_url",
                    "url": "https://example.com/custom_model.safetensors",
                    "sha256": "def456"
                }
            }
        }
        resolver = Resolver(config)
        
        asset = AssetReference(
            raw_name="custom_model.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver._check_config_mappings(asset)
        
        assert result is not None
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "direct_url"
        assert result.source.download_url == "https://example.com/custom_model.safetensors"
    
    @patch("comfy_resolve.resolver.HuggingFaceSource")
    def test_huggingface_check(self, mock_hf_class):
        """Test that HuggingFace check works."""
        mock_hf_instance = Mock()
        mock_hf_instance.search.return_value = [
            ResolutionResult(
                input_name="test.safetensors",
                resolved_name="test.safetensors",
                source=ResolvedSource(
                    source_type="huggingface",
                    repo="test/repo",
                    file="test.safetensors",
                    download_url="https://hf.co/test",
                    sha256=None,
                    file_size_bytes=None
                ),
                confidence="high",
                error=None
            )
        ]
        mock_hf_class.return_value = mock_hf_instance
        
        config = {"comfyui_dir": "C:/test"}
        resolver = Resolver(config)
        resolver.hf_source = mock_hf_instance
        
        asset = AssetReference(
            raw_name="test.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver._check_huggingface(asset)
        
        assert result is not None
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "huggingface"
    
    def test_heuristic_match(self):
        """Test that heuristic matching works for known models."""
        config = {"comfyui_dir": "C:/test"}
        resolver = Resolver(config)
        
        asset = AssetReference(
            raw_name="clip_l.safetensors",
            asset_type="clip",
            node_id="1",
            node_type="CLIPLoader"
        )
        
        result = resolver._heuristic_match(asset)
        
        assert result is not None
        assert result.confidence == "low"
        assert "clip_l.safetensors" in result.resolved_name.lower()
    
    @patch("comfy_resolve.resolver.get_model_path")
    def test_resolve_one_short_circuits_at_local(self, mock_get_model_path, tmp_path):
        """Test that resolve_one short-circuits when local file is found."""
        # Create a test file
        test_file = tmp_path / "local_model.safetensors"
        test_file.write_text("test content")
        
        mock_get_model_path.return_value = tmp_path
        
        config = {"comfyui_dir": str(tmp_path)}
        resolver = Resolver(config)
        
        # Mock other methods to ensure they're not called
        resolver._check_cache = Mock(return_value=None)
        resolver._check_huggingface = Mock(return_value=None)
        
        asset = AssetReference(
            raw_name="local_model.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver.resolve_one(asset)
        
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "local"
        # Cache and HF should not be called
        resolver._check_cache.assert_not_called()
        resolver._check_huggingface.assert_not_called()
    
    @patch("comfy_resolve.resolver.storage.get_mapping")
    @patch("comfy_resolve.resolver.get_model_path")
    def test_resolve_one_short_circuits_at_cache(self, mock_get_model_path, mock_get_mapping, tmp_path):
        """Test that resolve_one short-circuits when cache hit occurs."""
        mock_get_model_path.return_value = tmp_path  # Empty directory
        
        mock_get_mapping.return_value = {
            "resolved_name": "cached.safetensors",
            "confidence": "high",
            "source": {
                "type": "huggingface",
                "repo": "test/repo",
                "file": "cached.safetensors",
                "download_url": "https://example.com",
                "sha256": None,
                "file_size_bytes": None
            }
        }
        
        config = {"comfyui_dir": str(tmp_path)}
        resolver = Resolver(config)
        
        # Mock HF to ensure it's not called
        resolver._check_huggingface = Mock(return_value=None)
        
        asset = AssetReference(
            raw_name="cached.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver.resolve_one(asset)
        
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "huggingface"
        # HF should not be called
        resolver._check_huggingface.assert_not_called()
    
    @patch("comfy_resolve.resolver.HuggingFaceSource")
    @patch("comfy_resolve.resolver.storage.get_mapping")
    @patch("comfy_resolve.resolver.get_model_path")
    def test_resolve_one_reaches_huggingface(self, mock_get_model_path, mock_get_mapping, mock_hf_class, tmp_path):
        """Test that resolve_one reaches HuggingFace when nothing else matches."""
        mock_get_model_path.return_value = tmp_path  # Empty directory
        mock_get_mapping.return_value = None  # No cache
        
        mock_hf_instance = Mock()
        mock_hf_instance.search.return_value = [
            ResolutionResult(
                input_name="test.safetensors",
                resolved_name="test.safetensors",
                source=ResolvedSource(
                    source_type="huggingface",
                    repo="test/repo",
                    file="test.safetensors",
                    download_url="https://hf.co/test",
                    sha256=None,
                    file_size_bytes=None
                ),
                confidence="high",
                error=None
            )
        ]
        mock_hf_class.return_value = mock_hf_instance
        
        config = {"comfyui_dir": str(tmp_path)}
        resolver = Resolver(config)
        resolver.hf_source = mock_hf_instance
        
        asset = AssetReference(
            raw_name="test.safetensors",
            asset_type="checkpoints",
            node_id="1",
            node_type="CheckpointLoader"
        )
        
        result = resolver.resolve_one(asset)
        
        assert result.confidence == "high"
        assert result.source is not None
        assert result.source.source_type == "huggingface"
        mock_hf_instance.search.assert_called_once()
    
    def test_resolve_all(self):
        """Test that resolve_all processes multiple assets."""
        config = {"comfyui_dir": "C:/test"}
        resolver = Resolver(config)
        
        # Mock resolve_one
        resolver.resolve_one = Mock(side_effect=lambda asset: ResolutionResult(
            input_name=asset.raw_name,
            resolved_name=asset.raw_name,
            source=None,
            confidence="unknown",
            error=None
        ))
        
        assets = [
            AssetReference("model1.safetensors", "checkpoints", "1", "CheckpointLoader"),
            AssetReference("model2.safetensors", "checkpoints", "2", "CheckpointLoader"),
        ]
        
        results = resolver.resolve_all(assets)
        
        assert len(results) == 2
        assert resolver.resolve_one.call_count == 2
