"""Google — Gemini API docs (models, pricing, release notes) + Google DeepMind and Google AI blog feeds.

Sources (tier 1):
  * models         → ai.google.dev/gemini-api/docs/models : model cards (name, endpoint id, status, description) + "Model | Endpoint" tables;
                     per-model pages followed (Property/Description table: model code, data types, token limits, capabilities, versions, cutoff)
  * pricing        → …/docs/pricing : one section per model (h2 + endpoint id) with Standard / Batch / Flex / Priority tables
                     (paid tier, USD per 1M tokens; long-context tiers and scheduled changes kept in `features`)
  * release notes  → …/docs/changelog : dated sections → ANNOUNCEMENT events
  * DeepMind blog  → deepmind.google/blog/rss.xml (organization google-deepmind)
  * Google AI blog → blog.google/innovation-and-ai/technology/ai/rss/ (organization google)
"""
from __future__ import annotations

import re
from html import unescape as html_unescape
from typing import Any

from selectolax.parser import Node

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 (
    MODEL_WORDS,
    RELEASE_WORDS,
    announcement_events,
    claim_api_aliases,
    claim_modalities,
    claim_status,
    clean_cell,
    model_ref,
    money,
    month_year,
    normalize_capabilities,
    slug_of,
    tokens,
)

DOCS = "https://ai.google.dev/gemini-api/docs"
DEEPMIND_RSS = "https://deepmind.google/blog/rss.xml"
GOOGLE_AI_RSS = "https://blog.google/innovation-and-ai/technology/ai/rss/"
PROVIDER_KEY = "google-gemini-api"
ID_SCHEME = "gemini_model_id"
MAX_MODEL_PAGES = 45

ENDPOINT_ID = re.compile(r"^[a-z0-9][a-z0-9.\-]*$")
CAPABILITY = re.compile(r"(Audio generation|Caching|Code execution|Computer use|File search|Function calling|Grounding with Google Maps|Grounding with Google Search|"
                        r"Image generation|Live API|Search grounding|Structured outputs|Thinking|URL context|Batch API|Flex inference|Priority inference|Tuning|"
                        r"Video generation|Audio understanding|Image understanding)\s+(Supported|Not supported)(?:\s*\(([^)]*)\))?")
MODALITY_WORDS = {"text": "text", "image": "image", "images": "image", "video": "video", "audio": "audio", "pdf": "pdf", "code": "code", "embedding": "embedding"}


