import pytest

from aiatlas.connectors.hub.huggingface import HuggingFaceConnector, artifact_kind, canonical_name, pipeline_modalities
from aiatlas.sdk.facts import Target
from tests.conftest import claims_of, entity_names, extract_from_fixture, fixture_path


@pytest.fixture
def connector() -> HuggingFaceConnector:
    return HuggingFaceConnector({"models_per_org": 20})


async def test_discover_uses_registry_orgs(connector):
    targets = await connector.discover(None)  # type: ignore[arg-type]
    urls = {t.url for t in targets}
    assert "https://huggingface.co/papers" in urls
    assert "https://huggingface.co/models?author=Qwen&sort=downloads&p=0" in urls
    assert "https://huggingface.co/models?author=bartowski&sort=downloads&p=0" in urls
    assert len(targets) >= 40 and not any("/api/" in u for u in urls)


async def test_listing(connector):
    target = Target(url="https://huggingface.co/models?author=Qwen&sort=downloads&p=0", doc_type="listing", key="listing:Qwen:0", meta={"hf_org": "Qwen", "page": 0})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "listing-qwen.html"))
    models = [e for e in facts.entities if e.entity_type == "model"]
    artifacts = [e for e in facts.entities if e.entity_type == "artifact"]
    assert len(models) + len(artifacts) == 20 and len(artifacts) == 3            # models_per_org cap; Qwen's own FP8 repos are artifacts
    assert models[0].identifiers == {"hf_repo": "Qwen/Qwen3-0.6B"} and models[0].organization.name == "Qwen"
    assert models[0].family is not None and models[0].family.entity_type == "model_family" and models[0].family.name == "Qwen3"
    fp8 = next(a for a in artifacts if a.name == "Qwen/Qwen3.8-27B-FP8")
    assert fp8.artifact_kind == "quantization" and fp8.canonical is not None and fp8.canonical.name == "Qwen3.8-27B"
    assert fp8.canonical.identifiers == {"hf_repo": "Qwen/Qwen3.8-27B"}            # the official repo appears later in the same listing → one ref
    assert sum(1 for e in models if e.name == "Qwen3.8-27B") == 1
    c = claims_of(facts, "Qwen3-0.6B")
    assert c["metric.downloads"] == 20685071 and c["pipeline_tag"] == "text-generation" and c["access"] == "open" and c["weights_available"] is True
    assert c["modalities_input"] == ["text"] and c["modalities_output"] == ["text"] and "openness" not in c   # openness needs the licence (model page)
    assert c["parameter_count"] == 751632384
    assert [t.url for t in facts.targets][:2] == ["https://huggingface.co/Qwen/Qwen3-0.6B", "https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct"]
    assert all(t.doc_type == "model_page" for t in facts.targets) and len(facts.targets) == 20
    assert not any(r.predicate == "develops" and r.object.entity_type == "artifact" for r in facts.relations)


async def test_model_page(connector):
    target = Target(url="https://huggingface.co/Qwen/Qwen3-8B", doc_type="model_page", key="model:Qwen/Qwen3-8B", meta={"hf_repo": "Qwen/Qwen3-8B"})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "model-qwen3-8b.html"))
    c = claims_of(facts, "Qwen3-8B")
    assert c["license"] == "Apache-2.0" and c["license_raw"] == "apache-2.0" and c["library_name"] == "transformers" and c["architecture"] == "Qwen3ForCausalLM"
    assert c["parameter_count"] == 8190735360 and c["release_date"] == "2025-04-27" and c["openness"] == "open-weights"
    assert c["access"] == "open" and c["weights_available"] is True and "gated" not in c
    assert c["metric.likes"] == 1366 and c["file_size_gb"] == 16.38 and c["base_model"] == ["Qwen/Qwen3-8B-Base"]
    rels = {(r.subject.name, r.predicate, r.object.name) for r in facts.relations}
    assert ("Qwen3-8B", "fine_tuned_from", "Qwen3-8B-Base") in rels and ("Qwen", "develops", "Qwen3-8B") in rels
    assert ("Qwen3-8B", "described_by", "arXiv:2505.09388") in rels
    assert [t.url for t in facts.targets] == ["https://huggingface.co/Qwen/Qwen3-8B/raw/main/README.md"]
    assert facts.document_entity is not None and facts.document_entity.identifiers == {"hf_repo": "Qwen/Qwen3-8B"}
    assert facts.document_entity.entity_type == "model" and facts.document_entity.family.name == "Qwen3" and facts.document_entity.identity_confidence == "high"


