"""Official provider pricing pages that render server-side (tier 1). One connector class per provider domain so that every document,
claim and price row carries the right `source` (the SDK binds one `source_key` per connector) — module name kept as `provider_pricing`.

  * groq_pricing      https://console.groq.com/docs/models — HTML tables (Production / Systems / Preview): model id (`div[id]`), name,
                      speed, "$0.15 input $0.60 output" per 1M tokens, rate limits, context window, max completion tokens.
                      https://groq.com/pricing is a client-rendered Next.js page with no price in the HTML → fetched with `escalate=True`
                      (only used when a browser/Scrapfly key is configured) and otherwise recorded as "no data" in the run log.
  * together_pricing  https://www.together.ai/pricing — Webflow tables: serverless (input / cached / output per 1M tokens) and batch
                      tables share the model page slug (`/models/<slug>`), merged into one price row per model.
  * fireworks_pricing https://docs.fireworks.ai/serverless/pricing.md — the docs platform serves Markdown: "Standard" and "Priority"
                      cells are `input / cached input / output`; Fast/US variants are separate rows sharing the base model link.
                      https://fireworks.ai/pricing only lists size-based embedding tiers and GPU hours, so it is not fetched.
"""
from __future__ import annotations

import re
from typing import Any

from aiatlas.connectors._identity import family_ref, model_identity, org_ref_in
from aiatlas.registry import org_by_hf, organizations, provider_ref
from aiatlas.sdk.connector import BaseConnector, Parsed, RunContext
from aiatlas.sdk.extract.numbers import parse_context_length
from aiatlas.sdk.facts import EntityRef, Facts, Target
from aiatlas.sdk.fetch import FetchResult

MONEY = re.compile(r"\$\s*(\d+(?:\.\d+)?)")
INPUT_OUTPUT = re.compile(r"\$\s*(\d+(?:\.\d+)?)\s*input\s*\$\s*(\d+(?:\.\d+)?)\s*output", re.IGNORECASE)


def _money(s: str) -> float | None:
    m = MONEY.search((s or "").replace(",", ""))
    return float(m.group(1)) if m else None


def _org_for(facts: Facts, *candidates: str | None) -> EntityRef | None:
    """Registry organization from a vendor slug / display word: exact hf_org, key, name or alias match only — never a guess."""
    for c in candidates:
        if not c:
            continue
        low = c.strip().lower()
        known = org_by_hf(low)
        if known:
            return org_ref_in(facts, known["key"])
        for key, o in organizations().items():
            names = {key, o["name"].lower(), *(a.lower() for a in o.get("aliases", []))}
            if low in names:
                return org_ref_in(facts, "meta-ai" if key == "meta" and "meta-ai" in organizations() else key)   # models come from the lab, not the holding
    return None


def _org_for_id(facts: Facts, model_id: str | None, *fallback_words: str | None) -> EntityRef | None:
    """Developer organisation of a provider's model id (`openai/gpt-oss-120b`, `fireworks/kimi-k3`, `minimax-m3`): first-party prefix or
    family word through the shared identity helper, then an exact registry match on the display words."""
    ident = model_identity(model_id) if model_id else None
    if ident and ident.org_key:
        return org_ref_in(facts, ident.org_key)
    return _org_for(facts, *fallback_words)


class _PricingBase(BaseConnector):
    version = "2"
    parser_version = "2"
    interval_seconds = 6 * 3600
    min_interval_seconds = 3 * 3600
    max_interval_seconds = 3 * 86400
    tier = 1
    priority = 1
    concurrency = 1
    provider_key = ""

    def _model(self, facts: Facts, name: str, *, ids: dict[str, str], org: EntityRef | None, aliases: list[str]) -> EntityRef:
        for e in facts.entities:
            if e.entity_type == "model" and ids and any(e.identifiers.get(k) == v for k, v in ids.items()):
                return e
        ref = facts.entity("model", name, identifiers=ids, organization=org, aliases=[a for a in dict.fromkeys(aliases) if a and a != name],
                           family=family_ref(name, org), identity_confidence="medium")
        if org:
            facts.relate(org, "develops", ref)
        return ref