class GoogleConnector(BaseConnector):
    name = "google"
    label = "Google — Gemini API models, pricing, release notes; DeepMind & Google AI blogs"
    description = "Gemini API developer docs (model cards, model pages, pricing, release notes) plus the Google DeepMind and Google AI blog feeds."
    source_key = "ai.google.dev"
    version = "1"
    parser_version = "1"
    interval_seconds = 3600
    min_interval_seconds = 1800
    rate_per_min = 15
    tier = 1
    priority = 0
    expected_min_records = 30
    concurrency = 2

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [
            Target(url=f"{DOCS}/models", doc_type="model_docs", key="models", min_bytes=5000),
            Target(url=f"{DOCS}/pricing", doc_type="pricing", key="pricing", min_bytes=5000),
            Target(url=f"{DOCS}/changelog", doc_type="listing", key="changelog", min_bytes=3000),
            Target(url=DEEPMIND_RSS, doc_type="feed", key="deepmind_feed", min_bytes=1000),
            Target(url=GOOGLE_AI_RSS, doc_type="feed", key="google_ai_feed", min_bytes=1000),
        ]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        google = org_ref("google")
        key = target.key or ""
        if key == "deepmind_feed" and parsed.kind == "feed":
            deepmind = org_ref("google-deepmind")
            facts.entities.append(deepmind)
            announcement_events(facts, deepmind, parsed.feed_items, source_name="deepmind.google/blog", max_follow=15)
            facts.document_title, facts.document_entity = "Google DeepMind blog", deepmind
            return facts
        facts.entities.append(google)
        if key == "google_ai_feed" and parsed.kind == "feed":
            announcement_events(facts, google, parsed.feed_items, source_name="blog.google/technology/ai", max_follow=10)
            facts.document_title, facts.document_entity = "Google AI blog", google
        elif key == "models" and parsed.html:
            self._models(facts, google, parsed)
        elif key == "pricing" and parsed.html:
            self._pricing(facts, google, parsed)
        elif key == "changelog" and parsed.html:
            self._changelog(facts, google, parsed, res.final_url or res.url)
        elif target.doc_type == "model_page" and parsed.html:
            self._model_page(facts, google, target, parsed)
        return facts

    # ------------------------------------------------------------------------------------------ models overview
    def _models(self, facts: Facts, org: EntityRef, parsed: Parsed) -> None:
        html = parsed.html
        assert html
        facts.document_title, facts.document_entity = "Gemini models", org
        ids: list[str] = []
        for card in html.css("a.gemini-card-centered"):
            href = card.attributes.get("href") or ""
            api_id = slug_of(href)
            h3 = card.css_first("h3")
            if not h3 or not ENDPOINT_ID.match(api_id):
                continue
            name = _clean_name(node_text(h3))
            status_node = card.css_first("p.status-subtext")
            status = _status(node_text(status_node)) if status_node else None
            # two cards may share a title (GA endpoint + its preview): the later one takes its status as suffix so alias
            # resolution never fuses separately priced endpoints ("Gemini Omni Flash" / "Gemini Omni Flash Preview")
            if any(e.entity_type == "model" and e.name.lower() == name.lower() and e.identifiers.get(ID_SCHEME) != api_id for e in facts.entities):
                suffix = "Preview" if status == "preview" else api_id
                name = f"{name} {suffix}" if not name.lower().endswith(suffix.lower()) else name
            ref = _model(facts, org, name, api_id)
            desc = card.css_first("p.description-centered")
            facts.claim(ref, "description", node_text(desc) if desc else None)
            claim_status(facts, ref, status)
            if api_id not in ids:
                ids.append(api_id)
        api_aliases: dict[int, tuple[EntityRef, set[str]]] = {}
        for t in html.tables:
            hs = [clean_cell(h).lower() for h in t["headers"]]
            if not hs or hs[0] != "model" or "endpoint" not in hs:
                continue
            i_end = hs.index("endpoint")
            i_desc = hs.index("description") if "description" in hs else None
            for r in t["rows"]:
                if len(r) <= i_end:
                    continue
                raw_name = clean_cell(r[0]).replace("\xa0", " ")
                endpoints = [e for e in re.split(r"\s+", clean_cell(r[i_end])) if ENDPOINT_ID.match(e)]
                if not endpoints:
                    continue
                name = _clean_name(raw_name)
                # the same display name may carry several endpoints (GA + preview / live variants); the first one is the identity,
                # the others are recorded as a claim — never as aliases, which would merge separately priced endpoints
                same_name = next((e for e in facts.entities if e.entity_type == "model" and e.name.lower() == name.lower()), None)
                ref = same_name or _model(facts, org, name, endpoints[0])
                extra = {*endpoints} - {ref.identifiers.get(ID_SCHEME, "")}
                if extra:
                    api_aliases.setdefault(id(ref), (ref, set()))[1].update(extra)
                note = re.search(r"\(([^)]*)\)", raw_name)
                if note:
                    claim_status(facts, ref, _status(note.group(1)))
                if i_desc is not None and i_desc < len(r):
                    facts.claim(ref, "description", clean_cell(r[i_desc]))
                if endpoints[0] not in ids:
                    ids.append(endpoints[0])
        for ref, extra in api_aliases.values():
            claim_api_aliases(facts, ref, sorted(extra))
        for api_id in ids[:MAX_MODEL_PAGES]:
            facts.follow(f"{DOCS}/models/{api_id}", doc_type="model_page", key=f"model:{api_id}", meta={"api_id": api_id}, min_bytes=3000, priority=1)

    # ------------------------------------------------------------------------------------------ model page
    def _model_page(self, facts: Facts, org: EntityRef, target: Target, parsed: Parsed) -> None:
        html = parsed.html
        assert html
        table = next((t for t in html.tables if [clean_cell(h).lower() for h in t["headers"]][:2] == ["property", "description"]), None)
        if not table:
            return
        props: dict[str, str] = {}
        for r in table["rows"]:
            if len(r) >= 2:
                label = re.sub(r"^[a-z0-9_]+\s+", "", clean_cell(r[0]))          # drop the leading material-icon name ("id_card Model code")
                props[label.lower().replace("[*]", "").strip()] = clean_cell(r[1])
        api_id = props.get("model code") or target.meta.get("api_id")
        if not api_id or not ENDPOINT_ID.match(api_id):
            return
        name = _clean_name(next((t for lvl, t in html.headings if lvl == 1), None) or api_id)
        ref = _model(facts, org, name, api_id)
        facts.document_title, facts.document_entity = name, ref
        facts.claim(ref, "official_url", f"{DOCS}/models/{api_id}")
        types = props.get("supported data types", "")
        m = re.search(r"Inputs?\s+(.*?)\s+Outputs?\s+(.*)$", types, re.IGNORECASE)
        if m:
            claim_modalities(facts, ref, _modalities(m.group(1)), _modalities(m.group(2)))   # pdf → document, images → image (ontology)
        limits = props.get("token limits", "")
        m_in = re.search(r"Input token limit\s+([\d,]+)", limits)
        m_out = re.search(r"Output token limit\s+([\d,]+)", limits)
        if m_in:
            facts.claim(ref, "context_length", tokens(m_in.group(1)), unit="tokens")
        if m_out:
            facts.claim(ref, "max_output_tokens", tokens(m_out.group(1)), unit="tokens")
        caps = props.get("capabilities", "")
        if caps:
            supported: dict[str, Any] = {}
            for cap, state, detail in CAPABILITY.findall(caps):
                supported[cap] = (detail or True) if state == "Supported" else False
            if supported:
                labels = sorted(k for k, v in supported.items() if v)
                facts.claim(ref, "capabilities", normalize_capabilities(labels))
                facts.claim(ref, "capabilities_raw", labels)
                if "Function calling" in supported:
                    facts.claim(ref, "tool_calling", bool(supported["Function calling"]))
                if "Structured outputs" in supported:
                    facts.claim(ref, "structured_output", bool(supported["Structured outputs"]))
                if "Thinking" in supported:
                    facts.claim(ref, "reasoning", bool(supported["Thinking"]))
        versions: dict[str, list[str]] = {}
        for label, vid in re.findall(r"(Stable|Preview|Latest|Experimental):\s*([a-z0-9][a-z0-9.\-]*)", props.get("versions", "")):
            versions.setdefault(label.lower(), []).append(vid)
        facts.claim(ref, "versions", versions or None)
        facts.claim(ref, "knowledge_cutoff", month_year(props.get("knowledge cutoff", "")))
        facts.claim(ref, "latest_update", month_year(props.get("latest update", "")))
        for label, prop in (("release date", "release_date"), ("deprecation date", "deprecation_date"), ("shutdown date", "retirement_date"), ("retirement date", "retirement_date")):
            if props.get(label):
                dt = parse_datetime(props[label])
                facts.claim(ref, prop, dt.date().isoformat() if dt else None)

    # ------------------------------------------------------------------------------------------ pricing
    def _pricing(self, facts: Facts, org: EntityRef, parsed: Parsed) -> None:
        html = parsed.html
        assert html
        provider = provider_ref(PROVIDER_KEY)
        facts.entities.append(provider)
        facts.document_title, facts.document_entity = "Gemini Developer API pricing", provider
        for section in html.css("div.models-section"):
            h2 = section.css_first("h2")
            code = section.css_first("em a code") or section.css_first("em code")
            if not h2 or not code:
                continue
            api_id = node_text(code)
            if not ENDPOINT_ID.match(api_id):
                continue
            name = _clean_name(node_text(h2))
            ref = _model(facts, org, name, api_id)
            tiers = _tier_tables(section)
            if not tiers:
                continue
            obs = facts.price(model=ref, provider=provider, provider_model_id=api_id, meta={"from": "pricing page", "tier": "paid, standard"})
            per_unit = {k: v for k, v in tiers.items() if k.startswith("per_")}
            for unit_key, rows in per_unit.items():
                # per-second (Veo), per-request/song (Lyria), per-image tables: keep every row as stated
                table: dict[str, list[str]] = {label: lines for label, _free, lines in rows if lines and not label.startswith("used to")}
                if table:
                    obs.features[unit_key] = table
                    if unit_key == "per_request":
                        obs.per_request = next((money(lines[0]) for lines in table.values() if money(lines[0]) is not None), None)
                    if unit_key == "per_image":
                        obs.per_image = next((money(lines[0]) for lines in table.values() if money(lines[0]) is not None), None)
            standard = tiers.get("standard") or next((v for k, v in tiers.items() if not k.startswith("per_")), None)
            if standard:
                _apply_paid_rows(obs, standard, prefix="")
            if "batch" in tiers:
                batch: dict[str, Any] = {}
                _apply_paid_rows(batch, tiers["batch"], prefix="", as_dict=True)
                obs.batch_input_per_mtok = batch.get("input_per_mtok")
                obs.batch_output_per_mtok = batch.get("output_per_mtok")
            for tier in ("flex", "priority"):
                if tier in tiers:
                    extra: dict[str, Any] = {}
                    _apply_paid_rows(extra, tiers[tier], prefix="", as_dict=True)
                    for k in ("input_per_mtok", "output_per_mtok"):
                        if extra.get(k) is not None:
                            obs.features[f"{tier}_{k}"] = extra[k]
            obs.features = {k: v for k, v in obs.features.items() if v not in (None, {}, [])}
        facts.prices = [p for p in facts.prices if any(v is not None for v in p.price_tuple()[:-1]) or any(k.startswith("per_") for k in p.features)]

    # ------------------------------------------------------------------------------------------ release notes
    def _changelog(self, facts: Facts, org: EntityRef, parsed: Parsed, url: str) -> None:
        html = parsed.html
        assert html
        facts.document_title, facts.document_entity = "Gemini API release notes", org
        n_dates = 0
        for h2 in html.css("h2"):
            date = parse_datetime(node_text(h2)) if re.fullmatch(r"[A-Z][a-z]+ \d{1,2}, \d{4}", node_text(h2)) else None
            if not date:
                continue
            n_dates += 1
            if n_dates > 80:
                break
            anchor = h2.attributes.get("id") or date.date().isoformat()
            sib = h2.next
            i = 0
            while sib is not None and sib.tag != "h2":
                if sib.tag == "ul":
                    for li in sib.iter():
                        if li.tag != "li":
                            continue
                        strong = li.css_first("strong")
                        title = node_text(strong).rstrip(":") if strong else node_text(li)[:120]
                        text = node_text(li)
                        if not title:
                            continue
                        i += 1
                        codes = sorted({node_text(c) for c in li.css("code") if ENDPOINT_ID.match(node_text(c))})
                        is_release = bool(RELEASE_WORDS.search(text) and MODEL_WORDS.search(text)) or bool(re.search(r"\b(GA|generally available|released|preview|deprecat)", text, re.IGNORECASE))
                        facts.event("ANNOUNCEMENT", "release" if is_release else "company", f"Gemini API: {title}", entity=org, importance=2 if is_release else 1,
                                    effective_at=date, dedupe_key=f"ANNOUNCEMENT:{url}#{anchor}:{i}", source_url=f"{url}#{anchor}",
                                    meta={"source": "ai.google.dev/gemini-api/docs/changelog", "summary": text[:300], "models": codes[:10], "is_release": is_release})
                sib = sib.next