async def test_model_page_gguf_and_gated(connector):
    target = Target(url="https://huggingface.co/bartowski/Qwen3.8-27B-GGUF", doc_type="model_page", key="x", meta={"hf_repo": "bartowski/Qwen3.8-27B-GGUF"})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "model-bartowski-gguf.html"))
    art = next(e for e in facts.entities if e.entity_type == "artifact")
    assert art.name == "bartowski/Qwen3.8-27B-GGUF" and art.artifact_kind == "quantization" and art.identifiers == {"hf_repo": "bartowski/Qwen3.8-27B-GGUF"}
    assert art.canonical is not None and art.canonical.entity_type == "model" and art.canonical.name == "Qwen3.8-27B"
    assert art.canonical.organization.name == "Qwen" and art.canonical.identity_confidence == "medium" and art.canonical.family.name == "Qwen3.8"
    c = claims_of(facts, "bartowski/Qwen3.8-27B-GGUF")                     # artifacts keep the full id as name
    assert c["quant_format"] == "gguf" and "Q4_K_M" in c["quantization"] and c["is_quantized"] is True and c["quantized_by"] == "bartowski"
    assert c["artifact_kind"] == "quantization" and c["parameter_count"] == 27_000_000_000     # the packaged size stays on the artifact…
    assert "parameter_count" not in claims_of(facts, "Qwen3.8-27B")                          # …never on the canonical model
    assert c["modalities_input"] == ["image", "text"] and c["modalities_output"] == ["text"]  # image-text-to-text
    assert facts.entities[0].name == "bartowski" and facts.entities[0].identifiers["registry_org"] == "bartowski"
    rels = {(r.subject.name, r.predicate, r.object.name) for r in facts.relations}
    assert ("bartowski/Qwen3.8-27B-GGUF", "published_by", "bartowski") in rels and not any(p == "develops" for _s, p, _o in rels)

    target = Target(url="https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct", doc_type="model_page", key="x", meta={"hf_repo": "meta-llama/Llama-3.1-8B-Instruct"})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "model-llama-gated.html"))
    c = claims_of(facts, "Llama-3.1-8B-Instruct")
    assert c["access"] == "gated" and c["gated_mode"] == "manual" and "restricted" not in str(c.get("status"))
    assert c["license"] == "Llama-3.1-Community" and c["license_raw"] == "llama3.1" and c["openness"] == "restricted-weights" and c["weights_available"] is True
    assert c["languages"] == ["en", "de", "fr", "it", "pt", "hi", "es", "th"]
    assert facts.targets == []                                 # gated card is not fetched (401 without auth)
    assert any(r.predicate == "fine_tuned_from" and r.attributes.get("base_model_relation") == "finetune" for r in facts.relations)
    llama = next(e for e in facts.entities if e.name == "Llama-3.1-8B-Instruct")
    assert llama.entity_type == "model" and llama.family.name == "Llama 3.1" and llama.family.organization.name == "Meta AI"


async def test_model_card(connector):
    target = Target(url="https://huggingface.co/bartowski/Qwen3.8-27B-GGUF/raw/main/README.md", doc_type="model_card", key="x", meta={"hf_repo": "bartowski/Qwen3.8-27B-GGUF"})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "readme-bartowski-gguf.md"), content_type="text/markdown")
    assert claims_of(facts, "bartowski/Qwen3.8-27B-GGUF")["base_model"] == ["Qwen/Qwen3.8-27B"]
    assert ("bartowski/Qwen3.8-27B-GGUF", "quantized_from", "Qwen3.8-27B") in {(r.subject.name, r.predicate, r.object.name) for r in facts.relations}
    art = next(e for e in facts.entities if e.entity_type == "artifact")
    assert art.canonical is not None and art.canonical.identifiers == {"hf_repo": "Qwen/Qwen3.8-27B"} and art.identity_confidence == "high"
    target = Target(url="https://huggingface.co/Qwen/Qwen3-8B/raw/main/README.md", doc_type="model_card", key="x", meta={"hf_repo": "Qwen/Qwen3-8B"})
    facts = await extract_from_fixture(connector, target, fixture_path("huggingface", "readme-qwen3-8b.md"), content_type="text/markdown")
    assert claims_of(facts, "Qwen3-8B")["license"] == "Apache-2.0"
    assert not any(r.predicate in ("derived_from", "quantized_from") for r in facts.relations)   # typed relation comes from the page


