from pathlib import Path
import pytest
from comfy_resolve.parser import parse_workflow, extract_assets, AssetReference


# Get the fixtures directory
FIXTURES_DIR = Path(__file__).parent / "fixtures"


def test_standard_workflow_extracts_correct_assets():
    """Test that standard workflow format extracts correct assets and types."""
    workflow_path = FIXTURES_DIR / "standard_workflow.json"
    assets = parse_workflow(workflow_path)
    
    # Should find 3 assets: checkpoint, lora, vae
    assert len(assets) == 3
    
    # Check checkpoint
    checkpoint = next((a for a in assets if a.raw_name == "sd_xl_base_1.0.safetensors"), None)
    assert checkpoint is not None
    assert checkpoint.asset_type == "checkpoints"
    assert checkpoint.node_type == "CheckpointLoaderSimple"
    
    # Check lora
    lora = next((a for a in assets if a.raw_name == "detail_tweaker_xl.safetensors"), None)
    assert lora is not None
    assert lora.asset_type == "loras"
    assert lora.node_type == "LoraLoader"
    
    # Check vae
    vae = next((a for a in assets if a.raw_name == "sdxl_vae.safetensors"), None)
    assert vae is not None
    assert vae.asset_type == "vae"
    assert vae.node_type == "VAELoader"


def test_api_format_workflow_extracts_same_results():
    """Test that API format workflow extracts the same assets as standard format."""
    standard_path = FIXTURES_DIR / "standard_workflow.json"
    api_path = FIXTURES_DIR / "api_format_workflow.json"
    
    standard_assets = parse_workflow(standard_path)
    api_assets = parse_workflow(api_path)
    
    # Should have same number of assets
    assert len(standard_assets) == len(api_assets)
    
    # Extract asset names for comparison
    standard_names = {a.raw_name for a in standard_assets}
    api_names = {a.raw_name for a in api_assets}
    
    # Should have same asset names
    assert standard_names == api_names
    
    # Check that asset types match for each asset
    for std_asset in standard_assets:
        api_asset = next((a for a in api_assets if a.raw_name == std_asset.raw_name), None)
        assert api_asset is not None
        assert api_asset.asset_type == std_asset.asset_type


def test_unknown_nodes_fall_back_to_extension_heuristic():
    """Test that unknown node types fall back to extension heuristic."""
    workflow_path = FIXTURES_DIR / "unknown_nodes_workflow.json"
    assets = parse_workflow(workflow_path)
    
    # Should find 3 assets with known extensions (.safetensors, .pth)
    # The .txt file should not be detected as it's not in ASSET_EXTENSIONS
    assert len(assets) == 3
    
    # Check that assets were found
    asset_names = {a.raw_name for a in assets}
    assert "custom_model_v2.safetensors" in asset_names
    assert "my_custom_lora.safetensors" in asset_names
    assert "some_upscale_model.pth" in asset_names
    
    # Check that lora was inferred from filename
    lora = next((a for a in assets if "lora" in a.raw_name.lower()), None)
    assert lora is not None
    assert lora.asset_type == "loras"
    
    # Check that upscale was inferred from filename
    upscale = next((a for a in assets if "upscale" in a.raw_name.lower()), None)
    assert upscale is not None
    assert upscale.asset_type == "upscale_models"


def test_deduplication_same_filename_in_multiple_nodes():
    """Test that same filename in two nodes results in one entry."""
    # Create a workflow with duplicate assets
    workflow = {
        "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {
                "ckpt_name": "model.safetensors"
            }
        },
        "2": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {
                "ckpt_name": "model.safetensors"
            }
        },
        "3": {
            "class_type": "LoraLoader",
            "inputs": {
                "lora_name": "lora.safetensors"
            }
        },
        "4": {
            "class_type": "LoraLoader",
            "inputs": {
                "lora_name": "lora.safetensors"
            }
        }
    }
    
    assets = extract_assets(workflow)
    
    # Should only have 2 unique assets despite 4 references
    assert len(assets) == 2
    
    # Check that we have one checkpoint and one lora
    asset_names = {a.raw_name for a in assets}
    assert "model.safetensors" in asset_names
    assert "lora.safetensors" in asset_names


def test_empty_workflow():
    """Test that empty workflow returns empty list."""
    workflow = {}
    assets = extract_assets(workflow)
    assert len(assets) == 0


def test_workflow_with_no_assets():
    """Test workflow with nodes but no asset references."""
    workflow = {
        "1": {
            "class_type": "EmptyLatentImage",
            "inputs": {
                "width": 512,
                "height": 512,
                "batch_size": 1
            }
        },
        "2": {
            "class_type": "KSampler",
            "inputs": {
                "seed": 42,
                "steps": 20
            }
        }
    }
    
    assets = extract_assets(workflow)
    assert len(assets) == 0


def test_asset_reference_dataclass():
    """Test AssetReference dataclass creation."""
    asset = AssetReference(
        raw_name="test.safetensors",
        asset_type="checkpoints",
        node_id="1",
        node_type="CheckpointLoaderSimple"
    )
    
    assert asset.raw_name == "test.safetensors"
    assert asset.asset_type == "checkpoints"
    assert asset.node_id == "1"
    assert asset.node_type == "CheckpointLoaderSimple"