# ---------------------------------------------------------------------------------------------- helpers
def _model(facts: Facts, org: EntityRef, name: str, api_id: str) -> EntityRef:
    for e in facts.entities:
        if e.entity_type == "model" and e.identifiers.get(ID_SCHEME) == api_id:
            return e
    # `gemini_model_id` stays the primary scheme (existing rows resolve on it); `google_model_id` is emitted alongside for the provider key
    ref = model_ref(facts, name, org, api_id=api_id, provider_key="gemini", aliases=[api_id] if api_id != name else [], identifiers={"google_model_id": api_id})
    facts.claim(ref, "api_model_id", api_id)
    return ref


def _clean_name(name: str) -> str:
    name = name.replace("\xa0", " ").replace("🍌", "").strip()
    name = re.sub(r"\((Shut down|Deprecated|Retired)\)", "", name, flags=re.IGNORECASE).strip()
    return re.sub(r"\s+", " ", name)


def _status(text: str) -> str | None:
    low = (text or "").lower()
    if "shut down" in low or "retired" in low:
        return "retired"
    if "deprecat" in low:
        return "deprecated"
    if "experimental" in low or "preview" in low:
        return "preview"
    if "stable" in low or "ga" == low.strip():
        return "active"
    return None


def _modalities(text: str) -> list[str]:
    """Raw modality words of a 'Supported data types' cell (canonicalised by `claim_modalities`: pdf → document)."""
    out: list[str] = []
    for tok in re.split(r"[,/]|\band\b", text.lower()):
        tok = tok.strip().strip(".")
        mod = MODALITY_WORDS.get(tok)
        if mod and mod not in out:
            out.append(mod)
    return out