def test_identity_helpers():
    from aiatlas.connectors.hub.huggingface import _foreign_family

    assert _foreign_family("NousResearch", "Meta-Llama-3.1-70B-Instruct") is True and _foreign_family("meta-llama", "Llama-3.1-70B-Instruct") is False
    assert _foreign_family("NousResearch", "Hermes-4-70B") is False and _foreign_family("Qwen", "Qwen3-8B") is False
    assert _foreign_family("mlx-community", "Kimi-K2.5") is True
    assert canonical_name("Qwen3.8-27B-GGUF") == "Qwen3.8-27B" and canonical_name("Llama-3.1-8B-Instruct-bnb-4bit") == "Llama-3.1-8B-Instruct"
    assert canonical_name("DeepSeek-V3-0324-UD-Q4_K_XL") == "DeepSeek-V3-0324" and canonical_name("Kimi-K2.5-bf16") == "Kimi-K2.5"
    assert artifact_kind("zai-org/GLM-5-FP8", bases=[], quant_format=None) == "quantization"
    assert artifact_kind("mlx-community/Kimi-K2.5", bases=[], quant_format=None) == "packaging"
    assert artifact_kind("unsloth/Qwen3-8B-bf16", bases=[], quant_format=None) == "conversion"
    assert artifact_kind("Qwen/Qwen3-8B", bases=[("Qwen/Qwen3-8B-Base", "finetune")], quant_format=None) is None
    assert artifact_kind("NousResearch/Hermes-4-70B", bases=[("meta-llama/Llama-3.1-70B", "finetune")], quant_format=None) is None   # fine-tune = model
    assert pipeline_modalities("image-text-to-text") == (["image", "text"], ["text"]) and pipeline_modalities("automatic-speech-recognition") == (["audio"], ["text"])
    assert pipeline_modalities("text-to-image") == (["text"], ["image"]) and pipeline_modalities("unknown-tag") == ([], [])


def test_official_org_quantisations_are_artifacts(connector):
    """A quant/precision token makes the repo an artifact even under the developer's own org; canonical = the same-org model from the stripped
    name; None when the base cannot be named confidently."""
    from aiatlas.sdk.facts import Facts

    cases = {"nvidia/Gemma-4-31B-IT-NVFP4": ("quantization", "Gemma-4-31B-IT", "Google"),      # Gemma is Google's family: NVIDIA quantises it
             "tencent/HY-MT1.5-1.8B-FP8": ("quantization", "HY-MT1.5-1.8B", "Tencent"),
             "LiquidAI/LFM2.5-230M-GGUF": ("quantization", "LFM2.5-230M", "Liquid AI"), "black-forest-labs/FLUX.2-klein-4b-fp8": ("quantization", "FLUX.2-klein-4b", "Black Forest Labs"),
             "Qwen/Qwen3-8B-MLX-bf16": ("conversion", "Qwen3-8B", "Qwen")}
    for repo, (kind, base, org_name) in cases.items():
        facts = Facts()
        org = connector._org_ref(facts, repo.split("/")[0], None)
        ref = connector._model_ref(facts, repo, org)
        assert ref.entity_type == "artifact" and ref.artifact_kind == kind and ref.name == repo, repo
        assert ref.canonical is not None and ref.canonical.entity_type == "model" and ref.canonical.name == base, repo
        assert ref.canonical.organization.name == org_name and ref.canonical.identity_confidence == "medium", repo
        assert not any(r.predicate == "develops" and r.object is ref for r in facts.relations)
    # unsure → no invented model
    facts = Facts()
    ref = connector._model_ref(facts, "mradermacher/FooBar-GGUF", connector._org_ref(facts, "mradermacher", None))
    assert ref.entity_type == "artifact" and ref.canonical is None                        # redistributor, unknown family
    facts = Facts()
    ref = connector._model_ref(facts, "Qwen/Qwen3-8B", connector._org_ref(facts, "Qwen", None), quant_format="gguf")
    assert ref.entity_type == "artifact" and ref.canonical is None                        # quantised per tags only: nothing to strip from the name


async def test_daily_papers(connector):
    facts = await extract_from_fixture(connector, Target(url="https://huggingface.co/papers", doc_type="listing", key="papers"), fixture_path("huggingface", "papers.html"))
    papers = [e for e in facts.entities if e.entity_type == "paper"]
    assert len(papers) == 25 and all("arxiv" in p.identifiers for p in papers)
    first = claims_of(facts, papers[0].name)
    assert first["arxiv_id"] == "2609.10715" and first["metric.upvotes"] == 157 and first["pdf_url"] == "https://arxiv.org/pdf/2609.10715"
    assert len(first["authors"]) > 10 and first["published_at"].startswith("2026-09-09")
    assert entity_names(facts, "paper")