from aiatlas.connectors.hardware.spec_pages import AppleSpecsConnector, NvidiaSpecsConnector
from aiatlas.sdk.facts import Target
from tests.conftest import claims_of, extract_from_fixture, fixture_path


def _hw(facts, name):
    return next(e for e in facts.entities if e.entity_type == "hardware" and e.name == name)


async def test_apple_mac_mini_and_studio():
    c = AppleSpecsConnector()
    facts = await extract_from_fixture(c, Target(url="https://www.apple.com/mac-mini/specs/", doc_type="spec_page", key="apple:mac-mini", meta={"product": "Mac mini"}),
                                       fixture_path("hardware_specs", "apple-mac-mini.html"))
    m6 = _hw(facts, "Apple M6")
    assert m6.identifiers == {"registry_hardware": "apple-m6"}
    assert claims_of(facts, "Apple M6")["memory_bandwidth_gbs"] == 170.0 and claims_of(facts, "Apple M6")["memory_bandwidth_options_gbs"] == [153.0, 170.0]
    mini = claims_of(facts, "Mac mini (Apple M6)")
    assert mini["memory_gb"] == [16, 24, 32] and mini["memory_configurations"] == [{"memory_gb": 24, "memory_bandwidth_gbs": 170.0}, {"memory_gb": 32, "memory_bandwidth_gbs": 170.0}]
    assert claims_of(facts, "Mac mini (Apple M5 Pro)")["memory_gb"] == [24, 48, 64] and claims_of(facts, "Apple M5 Pro")["memory_bandwidth_gbs"] == 307.0
    assert ("Mac mini (Apple M6)", "uses", "Apple M6") in {(r.subject.name, r.predicate, r.object.name) for r in facts.relations}

    facts = await extract_from_fixture(c, Target(url="https://www.apple.com/mac-studio/specs/", doc_type="spec_page", key="apple:mac-studio", meta={"product": "Mac Studio"}),
                                       fixture_path("hardware_specs", "apple-mac-studio.html"))
    assert claims_of(facts, "Apple M5 Max")["memory_bandwidth_options_gbs"] == [460.0, 614.0] and claims_of(facts, "Apple M5 Ultra")["memory_bandwidth_gbs"] == 1200.0
    assert claims_of(facts, "Mac Studio (Apple M5 Max)")["memory_gb"] == [36, 48, 64, 128] and claims_of(facts, "Mac Studio (Apple M5 Ultra)")["memory_gb"] == [96, 256, 512]


async def test_apple_other_templates():
    c = AppleSpecsConnector()
    facts = await extract_from_fixture(c, Target(url="https://www.apple.com/macbook-pro/specs/", doc_type="spec_page", key="apple:macbook-pro", meta={"product": "MacBook Pro"}),
                                       fixture_path("hardware_specs", "apple-macbook-pro.html"))
    assert claims_of(facts, "MacBook Pro (Apple M5 Pro)")["memory_gb"] == [24, 48, 64]              # 36 / 128 GB require the M5 Max
    assert claims_of(facts, "MacBook Pro (Apple M5 Max)")["memory_gb"] == [36, 48, 64, 128]
    assert claims_of(facts, "Apple M5 Pro")["memory_bandwidth_gbs"] == 307.0 and "memory_bandwidth_options_gbs" not in claims_of(facts, "Apple M5 Pro")
    facts = await extract_from_fixture(c, Target(url="https://www.apple.com/macbook-air/specs/", doc_type="spec_page", key="apple:macbook-air", meta={"product": "MacBook Air"}),
                                       fixture_path("hardware_specs", "apple-macbook-air.html"))
    assert claims_of(facts, "MacBook Air (Apple M5)")["memory_gb"] == [16, 24, 32] and claims_of(facts, "Apple M5")["memory_bandwidth_gbs"] == 153.0
    facts = await extract_from_fixture(c, Target(url="https://www.apple.com/imac/specs/", doc_type="spec_page", key="apple:imac", meta={"product": "iMac"}),
                                       fixture_path("hardware_specs", "apple-imac.html"))
    assert _hw(facts, "Apple M4").identifiers == {"registry_hardware": "apple-m4"} and claims_of(facts, "Apple M4")["memory_bandwidth_gbs"] == 120.0
    assert claims_of(facts, "iMac (Apple M4)")["memory_gb"] == [16, 24, 32]


async def test_nvidia_pages():
    c = NvidiaSpecsConnector()
    targets = await c.discover(None)  # type: ignore[arg-type]
    by_key = {t.key: t for t in targets}
    facts = await extract_from_fixture(c, by_key["nvidia:h100"], fixture_path("hardware_specs", "nvidia-h100.html"))
    sxm = claims_of(facts, "NVIDIA H100 SXM")
    assert (sxm["memory_gb"], sxm["memory_bandwidth_gbs"], sxm["tdp_watts"]) == (80, 3350.0, 700)          # matches registry/hardware.yaml
    assert sxm["compute_fp16_tflops"] == 1979 and sxm["compute_fp8_tflops"] == 3958 and sxm["compute_sparsity"] is True and sxm["nvlink_bandwidth_gbs"] == 900
    nvl = claims_of(facts, "NVIDIA H100 NVL")
    assert (nvl["memory_gb"], nvl["memory_bandwidth_gbs"], nvl["tdp_watts"]) == (94, 3900.0, 400) and _hw(facts, "NVIDIA H100 NVL").identifiers == {"registry_hardware": "nvidia-h100-nvl"}
    facts = await extract_from_fixture(c, by_key["nvidia:h200"], fixture_path("hardware_specs", "nvidia-h200.html"))
    h200 = claims_of(facts, "NVIDIA H200")
    assert (h200["memory_gb"], h200["memory_bandwidth_gbs"], h200["tdp_watts"]) == (141, 4800.0, 700) and claims_of(facts, "NVIDIA H200 NVL")["tdp_watts"] == 600
    facts = await extract_from_fixture(c, by_key["nvidia:dgx-b200"], fixture_path("hardware_specs", "nvidia-dgx-b200.html"))
    dgx = claims_of(facts, "NVIDIA DGX B200")
    assert (dgx["gpu_count"], dgx["memory_gb"], dgx["memory_type"], dgx["memory_bandwidth_gbs"]) == (8, 1440, "HBM3e", 64000.0)
    assert dgx["compute_fp8_tflops"] == 72000 and dgx["system_power_watts"] == 14300
    assert ("NVIDIA DGX B200", "uses", "NVIDIA B200") in {(r.subject.name, r.predicate, r.object.name) for r in facts.relations}
    facts = await extract_from_fixture(c, by_key["nvidia:dgx-spark"], fixture_path("hardware_specs", "nvidia-dgx-spark.html"))
    spark = claims_of(facts, "NVIDIA DGX Spark")
    assert (spark["memory_gb"], spark["memory_type"], spark["memory_bandwidth_gbs"], spark["tdp_watts"], spark["power_supply_watts"]) == (128, "LPDDR5x", 273.0, 140, 240)
    assert spark["compute_fp4_tflops"] == 1000 and spark["architecture"] == "NVIDIA Grace Blackwell"