def _table_rows(table: Node) -> tuple[str, list[tuple[str, str, list[str]]]]:
    """(paid-tier unit: tokens|second|request|image, [(row label, free cell, paid cell lines)])"""
    unit = "tokens"
    for th in table.css("thead th"):
        m = re.search(r"per\s+(1M tokens|second|request|image|minute)", node_text(th), re.IGNORECASE)
        if m:
            unit = m.group(1).lower().replace("1m ", "")
    rows: list[tuple[str, str, list[str]]] = []
    for tr in table.css("tbody tr"):
        tds = tr.css("td")
        if len(tds) < 2:
            continue
        label = node_text(tds[0]).lower()
        free = node_text(tds[1]) if len(tds) >= 3 else ""
        paid_html = tds[-1].html or ""
        lines = [clean_cell(html_unescape(re.sub(r"<[^>]+>", " ", part))) for part in re.split(r"<br\s*/?>", paid_html)]
        rows.append((label, free, [ln for ln in lines if ln]))
    return unit, rows


def _tier_tables(section: Node) -> dict[str, list[tuple[str, str, list[str]]]]:
    """Tables that follow one `div.models-section` until the next one: {tier: [(row label, free cell, paid cell lines)]}.
    Token-priced models use a `devsite-selector` with one `section` per tier (Standard / Batch / Flex / Priority); per-second /
    per-request models (Veo, Lyria…) have a single bare table, exposed under the pseudo tier `per_<unit>`."""
    tiers: dict[str, list[tuple[str, str, list[str]]]] = {}
    sib = section.next
    while sib is not None:
        if sib.tag == "div" and "models-section" in (sib.attributes.get("class") or ""):
            break
        if sib.tag in ("h2",):
            break
        if sib.tag != "-text":
            sections = [sib] if sib.tag == "section" else sib.css("section")
            for sec in sections:
                h3 = sec.css_first("h3")
                table = sec.css_first("table")
                if not table:
                    continue
                unit, rows = _table_rows(table)
                tier = (node_text(h3).lower() if h3 else "standard").split()[0]
                if unit != "tokens":
                    tier = f"per_{unit}"
                if rows:
                    tiers.setdefault(tier, rows)
            if not sections:
                for table in ([sib] if sib.tag == "table" else sib.css("table")):
                    unit, rows = _table_rows(table)
                    if rows:
                        tiers.setdefault("standard" if unit == "tokens" else f"per_{unit}", rows)
        sib = sib.next
    return tiers


