"""xAI — developer docs models page (Next.js RSC payload embeds `globalThis.__XAI_PUBLIC_MODELS__`) + per-model Markdown pages.

Sources (tier 1):
  * models   → docs.x.ai/developers/models (the /docs/models URL redirects here). The RSC payload carries a JSON catalog per region cluster:
               languageModels (name, aliases, modalities, maxPromptLength, prices in 1e-4 USD per 1M tokens, long-context threshold, features),
               imageGenerationModels / videoGenerationModels (per-image / per-second prices in 1e-10 USD), audioModels.
  * model pages → docs.x.ai/developers/models/<id>.md : "At a glance" (context window, model name, batch), capabilities, knowledge cut-off
  * x.ai/news answers 403 to crawlers → targeted with `escalate=True` only; without an escalation transport it is recorded as `blocked`.

Price units are verified against the human-readable table of models.md (grok-4.6: 20000 → $2.00 / 1M input tokens).
"""
from __future__ import annotations

import re
from typing import Any

from aiatlas.registry import org_ref, provider_ref
from aiatlas.sdk.connector import BaseConnector, Parsed, RunContext
from aiatlas.sdk.extract.dates import parse_datetime
from aiatlas.sdk.facts import EntityRef, Facts, Target
from aiatlas.sdk.fetch import FetchResult

from ._common import json_after, model_ref, next_flight_payload, tokens

DOCS = "https://docs.x.ai/developers"
NEWS = "https://x.ai/news"
PROVIDER_KEY = "xai"
ID_SCHEME = "xai_model_id"
TOKEN_PRICE_DIVISOR = 10_000          # "20000" → 2.0 USD per 1M tokens
UNIT_PRICE_DIVISOR = 10_000_000_000   # "400000000" → 0.04 USD per image / per second
MAX_MODEL_PAGES = 20
MODALITY = {"TEXT": "text", "IMAGE": "image", "AUDIO": "audio", "VIDEO": "video"}


