"""Mistral AI — docs models overview (Next.js, server-rendered cards + deprecation table), model pages (embedded pricing JSON), news RSS.

Sources (tier 1):
  * models overview → docs.mistral.ai/models (the old /getting-started/models/models_overview/ URL redirects here):
                      cards `a[href^=/models/]` (name, weights-license badge, description, version) and the
                      "Deprecated & retired models" table (name, version, API id, deprecation / retirement dates, alternative)
  * model pages     → docs.mistral.ai/models/<slug> : context window, "Released as open weights under a … license", features,
                      RSC payload `"pricing":{input:[{price}], output:[{price}]}` in USD per 1M tokens (EUR kept in features), `isRetired`
  * news            → mistral.ai/news/rss (RSS feed linked from the news listing) → ANNOUNCEMENT events

mistral.ai/pricing is server-rendered without any per-model price (plans only; model prices load client-side), so it is not crawled:
provider prices come from the docs model pages instead.
"""
from __future__ import annotations

import re
from datetime import datetime
from typing import Any

from selectolax.parser import HTMLParser

from aiatlas.connectors._identity import family_ref
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.extract.html import node_text
from aiatlas.sdk.facts import EntityRef, Facts, Target
from aiatlas.sdk.fetch import FetchResult

from ._common import (
    announcement_events,
    claim_api_aliases,
    claim_license,
    claim_status,
    json_after,
    model_ref,
    next_flight_payload,
    normalize_capabilities,
    slug_of,
    tokens,
)

DOCS = "https://docs.mistral.ai"
NEWS_RSS = "https://mistral.ai/news/rss"
PROVIDER_KEY = "mistral"
SLUG_SCHEME = "mistral_docs_slug"
ID_SCHEME = "mistral_model_id"
MAX_MODEL_PAGES = 50

API_ID = re.compile(r"^[a-z0-9][a-z0-9.\-]*$")
OPEN_LICENSES = re.compile(r"apache|mit\b|mrl|research licen[cs]e|open[- ]weight|cc-by|gpl|bsd", re.IGNORECASE)
LICENSE_SENTENCE = re.compile(r"(?:released|available|distributed)\s+(?:as\s+)?open[- ]weights?\s+under\s+(?:a\s+|an\s+|the\s+)?(.+?)\s+licen[cs]e", re.IGNORECASE)
US_DATE = re.compile(r"\b(\d{1,2})/(\d{1,2})/(\d{4})\b")