def _apply_paid_rows(target: Any, rows: list[tuple[str, str, list[str]]], *, prefix: str, as_dict: bool = False) -> None:
    """Map 'Input price' / 'Output price' / 'Context caching price' rows of a paid-tier table onto a PriceObs (or a dict)."""
    def put(key: str, value: Any) -> None:
        if value is None:
            return
        if as_dict:
            target[key] = value
        elif hasattr(target, key):
            setattr(target, key, value)
        else:
            target.features[key] = value

    unit_rx = re.compile(r"per (image|second|video|request|query|minute|song)|/\s*(image|second|sec|video|request|query|min)\b", re.IGNORECASE)
    qualifier_rx = re.compile(r"\(([^)]*)\)")
    for label, free, lines in rows:
        if not lines:
            continue
        first = lines[0]
        bare_first = qualifier_rx.sub("", first)                    # "(text / image / video)" is a modality list, not a unit
        primary = money(bare_first) if not unit_rx.search(bare_first) else None
        if re.match(r"(text )?(input|output) price", label):
            kind = "input" if "input" in label.split(" price")[0] else "output"
            put(f"{kind}_per_mtok", primary)
            if unit_rx.search(bare_first) and re.search(r"per image|/\s*image", bare_first, re.IGNORECASE):
                put("per_image", money(bare_first))
            for ln in lines[1:]:
                if ln.lower().startswith("equivalent"):
                    continue
                v = money(qualifier_rx.sub("", ln)) if not unit_rx.search(qualifier_rx.sub("", ln)) else None
                if v is None:
                    continue
                qual = qualifier_rx.search(ln)
                if re.search(r">\s*200k|>\s*128k|long context|prompts >", ln, re.IGNORECASE):
                    put(f"long_context_{kind}_per_mtok", v)
                elif re.search(r"\bstarting\b", ln, re.IGNORECASE):
                    put(f"scheduled_{kind}_per_mtok", v)
                    put(f"scheduled_{kind}_effective", re.sub(r"^\$[\d.,]+\s*", "", qualifier_rx.sub("", ln)).strip())
                elif qual:
                    mod = re.sub(r"[^a-z]+", "_", qual.group(1).lower()).strip("_")
                    put(f"{mod}_{kind}_per_mtok", v)
            if "through" in bare_first.lower() and not as_dict:
                target.features.setdefault("promotional_until", re.sub(r"^\$[\d.,]+\s*(through)?\s*", "", bare_first).strip().rstrip("."))
            if free.lower().startswith("free of charge") and not as_dict:
                target.features["free_tier"] = True
        elif re.match(r"(image|audio|video) input price", label):
            put(f"{label.split()[0]}_input_per_mtok", primary)
        elif label.startswith("context caching price"):
            for ln in lines:
                v = money(ln)
                if v is None:
                    continue
                if "storage" in ln.lower():
                    put("cache_storage_per_mtok_hour", v)
                    break
                if not (as_dict and "cached_input_per_mtok" in target) and not (not as_dict and target.cached_input_per_mtok is not None):
                    put("cached_input_per_mtok", v)
                elif re.search(r">\s*200k|prompts >", ln, re.IGNORECASE):
                    put("long_context_cached_input_per_mtok", v)
        elif label.startswith("grounding with google search"):
            m = re.search(r"\$\s*([\d.]+)\s*(?:per|/)\s*1,?000", " ".join(lines))
            if m:
                put("search_grounding_per_1k_requests", float(m.group(1)))


CONNECTORS = [GoogleConnector]