import pytest
import json
from py.recipes.parsers.comfy import ComfyMetadataParser


class LocalRecipeScanner:
    class LoraScanner:
        @staticmethod
        def has_hash(model_hash):
            return False

    def __init__(self, models):
        self.models = models
        self.queries = []
        self.base_models = []
        self._lora_scanner = self.LoraScanner()

    async def get_local_lora(self, name, base_model=None):
        self.queries.append(name)
        self.base_models.append(base_model)
        return self.models.get(name)


def local_lora(file_name):
    return {
        "file_path": f"/models/loras/{file_name}.safetensors",
        "file_name": file_name.rsplit("/", 1)[-1],
        "model_name": file_name.rsplit("/", 1)[-1],
        "sha256": file_name[0] * 64,
        "size": 4096,
        "base_model": "SDXL 1.0",
        "preview_url": "",
        "civitai": None,
    }


@pytest.mark.asyncio
async def test_parse_metadata_without_loras(monkeypatch):
    checkpoint_info = {
        "id": 2224012,
        "modelId": 1908679,
        "model": {"name": "SDXL Checkpoint", "type": "checkpoint"},
        "name": "v1.0",
        "images": [{"url": "https://image.civitai.com/checkpoints/original=true"}],
        "baseModel": "sdxl",
        "downloadUrl": "https://civitai.com/api/download/checkpoint",
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                assert version_id == "2224012"
                return checkpoint_info, None
        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = ComfyMetadataParser()

    # User provided metadata
    metadata_json = {
        "resource-stack": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": "urn:air:sdxl:checkpoint:civitai:1908679@2224012"}
        },
        "6": {
            "class_type": "smZ CLIPTextEncode",
            "inputs": {"text": "Positive prompt content"},
            "_meta": {"title": "Positive"}
        },
        "7": {
            "class_type": "smZ CLIPTextEncode",
            "inputs": {"text": "Negative prompt content"},
            "_meta": {"title": "Negative"}
        },
        "11": {
            "class_type": "KSampler",
            "inputs": {
                "sampler_name": "euler_ancestral",
                "scheduler": "normal",
                "seed": 904124997,
                "steps": 35,
                "cfg": 6,
                "denoise": 0.1,
                "model": ["resource-stack", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["21", 0]
            },
            "_meta": {"title": "KSampler"}
        },
        "extraMetadata": json.dumps({
            "prompt": "One woman, (solo:1.3), ...",
            "negativePrompt": "lowres, worst quality, ...",
            "steps": 35,
            "cfgScale": 6,
            "sampler": "euler_ancestral",
            "seed": 904124997,
            "width": 1024,
            "height": 1024
        })
    }

    result = await parser.parse_metadata(json.dumps(metadata_json))

    assert "error" not in result
    assert result["loras"] == []
    assert result["checkpoint"] is not None
    assert int(result["checkpoint"]["modelId"]) == 1908679
    assert int(result["checkpoint"]["id"]) == 2224012
    assert result["gen_params"]["prompt"] == "One woman, (solo:1.3), ..."
    assert result["gen_params"]["steps"] == 35
    assert result["gen_params"]["size"] == "1024x1024"
    assert result["from_comfy_metadata"] is True

@pytest.mark.asyncio
async def test_parse_metadata_resolves_standard_and_manager_local_loras(monkeypatch):
    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                raise AssertionError("Local LoRAs must not query Civitai")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({
        "styles/standard.safetensors": local_lora("standard"),
        "manager": local_lora("manager"),
    })
    metadata_json = {
        "1": {
            "class_type": "LoraLoader",
            "inputs": {
                "lora_name": "styles/standard.safetensors",
                "strength_model": 0.55,
            },
        },
        "2": {
            "class_type": "LoraLoaderLM",
            "inputs": {
                "loras": {
                    "__value__": [
                        {"name": "manager", "strength": "0.80", "active": True},
                        {"name": "disabled", "strength": 1.0, "active": False},
                        {"name": "dummy", "strength": 1.0, "active": True, "_isDummy": True},
                    ]
                }
            },
        },
    }

    result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)

    assert [entry["file_name"] for entry in result["loras"]] == ["standard", "manager"]
    assert [entry["weight"] for entry in result["loras"]] == [0.55, 0.8]
    assert all(isinstance(entry["weight"], float) for entry in result["loras"])
    assert all(entry["existsLocally"] is True for entry in result["loras"])
    assert all(entry["isDeleted"] is False for entry in result["loras"])
    assert scanner.queries == ["styles/standard.safetensors", "manager"]