class MistralConnector(BaseConnector):
    name = "mistral"
    label = "Mistral AI — models, model pages (pricing), news"
    description = "Mistral docs models overview (cards + deprecation table), model pages with embedded pricing, and the Mistral news RSS."
    source_key = "docs.mistral.ai"
    version = "1"
    parser_version = "1"
    interval_seconds = 3600
    min_interval_seconds = 1800
    rate_per_min = 12
    tier = 1
    priority = 0
    expected_min_records = 40
    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_RSS, doc_type="feed", key="news", min_bytes=1000),
        ]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        org = org_ref("mistral")
        facts.entities.append(org)
        key = target.key or ""
        if key == "models" and parsed.html:
            self._models(facts, org, parsed, res)
        elif key == "news" and parsed.kind == "feed":
            announcement_events(facts, org, parsed.feed_items, source_name="mistral.ai/news", max_follow=15)
            facts.document_title, facts.document_entity = "Mistral AI news", org
        elif target.doc_type == "model_page" and parsed.html:
            self._model_page(facts, org, target, res, parsed)
        return facts

    # ------------------------------------------------------------------------------------------ models overview
    def _models(self, facts: Facts, org: EntityRef, parsed: Parsed, res: FetchResult) -> None:
        observed: datetime = res.fetched_at
        html = parsed.html
        assert html
        facts.document_title, facts.document_entity = "Mistral models overview", org
        # the SDK's parse_html strips <button> (and nav/footer/form) from the tree it keeps; the licence badge lives inside a button,
        # so the cards are read from a fresh parse of the raw document
        tree = HTMLParser(res.content)
        active: list[str] = []
        names: dict[str, str] = {}            # display name (lower) -> slug of the first model carrying it (Mistral reuses names across versions)
        for card in tree.css('a[href^="/models/"]'):
            slug = slug_of(card.attributes.get("href") or "")
            h3 = card.css_first("h3")
            if not h3 or not slug:
                continue
            version = next((m.group(1) for m in (re.fullmatch(r"v\s*(\d[\w.\-]*)", node_text(div)) for div in card.css("div")) if m), None)
            name = _unique_name(names, node_text(h3), slug, version)
            desc = card.css_first("p")
            desc_text = node_text(desc) if desc else ""
            developer = _third_party(name, desc_text)
            ref = _model(facts, developer or org, name, slug, third_party=developer is not None)
            badge = card.css_first("span[data-slot='badge']")
            if badge:
                _license_claims(facts, ref, node_text(badge))
            facts.claim(ref, "description", desc_text or None)
            facts.claim(ref, "version", version)
            if slug not in active:
                active.append(slug)
        deprecated: list[str] = []
        for tr in tree.css("table tr"):
            tds = tr.css("td")
            if len(tds) != 5:
                continue
            link = tds[0].css_first("a[href]")
            name = re.sub(r"\s*[↗→]\s*$", "", node_text(tds[0])).strip()
            slug = slug_of(link.attributes.get("href") or "") if link else ""
            api_id = node_text(tds[2]).strip()
            if not name:
                continue
            name = _unique_name(names, name, slug, node_text(tds[1]) or None)
            ref = _model(facts, org, name, slug or None, api_id if API_ID.match(api_id) else None)
            facts.claim(ref, "version", node_text(tds[1]) or None)
            dates = [_us_date(m) for m in US_DATE.finditer(node_text(tds[3]))]
            dep, ret = (dates + [None, None])[:2]
            facts.claim(ref, "deprecation_date", dep)
            facts.claim(ref, "retirement_date", ret)
            claim_status(facts, ref, "retired" if ret and ret < observed.date().isoformat() else "deprecated")
            alt_link = tds[4].css_first("a[href]")
            alt_name = node_text(tds[4]).strip()
            if alt_name and alt_name != name:
                alt = _model(facts, org, alt_name, slug_of(alt_link.attributes.get("href") or "") if alt_link else None)
                facts.relate(ref, "superseded_by", alt, attributes={"deprecated": dep, "retired": ret})
            if slug:
                deprecated.append(slug)
        for slug in [s for s in active if s not in deprecated][:MAX_MODEL_PAGES]:
            facts.follow(f"{DOCS}/models/{slug}", doc_type="model_page", key=f"model:{slug}", meta={"slug": slug}, min_bytes=20000, priority=1)

    # ------------------------------------------------------------------------------------------ model page
    def _model_page(self, facts: Facts, org: EntityRef, target: Target, res: FetchResult, parsed: Parsed) -> None:
        html = parsed.html
        assert html
        slug = target.meta.get("slug") or slug_of(res.final_url or res.url)
        name = next((t for lvl, t in html.headings if lvl == 1), None)
        if not name or not slug:
            return
        ref = _model(facts, org, name, slug)
        facts.document_title, facts.document_entity = name, ref
        facts.claim(ref, "official_url", f"{DOCS}/models/{slug}")
        text = html.text
        m = re.search(r"\bContext\s+(\d+(?:\.\d+)?\s*[kKmM]\b|\d{4,})", text)
        if m:
            facts.claim(ref, "context_length", tokens(m.group(1)), unit="tokens")
        intro = next((node_text(p) for p in html.css("p") if LICENSE_SENTENCE.search(node_text(p))), None)
        m = LICENSE_SENTENCE.search(intro or "")
        if m:
            facts.claim(ref, "description", intro.split(" Released")[0].strip() if intro and " Released" in intro else None)
            claim_license(facts, ref, m.group(1).strip(), weights_available=True)
        features_section = html.css_first("h3#features, h3[id*='features']")
        feats = [node_text(n) for n in html.css("[class*='LinkItem_link'] span, [class*='LinkItem_link'] p")] if features_section is None else []
        blob = " | ".join(feats) if feats else text
        caps: list[str] = []
        for label, prop in (("Function Calling", "tool_calling"), ("Structured Outputs", "structured_output"), ("Vision", "vision"), ("Reasoning", "reasoning")):
            if re.search(rf"\b{label}\b", blob):
                facts.claim(ref, prop, True)
                caps.append(label)
        if caps:
            facts.claim(ref, "capabilities", normalize_capabilities(caps))
        payload = next_flight_payload(res.content)
        if payload:
            pricing = json_after(payload, '"pricing":')
            if isinstance(pricing, dict):
                provider = provider_ref(PROVIDER_KEY)
                facts.entities.append(provider)
                p_in = _first_price(pricing.get("input"))
                p_out = _first_price(pricing.get("output"))
                if p_in is not None or p_out is not None:
                    feats_d: dict[str, Any] = {"free": bool(pricing.get("free"))}
                    eur_in, eur_out = _first_price(pricing.get("input"), key="priceEur"), _first_price(pricing.get("output"), key="priceEur")
                    if eur_in is not None:
                        feats_d["eur_input_per_mtok"] = eur_in
                    if eur_out is not None:
                        feats_d["eur_output_per_mtok"] = eur_out
                    facts.price(model=ref, provider=provider, input_per_mtok=p_in, output_per_mtok=p_out, features=feats_d, meta={"from": "model page"})
            if re.search(r'"isRetired":true', payload):
                claim_status(facts, ref, "retired")
            names = json_after(payload, '"names":')
            if isinstance(names, list):
                latest = [n for n in names if isinstance(n, str) and n.endswith("-latest")]
                claim_api_aliases(facts, ref, latest)