# ================================================================================================ Groq
class GroqPricingConnector(_PricingBase):
    name = "groq_pricing"
    label = "GroqCloud — supported models, limits and prices"
    description = "console.groq.com/docs/models tables (server-rendered) and the marketing pricing page (client-rendered, escalation only)."
    source_key = "groq.com"
    rate_per_min = 12
    expected_min_records = 8
    provider_key = "groq"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url="https://console.groq.com/docs/models", doc_type="model_docs", key="groq_models", min_bytes=20000, priority=1),
                Target(url="https://groq.com/pricing", doc_type="pricing", key="groq_pricing", min_bytes=5000, priority=3, escalate=True)]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        provider = provider_ref(self.provider_key)
        facts.entities.append(provider)
        facts.document_entity = provider
        if not parsed.html:
            return facts
        if target.key == "groq_pricing":
            if "$" not in parsed.html.text:
                ctx.log.warning("groq.com/pricing has no server-rendered prices (client-side app); nothing extracted", extra={"url": target.url})
            facts.document_title = parsed.html.title
            return facts
        html = parsed.html
        section = "production"
        for node in html.tree.body.traverse() if html.tree and html.tree.body else []:
            if node.tag == "h2":
                t = node.text(strip=True).lower()
                section = "preview" if "preview" in t else "deprecated" if "deprecated" in t else "system" if "system" in t else "production" if "production" in t else section
            elif node.tag == "tr" and node.css_first("td"):
                self._row(facts, provider, node, section)
        facts.document_title = html.title
        return facts

    def _row(self, facts: Facts, provider: EntityRef, tr: Any, section: str) -> None:
        cells = tr.css("td")
        if len(cells) < 6:
            return
        first = cells[0]
        id_node = first.css_first("div[id]")
        model_id = (id_node.attributes.get("id") if id_node else None) or (first.css_first("span.font-mono").text(strip=True) if first.css_first("span.font-mono") else None)
        link = first.css_first("a")
        name = link.text(strip=True) if link else None
        if not model_id or not name:
            return
        if section == "deprecated":
            return
        img = first.css_first("img[alt]")
        vendor = model_id.split("/")[0] if "/" in model_id else None
        org = _org_for_id(facts, model_id, vendor, img.attributes.get("alt") if img else None)
        ref = self._model(facts, name, ids={"groq_model_id": model_id}, org=org, aliases=[model_id, model_id.split("/")[-1]])
        facts.claim(ref, "groq_model_id", model_id)
        texts = [c.text(separator=" ", strip=True) for c in cells]
        ctx_len = parse_context_length(texts[4]) if texts[4] not in ("-", "") else None
        max_out = parse_context_length(texts[5]) if texts[5] not in ("-", "") else None
        badge = first.css_first("span.select-none")
        tier_label = badge.text(strip=True).lower() if badge else None
        features: dict[str, Any] = {"groq_section": section}
        if tier_label:
            features["groq_tier"] = tier_label
        speed = texts[1].replace(",", "")
        if speed.isdigit():
            features["output_tokens_per_second"] = int(speed)
        if texts[3] and texts[3] != "-":
            features["rate_limit_developer_plan"] = texts[3]
        m = INPUT_OUTPUT.search(texts[2])
        if ctx_len:
            features["context_length"] = ctx_len            # Groq's deployment limit — the model's own context is claimed by the lab's docs
        if max_out:
            features["max_completion_tokens"] = max_out
        if m:
            facts.price(model=ref, provider=provider, provider_model_id=model_id, input_per_mtok=float(m.group(1)), output_per_mtok=float(m.group(2)),
                        context_length=ctx_len, max_output_tokens=max_out, features=features, meta={"section": section})
        else:
            low = texts[2].lower()
            if "per hour" in low or "per 1m characters" in low or "contact" in low:
                facts.claim(ref, "groq_pricing_note", texts[2][:120])
            facts.relate(ref, "available_through", provider, attributes={"provider_model_id": model_id, **features})
        if section == "preview":
            facts.claim(ref, "groq_status", "preview")        # Groq's own availability tier; the model's `status` belongs to its lab


# ================================================================================================ Together AI
class TogetherPricingConnector(_PricingBase):
    name = "together_pricing"
    label = "Together AI — serverless inference pricing"
    description = "together.ai/pricing tables (serverless input / cached / output and batch prices per 1M tokens)."
    source_key = "together.ai"
    rate_per_min = 12
    expected_min_records = 10
    provider_key = "together-ai"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url="https://www.together.ai/pricing", doc_type="pricing", key="together_pricing", min_bytes=50000, priority=1)]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        provider = provider_ref(self.provider_key)
        facts.entities.append(provider)
        facts.document_entity = provider
        html = parsed.html
        if not html:
            return facts
        rows: dict[str, dict[str, Any]] = {}
        for table in html.css("table"):
            headers = [th.text(strip=True).lower() for th in table.css("thead th")]
            if not headers or headers[0] != "model" or "input" not in headers:
                continue
            is_batch = "batch" in (table.parent.text(separator=" ", strip=True).lower()[:400] if table.parent else "") and "cached" not in table.text().lower()
            for tr in table.css("tbody tr"):
                tds = tr.css("td")
                if len(tds) < 3:
                    continue
                a = tds[0].css_first("a[href*='/models/']")
                slug = a.attributes.get("href", "").rstrip("/").rsplit("/", 1)[-1] if a else None
                name = tds[0].text(strip=True)
                if not name:
                    continue
                key = slug or name.lower()
                entry = rows.setdefault(key, {"name": name, "slug": slug, "input": None, "cached": None, "output": None, "batch_in": None, "batch_out": None})
                in_txt = tds[1].text(separator=" ", strip=True)
                out_txt = tds[2].text(separator=" ", strip=True)
                cached = None
                if "(cached)" in in_txt.lower():
                    parts = [p for p in MONEY.findall(in_txt)]
                    if len(parts) >= 2:
                        cached = float(parts[1])
                if is_batch or (entry["input"] is not None and cached is None and entry["cached"] is not None):
                    entry["batch_in"], entry["batch_out"] = _money(in_txt), _money(out_txt)
                else:
                    entry["input"], entry["output"], entry["cached"] = _money(in_txt), _money(out_txt), cached
        for key, e in rows.items():
            if e["input"] is None and e["output"] is None:
                continue
            ids = {"together_ai_model_slug": e["slug"]} if e["slug"] else {}
            org = _org_for_id(facts, e["slug"], e["name"].split(" ")[0])
            ref = self._model(facts, e["name"], ids=ids, org=org, aliases=[e["slug"] or ""])
            facts.price(model=ref, provider=provider, provider_model_id=e["slug"], input_per_mtok=e["input"], output_per_mtok=e["output"],
                        cached_input_per_mtok=e["cached"], batch_input_per_mtok=e["batch_in"], batch_output_per_mtok=e["batch_out"],
                        features={"serverless": True}, source_url=f"https://www.together.ai/models/{e['slug']}" if e["slug"] else None)
        facts.document_title = html.title
        return facts