@pytest.mark.asyncio
async def test_parse_metadata_defaults_malformed_weight_and_passes_checkpoint_base_model(monkeypatch):
    checkpoint_info = {
        "id": 456,
        "modelId": 123,
        "model": {"name": "Checkpoint", "type": "checkpoint"},
        "name": "v1",
        "baseModel": "SDXL 1.0",
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                return checkpoint_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"style": local_lora("style")})
    metadata_json = {
        "1": {"class_type": "LoraLoader", "inputs": {"lora_name": "style", "strength_model": "invalid"}},
        "2": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "civitai:123@456"}},
    }

    result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)

    assert result["loras"][0]["weight"] == 1.0
    assert scanner.base_models == ["SDXL 1.0"]


@pytest.mark.asyncio
async def test_parse_metadata_keeps_civitai_urn_with_local_lora(monkeypatch):
    remote_info = {
        "id": 456,
        "modelId": 123,
        "model": {"name": "Remote LoRA", "type": "LORA"},
        "name": "v1",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "name": "remote.safetensors",
                "hashes": {"SHA256": "c" * 64},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                assert version_id == "456"
                return remote_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local": local_lora("local")})
    metadata_json = {
        "1": {
            "class_type": "LoraLoader",
            "inputs": {"lora_name": "local", "strength_model": 0.4},
        },
        "2": {
            "class_type": "LoraLoader",
            "inputs": {
                "lora_name": "urn:air:sdxl:lora:civitai:123@456",
                "strength_model": 0.9,
            },
        },
    }

    result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)

    assert [entry["name"] for entry in result["loras"]] == ["local", "Remote LoRA"]
    assert [entry["weight"] for entry in result["loras"]] == [0.4, 0.9]
    assert result["loras"][0]["existsLocally"] is True
    assert result["loras"][1]["id"] == 456


@pytest.mark.asyncio
async def test_parse_metadata_without_extra_metadata(monkeypatch):
    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                return {"model": {"name": "Test"}, "id": version_id}, None
        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = ComfyMetadataParser()

    metadata_json = {
        "node_1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": "urn:air:sdxl:checkpoint:civitai:123@456"}
        }
    }

    result = await parser.parse_metadata(json.dumps(metadata_json))

    assert "error" not in result
    assert result["loras"] == []
    assert result["checkpoint"]["id"] == "456"


@pytest.mark.asyncio
async def test_parse_metadata_with_list_ckpt_name(monkeypatch):
    """ckpt_name serialized as a single-element list must not crash re.search."""
    checkpoint_info = {
        "id": 456,
        "modelId": 123,
        "model": {"name": "Checkpoint", "type": "checkpoint"},
        "name": "v1",
        "baseModel": "SDXL 1.0",
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                assert version_id == "456"
                return checkpoint_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )

    metadata_json = {
        "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": ["urn:air:sdxl:checkpoint:civitai:123@456"]},
        }
    }

    result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json))

    assert "error" not in result
    assert result["checkpoint"] is not None
    assert int(result["checkpoint"]["id"]) == 456
    assert int(result["checkpoint"]["modelId"]) == 123


@pytest.mark.asyncio
async def test_parse_metadata_with_none_ckpt_name(monkeypatch):
    """Missing (None) ckpt_name must not crash re.search with a TypeError."""

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                raise AssertionError("Checkpoint lookup must be skipped")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.comfy.get_default_metadata_provider",
        fake_metadata_provider,
    )

    metadata_json = {
        "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": None},
        }
    }

    result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json))

    assert "error" not in result
    assert result["checkpoint"] is None
    assert result["loras"] == []