# ---------------------------------------------------------------------------------------------- helpers
def _model(facts: Facts, org: EntityRef, name: str, slug: str | None, api_id: str | None = None, *, third_party: bool = False) -> EntityRef:
    for e in facts.entities:
        if e.entity_type != "model":
            continue
        if slug and e.identifiers.get(SLUG_SCHEME) == slug:
            if api_id and ID_SCHEME not in e.identifiers:
                e.identifiers[ID_SCHEME] = api_id
                facts.claim(e, "api_model_id", api_id)
            return e
        if not slug and e.name.lower() == name.lower():
            return e
    ids: dict[str, str] = {}
    if slug:
        ids[SLUG_SCHEME] = slug
    if api_id:
        ids[ID_SCHEME] = api_id
    ref = facts.entity("model", name, identifiers=ids, organization=org, aliases=[api_id] if api_id else [], family=family_ref(name, org),
                       identity_confidence="high" if ids else "medium")
    if org not in facts.entities:
        facts.entities.append(org)
    facts.relate(org, "develops", ref)
    if third_party:
        facts.relate(ref, "available_through", provider_ref(PROVIDER_KEY))
    facts.claim(ref, "api_model_id", api_id)
    if slug:
        facts.claim(ref, "official_url", f"{DOCS}/models/{slug}")
    return ref


def _unique_name(names: dict[str, str], name: str, slug: str, version: str | None) -> str:
    """'Codestral' exists as 24.05, 25.01 and 25.08: the first keeps the bare name, later versions get it appended so that
    alias resolution never folds distinct versions into one entity."""
    key = name.lower()
    if slug and names.setdefault(key, slug) != slug:
        return f"{name} {version}" if version and not name.endswith(version) else f"{name} ({slug})"
    return name


THIRD_PARTY_ORGS = {"z.ai": "zhipu", "glm": "zhipu"}


def _third_party(name: str, description: str) -> EntityRef | None:
    """Models hosted by Mistral but developed elsewhere ("A third-party open source text model from Z.ai") → their own registry org."""
    if "third-party" not in description.lower() and "third party" not in description.lower():
        return None
    low = name.lower()
    for prefix, key in THIRD_PARTY_ORGS.items():
        if low.startswith(prefix):
            return org_ref(key)
    return None


def _license_claims(facts: Facts, ref: EntityRef, badge: str) -> None:
    """Weights-licence badge ("Apache 2.0", "MRL", "Modified MIT", "Proprietary") → canonical `license` (+ `license_raw`) and the openness
    category derived from the ontology (MRL → restricted-weights, Apache 2.0 → open-weights)."""
    badge = badge.strip()
    if not badge:
        return
    low = badge.lower()
    if OPEN_LICENSES.search(badge):
        claim_license(facts, ref, badge, weights_available=True)
    elif "proprietary" in low or "commercial" in low or "api only" in low or "closed" in low:
        facts.claim(ref, "openness", "proprietary")
        facts.claim(ref, "weights_available", False)
        if "licen" in low:
            claim_license(facts, ref, badge)


def _us_date(m: re.Match[str]) -> str | None:
    dt = parse_datetime(f"{m.group(3)}-{int(m.group(1)):02d}-{int(m.group(2)):02d}")
    return dt.date().isoformat() if dt else None


def _first_price(entries: Any, *, key: str = "price") -> float | None:
    if not isinstance(entries, list):
        return None
    for e in entries:
        if isinstance(e, dict) and isinstance(e.get(key), (int, float)) and "token" in str(e.get("denominator", "/M Tokens")).lower():
            return float(e[key])
    return None


_ = model_ref  # shared helper kept importable for symmetry with the other lab connectors

CONNECTORS = [MistralConnector]