import pytest
import hashlib
from pathlib import Path
from comfy_resolve.downloader import Downloader, DownloadTask, DownloadResult
from comfy_resolve.sources.base import ResolutionResult, ResolvedSource


class TestDownloader:
    """Tests for the Downloader class."""
    
    def test_downloader_initialization(self):
        """Test that downloader initializes with config."""
        config = {"download": {"max_retries": 5, "concurrent": 3}}
        downloader = Downloader(config)
        assert downloader.max_retries == 5
        assert downloader.concurrent == 3
    
    def test_should_skip_file_not_exists(self, tmp_path):
        """Test that _should_skip returns False when file doesn't exist."""
        config = {}
        downloader = Downloader(config)
        
        source = ResolvedSource(
            source_type="direct_url",
            repo="",
            file="test.safetensors",
            download_url="https://example.com/test.safetensors",
            sha256="abc123",
            file_size_bytes=1000,
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        assert not downloader._should_skip(task)
    
    def test_should_skip_file_exists_no_hash(self, tmp_path):
        """Test that _should_skip returns True when file exists and no hash available."""
        config = {}
        downloader = Downloader(config)
        
        # Create a test file
        test_file = tmp_path / "test.safetensors"
        test_file.write_text("test content")
        
        source = ResolvedSource(
            source_type="direct_url",
            repo="",
            file="test.safetensors",
            download_url="https://example.com/test.safetensors",
            sha256=None,  # No hash
            file_size_bytes=1000,
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        assert downloader._should_skip(task)
    
    def test_should_skip_file_exists_hash_match(self, tmp_path):
        """Test that _should_skip returns True when file exists and hash matches."""
        config = {}
        downloader = Downloader(config)
        
        # Create a test file
        test_file = tmp_path / "test.safetensors"
        content = b"test content"
        test_file.write_bytes(content)
        
        # Calculate actual hash
        actual_hash = hashlib.sha256(content).hexdigest()
        
        source = ResolvedSource(
            source_type="direct_url",
            repo="",
            file="test.safetensors",
            download_url="https://example.com/test.safetensors",
            sha256=actual_hash,
            file_size_bytes=len(content),
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        assert downloader._should_skip(task)
    
    def test_should_skip_file_exists_hash_mismatch(self, tmp_path):
        """Test that _should_skip returns False when file exists but hash doesn't match."""
        config = {}
        downloader = Downloader(config)
        
        # Create a test file
        test_file = tmp_path / "test.safetensors"
        test_file.write_bytes(b"test content")
        
        # Use a different hash
        wrong_hash = "0" * 64
        
        source = ResolvedSource(
            source_type="direct_url",
            repo="",
            file="test.safetensors",
            download_url="https://example.com/test.safetensors",
            sha256=wrong_hash,
            file_size_bytes=1000,
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        assert not downloader._should_skip(task)
    
    def test_verify_integrity_correct_hash(self, tmp_path):
        """Test that _verify_integrity returns True for correct hash."""
        config = {}
        downloader = Downloader(config)
        
        # Create a test file
        test_file = tmp_path / "test.txt"
        content = b"test content for integrity check"
        test_file.write_bytes(content)
        
        # Calculate hash
        expected_hash = hashlib.sha256(content).hexdigest()
        
        assert downloader._verify_integrity(test_file, expected_hash)
    
    def test_verify_integrity_wrong_hash(self, tmp_path):
        """Test that _verify_integrity returns False for wrong hash."""
        config = {}
        downloader = Downloader(config)
        
        # Create a test file
        test_file = tmp_path / "test.txt"
        test_file.write_bytes(b"test content")
        
        # Use wrong hash
        wrong_hash = "0" * 64
        
        assert not downloader._verify_integrity(test_file, wrong_hash)
    
    def test_verify_integrity_file_not_exists(self, tmp_path):
        """Test that _verify_integrity returns False when file doesn't exist."""
        config = {}
        downloader = Downloader(config)
        
        non_existent_file = tmp_path / "nonexistent.txt"
        
        assert not downloader._verify_integrity(non_existent_file, "abc123")
    
    def test_check_disk_space(self, tmp_path):
        """Test that check_disk_space returns sensible values."""
        config = {}
        downloader = Downloader(config)
        
        source = ResolvedSource(
            source_type="direct_url",
            repo="",
            file="test.safetensors",
            download_url="https://example.com/test.safetensors",
            sha256=None,
            file_size_bytes=1000000,  # 1 MB
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        enough, required, available = downloader.check_disk_space([task])
        
        # Should have enough space for 1 MB
        assert enough
        assert required == 1000000
        assert available > 0
    
    def test_check_disk_space_empty_tasks(self):
        """Test that check_disk_space handles empty task list."""
        config = {}
        downloader = Downloader(config)
        
        enough, required, available = downloader.check_disk_space([])
        
        assert enough
        assert required == 0
        assert available == 0
    
    def test_get_available_space(self, tmp_path):
        """Test that _get_available_space returns positive value."""
        config = {}
        downloader = Downloader(config)
        
        available = downloader._get_available_space(tmp_path)
        
        assert available > 0
    
    def test_download_one_local_source_skipped(self, tmp_path):
        """Test that local sources are skipped."""
        config = {}
        downloader = Downloader(config)
        
        source = ResolvedSource(
            source_type="local",
            repo=str(tmp_path),
            file="test.safetensors",
            download_url=None,
            sha256=None,
            file_size_bytes=1000,
        )
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=source,
            confidence="high",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        download_result = downloader.download_one(task)
        
        assert download_result.success
        assert download_result.skipped
    
    def test_download_one_no_source(self, tmp_path):
        """Test that download fails when no source is available."""
        config = {}
        downloader = Downloader(config)
        
        result = ResolutionResult(
            input_name="test.safetensors",
            resolved_name="test.safetensors",
            source=None,
            confidence="unknown",
            error=None,
        )
        
        task = DownloadTask(
            result=result,
            destination=tmp_path,
            asset_type="checkpoints",
        )
        
        download_result = downloader.download_one(task)
        
        assert not download_result.success
        assert not download_result.skipped
        assert download_result.error is not None
        assert "No source available" in download_result.error