class XAIConnector(BaseConnector):
    name = "xai"
    label = "xAI — Grok models & pricing"
    description = "xAI developer docs: public model catalog embedded in the models page (ids, aliases, modalities, context, prices, features) and model pages."
    source_key = "x.ai"
    version = "1"
    parser_version = "1"
    interval_seconds = 7200
    min_interval_seconds = 3600
    rate_per_min = 10
    tier = 1
    priority = 1
    expected_min_records = 10
    concurrency = 2

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [
            Target(url=f"{DOCS}/models", doc_type="model_docs", key="models", min_bytes=20000),
            Target(url=NEWS, doc_type="listing", key="news", min_bytes=2000, escalate=True, priority=3),
        ]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        org = org_ref("xai")
        facts.entities.append(org)
        key = target.key or ""
        if key == "models" and parsed.html:
            self._models(facts, org, res)
        elif target.doc_type == "model_page" and parsed.markdown:
            self._model_page(facts, org, target, parsed)
        elif key == "news" and parsed.html:
            from aiatlas.sdk.extract.feeds import FeedItem

            from ._common import announcement_events

            items = [FeedItem(id=h, url=h, title=t, summary=None, published_at=None, updated_at=None)
                     for h, t in parsed.html.links_matching(r"^https://x\.ai/news/[^/?#]+$") if len(t) > 12]
            announcement_events(facts, org, items, source_name="x.ai/news", max_follow=10)
            facts.document_title, facts.document_entity = "xAI news", org
        return facts

    # ------------------------------------------------------------------------------------------ models catalog (embedded JSON)
    def _models(self, facts: Facts, org: EntityRef, res: FetchResult) -> None:
        payload = next_flight_payload(res.content)
        data = json_after(payload, "__XAI_PUBLIC_MODELS__=") if payload else None
        if not isinstance(data, dict):
            data = json_after(res.text.replace('\\"', '"'), "__XAI_PUBLIC_MODELS__=")
        if not isinstance(data, dict):
            return
        provider = provider_ref(PROVIDER_KEY)
        facts.entities.append(provider)
        facts.document_title, facts.document_entity = "xAI models", provider
        merged: dict[str, dict[str, Any]] = {}
        for cluster in data.get("clusterConfigs") or []:
            region = cluster.get("clusterName")
            for kind in ("languageModels", "imageGenerationModels", "videoGenerationModels", "audioModels"):
                for m in cluster.get(kind) or []:
                    name_ = m.get("name")
                    if not name_:
                        continue
                    entry = merged.setdefault(name_, {"kind": kind, "model": m, "regions": []})
                    if region and region not in entry["regions"]:
                        entry["regions"].append(region)
                    if len(m) > len(entry["model"]):
                        entry["model"] = m
        followed = 0
        for name_, entry in merged.items():
            m, kind = entry["model"], entry["kind"]
            ref = model_ref(facts, _display(name_), org, api_id=name_, provider_key=PROVIDER_KEY, family="Grok" if name_.startswith("grok") else None,
                            aliases=[name_] + [a for a in (m.get("aliases") or []) if isinstance(a, str)])
            facts.claim(ref, "api_model_id", name_)
            facts.claim(ref, "official_url", f"{DOCS}/models/{name_}")
            facts.claim(ref, "regions", sorted(entry["regions"]))
            mods_in = [MODALITY[x] for x in m.get("inputModalities") or [] if x in MODALITY]
            mods_out = [MODALITY[x] for x in m.get("outputModalities") or [] if x in MODALITY]
            facts.claim(ref, "modalities_input", mods_in or None)
            facts.claim(ref, "modalities_output", mods_out or None)
            facts.claim(ref, "modalities", sorted(set(mods_in) | set(mods_out)) or None)
            if mods_in:
                facts.claim(ref, "vision", "image" in mods_in)
            feats = m.get("features") or {}
            if kind == "languageModels":
                facts.claim(ref, "context_length", m.get("maxPromptLength") if isinstance(m.get("maxPromptLength"), int) else None, unit="tokens")
                if "functionCalling" in feats:
                    facts.claim(ref, "tool_calling", bool(feats["functionCalling"]))
                if "structuredOutputs" in feats:
                    facts.claim(ref, "structured_output", bool(feats["structuredOutputs"]))
                if "reasoning" in feats:
                    facts.claim(ref, "reasoning", bool(feats["reasoning"]))
                efforts = (feats.get("reasoningEffortOptions") or {}).get("supportedEfforts")
                facts.claim(ref, "reasoning_effort_options", efforts if isinstance(efforts, list) else None)
                self._language_price(facts, ref, provider, name_, m, entry["regions"])
                if followed < MAX_MODEL_PAGES:
                    followed += 1
                    facts.follow(f"{DOCS}/models/{name_}.md", doc_type="model_page", key=f"model:{name_}", meta={"api_id": name_, "content_type": "text/markdown"},
                                 min_bytes=200, priority=1)
            elif kind == "imageGenerationModels":
                per_image = [(_unit(r.get("pricePerImage")), r.get("resolution"), r.get("quality")) for r in m.get("resolutionPricing") or []]
                per_image = [p for p in per_image if p[0] is not None]
                base = min((p[0] for p in per_image), default=_unit(m.get("imagePrice")))
                if base is not None:
                    facts.price(model=ref, provider=provider, provider_model_id=name_, per_image=base,
                                features={"per_image_by_resolution": [{"resolution": r, "quality": q, "usd": p} for p, r, q in per_image],
                                          "per_input_image": _unit(m.get("pricePerInputImage")), "regions": sorted(entry["regions"])}, meta={"from": "models page JSON"})
            elif kind == "videoGenerationModels":
                per_sec = {str(r.get("resolution")).replace("VIDEO_RESOLUTION_", "").lower(): _unit(r.get("pricePerSecond")) for r in m.get("resolutionPricing") or []}
                facts.claim(ref, "pricing_per_second_usd", {k: v for k, v in per_sec.items() if v is not None} or None)
            if m.get("version") and str(m["version"]) != "1.0":
                facts.claim(ref, "version", str(m["version"]))

    def _language_price(self, facts: Facts, ref: EntityRef, provider: EntityRef, name_: str, m: dict[str, Any], regions: list[str]) -> None:
        p_in, p_cached, p_out = _tok(m.get("promptTextTokenPrice")), _tok(m.get("cachedPromptTokenPrice")), _tok(m.get("completionTextTokenPrice"))
        if p_in is None and p_out is None:
            return
        features: dict[str, Any] = {"regions": sorted(regions)}
        threshold = m.get("longContextThreshold")
        if threshold and str(threshold).isdigit():
            features["long_context_threshold_tokens"] = int(threshold)
            for key, prop in (("promptTextTokenPriceLongContext", "long_context_input_per_mtok"), ("cachedPromptTokenPriceLongContext", "long_context_cached_input_per_mtok"),
                              ("completionTokenPriceLongContext", "long_context_output_per_mtok")):
                v = _tok(m.get(key))
                if v is not None:
                    features[prop] = v
        img = _tok(m.get("promptImageTokenPrice"))
        if img is not None and img != p_in:
            features["image_input_per_mtok"] = img
        if m.get("batchDiscountPercent"):
            features["batch_discount_percent"] = m["batchDiscountPercent"]
            if p_in is not None:
                pass
        features["batch_enabled"] = bool(m.get("batchEnabled"))
        batch_in = batch_out = None
        if m.get("batchEnabled") and m.get("batchDiscountPercent"):
            factor = 1 - float(m["batchDiscountPercent"]) / 100
            batch_in = round(p_in * factor, 6) if p_in is not None else None
            batch_out = round(p_out * factor, 6) if p_out is not None else None
        facts.price(model=ref, provider=provider, provider_model_id=name_, input_per_mtok=p_in, cached_input_per_mtok=p_cached, output_per_mtok=p_out,
                    batch_input_per_mtok=batch_in, batch_output_per_mtok=batch_out, context_length=m.get("maxPromptLength") if isinstance(m.get("maxPromptLength"), int) else None,
                    features=features, meta={"from": "models page JSON"})

    # ------------------------------------------------------------------------------------------ model page (.md)
    def _model_page(self, facts: Facts, org: EntityRef, target: Target, parsed: Parsed) -> None:
        md = parsed.markdown
        assert md
        api_id = target.meta.get("api_id")
        m = re.search(r"\*\*Model name:\*\*\s*`([^`]+)`", md.body)
        if m:
            api_id = m.group(1).strip()
        title = next((t for lvl, t in md.headings if lvl == 1), None)
        if not api_id or not title:
            return
        ref = model_ref(facts, title, org, api_id=api_id, provider_key=PROVIDER_KEY, aliases=[api_id])
        facts.document_title, facts.document_entity = title, ref
        first_para = next((ln.strip() for ln in md.body.split("\n")[1:] if ln.strip() and not ln.startswith("#")), None)
        facts.claim(ref, "description", first_para)
        m = re.search(r"\*\*Context window:\*\*\s*([\d,]+)", md.body)
        if m:
            facts.claim(ref, "context_length", tokens(m.group(1)), unit="tokens")
        m = re.search(r"\*\*Modalities:\*\*\s*([^→\n]+)→\s*([^\n]+)", md.body)
        if m:
            mods_in = [x.strip() for x in m.group(1).split(",") if x.strip() in MODALITY.values()]
            mods_out = [x.strip() for x in m.group(2).split(",") if x.strip() in MODALITY.values()]
            facts.claim(ref, "modalities_input", mods_in or None)
            facts.claim(ref, "modalities_output", mods_out or None)
        for label, prop in (("Function calling", "tool_calling"), ("Structured outputs", "structured_output"), ("Reasoning", "reasoning")):
            mm = re.search(rf"\*\*{label}:\*\*\s*(Yes|No)", md.body)
            if mm:
                facts.claim(ref, prop, mm.group(1) == "Yes")
        mm = re.search(r"knowledge cut-?off date of .*? is ([A-Z][a-z]+ \d{1,2}, \d{4})", md.body)
        if mm:
            dt = parse_datetime(mm.group(1))
            facts.claim(ref, "knowledge_cutoff", f"{dt.year:04d}-{dt.month:02d}" if dt else None)
        mm = re.search(r"\*\*Batch API:\*\*\s*(Supported|Not supported)", md.body)
        if mm:
            facts.claim(ref, "batch_api", mm.group(1) == "Supported")


# ---------------------------------------------------------------------------------------------- helpers
def _tok(v: Any) -> float | None:
    try:
        return round(int(str(v)) / TOKEN_PRICE_DIVISOR, 6) if v not in (None, "") else None
    except ValueError:
        return None


def _unit(v: Any) -> float | None:
    try:
        return round(int(str(v)) / UNIT_PRICE_DIVISOR, 6) if v not in (None, "") else None
    except ValueError:
        return None


def _display(api_id: str) -> str:
    """grok-4.6 → Grok 4.6 ; grok-imagine-image-2.0 → Grok Imagine Image 2.0 ; dated / variant ids are kept verbatim."""
    parts = api_id.split("-")
    if all(re.fullmatch(r"[a-z]+|\d+(?:\.\d+)?", p) for p in parts) and sum(bool(re.fullmatch(r"\d+(?:\.\d+)?", p)) for p in parts) <= 1:
        return " ".join((p.upper() if p in ("tts", "stt") else p.capitalize()) if p.isalpha() else p for p in parts)
    return api_id


CONNECTORS = [XAIConnector]