# ================================================================================================ Fireworks AI
class FireworksPricingConnector(_PricingBase):
    name = "fireworks_pricing"
    label = "Fireworks AI — serverless per-token pricing"
    description = "docs.fireworks.ai/serverless/pricing (Markdown): Standard and Priority input / cached / output prices per 1M tokens."
    source_key = "fireworks.ai"
    rate_per_min = 12
    expected_min_records = 10
    provider_key = "fireworks-ai"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url="https://docs.fireworks.ai/serverless/pricing.md", doc_type="pricing", key="fireworks_serverless", min_bytes=2000, priority=1,
                       accept="text/markdown, text/plain;q=0.9, */*;q=0.5", meta={"content_type": "text/markdown"})]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        provider = provider_ref(self.provider_key)
        facts.entities.append(provider)
        facts.document_entity = provider
        md = parsed.markdown
        if not md:
            return facts
        base_names: dict[str, str] = {}
        refs: dict[str, EntityRef] = {}
        for table in md.tables:
            headers = [h.strip().lower() for h in table["headers"]]
            if not headers or headers[0] != "model" or "standard" not in headers:
                continue
            i_std, i_pri = headers.index("standard"), headers.index("priority") if "priority" in headers else None
            i_res = next((i for i, h in enumerate(headers) if "reserved" in h), None)
            for row in table["rows"]:
                if len(row) <= i_std:
                    continue
                label, link = _md_link(row[0])
                if not label:
                    continue
                slug = None
                if link:
                    m = re.search(r"app\.fireworks\.ai/models/([\w.-]+/[\w.-]+)", link)
                    slug = m.group(1) if m else None
                std = [float(x) for x in MONEY.findall(row[i_std].replace("\\$", "$"))]
                if len(std) < 2:
                    continue
                inp, cached, out = (std[0], std[1], std[2]) if len(std) >= 3 else (std[0], None, std[1])
                base_key = slug or label.lower()
                base_name = base_names.setdefault(base_key, label)
                variant = label[len(base_name):].strip().lower().replace(" ", "-") if label.startswith(base_name) and label != base_name else None
                ref = refs.get(base_key)
                if ref is None:
                    org = _org_for_id(facts, slug, base_name.split(" ")[0])
                    # canonical scheme `fireworks_model_id`; the historical hyphenated scheme is kept as a second identifier so existing rows keep resolving
                    ids = {"fireworks_model_id": slug, "fireworks-ai_model_id": slug} if slug else {}
                    ref = refs[base_key] = self._model(facts, base_name, ids=ids, org=org, aliases=[slug or "", slug.split("/")[-1] if slug else ""])
                features: dict[str, Any] = {"serving_path": variant or "standard"}
                if i_pri is not None and len(row) > i_pri:
                    pri = [float(x) for x in MONEY.findall(row[i_pri].replace("\\$", "$"))]
                    if len(pri) >= 2:
                        features["priority"] = {"input_per_mtok": pri[0], "cached_input_per_mtok": pri[1] if len(pri) >= 3 else None, "output_per_mtok": pri[-1]}
                if i_res is not None and len(row) > i_res and "✓" in row[i_res]:
                    features["reserved_throughput"] = True
                pmid = f"{slug}:{variant}" if slug and variant else slug or (f"{label.lower()}" if variant else None)
                facts.price(model=ref, provider=provider, provider_model_id=pmid, input_per_mtok=inp, output_per_mtok=out, cached_input_per_mtok=cached,
                            features=features, source_url=link or None, meta={"row_label": label})
        facts.claim(provider, "batch_discount", "50% of serverless pricing") if "50% of serverless" in md.body else None
        facts.document_title = next((t for lvl, t in md.headings if lvl == 1), "Fireworks serverless pricing")
        return facts


def _md_link(cell: str) -> tuple[str | None, str | None]:
    m = re.match(r"\s*\[([^\]]+)\]\(([^)]+)\)", cell or "")
    if m:
        return m.group(1).strip(), m.group(2).strip()
    text = re.sub(r"[*_`]", "", cell or "").strip()
    return (text or None), None


CONNECTORS = [GroqPricingConnector, TogetherPricingConnector, FireworksPricingConnector]