import pytest

from aiatlas.connectors.labs.xai import DOCS, XAIConnector, _display, _tok, _unit
from aiatlas.sdk.facts import Target
from tests.conftest import claims_of, entity_names, extract_from_fixture, fixture_path


@pytest.fixture
def connector() -> XAIConnector:
    return XAIConnector()


async def test_models_catalog(connector):
    facts = await extract_from_fixture(connector, Target(url=f"{DOCS}/models", doc_type="model_docs", key="models"), fixture_path("xai", "models.html"))
    models = entity_names(facts, "model")
    assert len(models) == 16 and {"Grok 4.6", "Grok 4.5", "Grok 4.3", "Grok Imagine Image 2.0", "Grok Imagine Video 1.5", "Grok TTS"} <= models
    g46 = claims_of(facts, "Grok 4.6")
    assert g46["api_model_id"] == "grok-4.6" and g46["context_length"] == 500_000 and g46["reasoning"] is True and g46["tool_calling"] is True
    assert g46["modalities_input"] == ["text", "image"] and g46["vision"] is True and g46["reasoning_effort_options"] == ["low", "medium", "high", "xhigh"]
    assert set(g46["regions"]) == {"us-east-1", "us-west-2", "us-central-1"}
    g45 = next(e for e in facts.entities if e.name == "Grok 4.5")
    assert g45.identifiers == {"xai_model_id": "grok-4.5"} and {"grok-4.5-latest", "grok-build-latest"} <= set(g45.aliases)
    by_id = {p.provider_model_id: p for p in facts.prices}
    p46 = by_id["grok-4.6"]
    assert (p46.input_per_mtok, p46.cached_input_per_mtok, p46.output_per_mtok) == (2.0, 0.5, 6.0) and p46.context_length == 500_000
    assert p46.features["long_context_threshold_tokens"] == 200_000 and p46.features["long_context_input_per_mtok"] == 4.0 and p46.features["long_context_output_per_mtok"] == 12.0
    assert p46.provider.identifiers["registry_provider"] == "xai"
    p43 = by_id["grok-4.3"]
    assert (p43.input_per_mtok, p43.output_per_mtok, p43.batch_input_per_mtok, p43.batch_output_per_mtok) == (1.25, 2.5, 1.0, 2.0)
    assert by_id["grok-imagine-image-2.0"].per_image == 0.04
    assert claims_of(facts, "Grok Imagine Video 1.5")["pricing_per_second_usd"] == {"480p": 0.08, "720p": 0.14, "1080p": 0.25}
    pages = [t for t in facts.targets if t.doc_type == "model_page"]
    assert len(pages) == 7 and pages[0].url == f"{DOCS}/models/grok-4.6.md"
    assert facts.document_entity is not None and facts.document_entity.entity_type == "provider"


async def test_model_page(connector):
    target = Target(url=f"{DOCS}/models/grok-4.6.md", doc_type="model_page", key="model:grok-4.6", meta={"api_id": "grok-4.6"})
    facts = await extract_from_fixture(connector, target, fixture_path("xai", "model-grok-4.6.md"), content_type="text/markdown")
    m = claims_of(facts, "Grok 4.6")
    assert m["context_length"] == 500_000 and m["batch_api"] is False and m["reasoning"] is True and m["structured_output"] is True
    assert m["modalities_input"] == ["text", "image"] and m["modalities_output"] == ["text"]
    assert m["description"].startswith("SpaceXAI's frontier model")
    assert next(e for e in facts.entities if e.entity_type == "model").identifiers == {"xai_model_id": "grok-4.6"}


def test_units_and_names():
    assert _tok("20000") == 2.0 and _tok("125") == 0.0125 and _tok(None) is None
    assert _unit("400000000") == 0.04 and _unit("800000000") == 0.08
    assert _display("grok-4.6") == "Grok 4.6" and _display("grok-imagine-image-2.0") == "Grok Imagine Image 2.0"
    assert _display("grok-4.20-0309-reasoning") == "grok-4.20-0309-reasoning" and _display("grok-tts") == "Grok TTS"