"""Cohere — docs models overview (served as Markdown by appending `.md`) + blog listing (Next.js RSC payload).

Sources (tier 1):
  * models → docs.cohere.com/docs/models.md : one table per family (Command, North, Embed, Rerank, Parse, Audio, Aya…) with model id,
             status, description, modalities, context length, max output, endpoints; platform tables (Bedrock / SageMaker / Azure / OCI ids)
  * blog   → cohere.com/blog : post objects (title, slug, date, tags) embedded in the React Server Components payload → ANNOUNCEMENT events
"""
from __future__ import annotations

import re
from typing import Any

from aiatlas.ontology.openness import normalize_openness
from aiatlas.ontology.taxonomy import normalize_modalities
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.feeds import FeedItem
from aiatlas.sdk.extract.numbers import parse_active_params, parse_param_count
from aiatlas.sdk.facts import EntityRef, Facts, Target
from aiatlas.sdk.fetch import FetchResult

from ._common import (
    announcement_events,
    claim_api_aliases,
    claim_status,
    clean_cell,
    model_ref,
    next_flight_payload,
    tokens,
)

DOCS = "https://docs.cohere.com/docs/models"
BLOG = "https://cohere.com/blog"
PROVIDER_KEY = "cohere"
ID_SCHEME = "cohere_model_id"
API_ID = re.compile(r"^[a-z0-9][a-z0-9.\-]*$")
ALIAS_FOR = re.compile(r"^Alias for `?([a-z0-9.\-]+)`?", re.IGNORECASE)
PLATFORM_COLUMNS = {"amazon bedrock model id": "bedrock_model_id", "azure ai foundry": "foundry_model_id", "oracle oci generative ai service": "oci_model_id"}
NOT_AN_ID = re.compile(r"n/a|unique per deployment|coming soon", re.IGNORECASE)
MODALITIES = {"text": "text", "images": "image", "image": "image", "audio": "audio", "video": "video"}


class CohereConnector(BaseConnector):
    name = "cohere"
    label = "Cohere — models overview, blog"
    description = "Cohere docs models overview (Markdown tables per family + platform ids) and the Cohere blog listing."
    source_key = "cohere.com"
    version = "1"
    parser_version = "1"
    interval_seconds = 7200
    min_interval_seconds = 3600
    rate_per_min = 12
    tier = 1
    priority = 1
    expected_min_records = 25
    concurrency = 2

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [
            Target(url=f"{DOCS}.md", doc_type="model_docs", key="models", min_bytes=5000, meta={"content_type": "text/markdown"}),
            Target(url=BLOG, doc_type="listing", key="blog", min_bytes=20000),
        ]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        org = org_ref("cohere")
        facts.entities.append(org)
        key = target.key or ""
        if key == "models" and parsed.markdown:
            self._models(facts, org, parsed)
        elif key == "blog" and parsed.html:
            items = blog_items(res.content) or [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://cohere\.com/blog/(?!tag/)[^/?#]+$") if len(t) > 12]
            announcement_events(facts, org, items, source_name="cohere.com/blog", max_follow=15)
            facts.document_title, facts.document_entity = "The Cohere Blog", org
        return facts

    # ------------------------------------------------------------------------------------------ models overview
    def _models(self, facts: Facts, org: EntityRef, parsed: Parsed) -> None:
        md = parsed.markdown
        assert md
        facts.document_title, facts.document_entity = "Cohere models", org
        provider = provider_ref(PROVIDER_KEY)
        facts.entities.append(provider)
        families = _table_families(md.body)
        aliases: list[tuple[str, str, dict[str, str]]] = []       # (alias id, target id, row)
        for i, t in enumerate(md.tables):
            hs = [clean_cell(h).lower() for h in t["headers"]]
            if not hs or "model name" not in hs[0]:
                continue
            family = families[i] if i < len(families) else None
            is_platform = any(h in PLATFORM_COLUMNS or "sagemaker" in h for h in hs[1:])
            for r in t["rows"]:
                if len(r) < 2:
                    continue
                row = {h: clean_cell(c) for h, c in zip(hs, r, strict=False)}
                api_id = row.get(hs[0], "")
                if not API_ID.match(api_id):
                    continue
                if is_platform:
                    # alias rows (command-r-plus → command-r-plus-04-2024) carry the platform ids of the model they point at
                    target_id = next((t for a, t, _row in aliases if a == api_id), api_id)
                    ref = _model(facts, org, target_id, family)
                    if target_id != api_id and api_id not in ref.aliases:
                        ref.aliases.append(api_id)
                    for h, prop in PLATFORM_COLUMNS.items():
                        v = row.get(h, "")
                        if v and not NOT_AN_ID.search(v):
                            facts.claim(ref, prop, v.strip("'’‘"))
                    continue
                desc = row.get("description", "")
                m_alias = ALIAS_FOR.match(desc)
                if m_alias:
                    aliases.append((api_id, m_alias.group(1), row))
                    continue
                ref = _model(facts, org, api_id, family)
                self._row_claims(facts, ref, row, desc)
                facts.relate(provider, "available_through", ref) if False else None  # (relation direction: model available_through provider)
                facts.relate(ref, "available_through", provider, attributes={"endpoints": _endpoints(row.get("endpoints", ""))})
        for alias, target_id, row in aliases:
            ref = next((e for e in facts.entities if e.entity_type == "model" and e.identifiers.get(ID_SCHEME) == target_id), None)
            if ref is None:
                ref = _model(facts, org, target_id, None)
            if alias not in ref.aliases:
                ref.aliases.append(alias)
            existing = next((c.value for c in facts.claims if c.entity is ref and c.property == "api_aliases"), [])
            facts.claims = [c for c in facts.claims if not (c.entity is ref and c.property in ("api_aliases", "api_alias"))]
            claim_api_aliases(facts, ref, [*existing, alias])
            status, dep = _status(row.get("status", ""))
            if status == "deprecated" and dep:
                facts.claim(ref, "alias_deprecation_date", dep)

    def _row_claims(self, facts: Facts, ref: EntityRef, row: dict[str, str], desc: str) -> None:
        facts.claim(ref, "description", desc or None)
        status, dep = _status(row.get("status", ""))
        claim_status(facts, ref, status)
        facts.claim(ref, "deprecation_date", dep)
        mods_raw = row.get("modality") or row.get("modalities") or ""
        mods = normalize_modalities(_modalities(mods_raw))
        if mods:
            facts.claim(ref, "modalities", mods)
            facts.claim(ref, "modalities_input", mods)
            facts.claim(ref, "vision", "image" in mods)
        ctx = row.get("context length", "")
        if ctx:
            facts.claim(ref, "context_length", tokens(ctx), unit="tokens")
        out = row.get("maximum output tokens", "")
        if out:
            facts.claim(ref, "max_output_tokens", tokens(out), unit="tokens")
        dims = row.get("dimensions", "")
        if dims:
            nums = [int(x) for x in re.findall(r"\d{2,5}", dims)]
            facts.claim(ref, "embedding_dimensions", sorted(set(nums)) if nums else None)
            facts.claim(ref, "modalities_output", ["embedding"])
        facts.claim(ref, "endpoints", _endpoints(row.get("endpoints", "")) or None)
        low = desc.lower()
        if re.search(r"\bMoE\b|mixture[- ]of[- ]experts", desc, re.IGNORECASE):
            facts.claim(ref, "is_moe", True)
        if "open-weight" in low or "open weight" in low or "open source" in low or "open-source" in low:
            # the docs prose says "open"; the weights exist on the hub (Cohere Labs) — the category itself is derived from the licence elsewhere
            facts.claim(ref, "weights_available", True)
            facts.claim(ref, "openness", normalize_openness("open-weights"))
            facts.claim(ref, "openness_raw", "open-source" if "open source" in low or "open-source" in low else "open-weights")
        if re.search(r"\d+(?:\.\d+)?B\s+(total|parameter|params|instruct|model)", desc):
            facts.claim(ref, "parameter_count", parse_param_count(desc))
            active = parse_active_params(desc)
            facts.claim(ref, "active_parameter_count", active)
            if active or re.search(r"\bMoE\b|mixture of experts", desc, re.IGNORECASE):
                facts.claim(ref, "is_moe", True)
        if re.search(r"\breasoning model\b|able to 'think'|think before", low):
            facts.claim(ref, "reasoning", True)
        if "tool use" in low or "tool-use" in low or "function calling" in low:
            facts.claim(ref, "tool_calling", True)


# ---------------------------------------------------------------------------------------------- helpers
def _model(facts: Facts, org: EntityRef, api_id: str, family: str | None) -> EntityRef:
    for e in facts.entities:
        if e.entity_type == "model" and e.identifiers.get(ID_SCHEME) == api_id:
            return e
    ref = model_ref(facts, api_id, org, api_id=api_id, provider_key=PROVIDER_KEY, family=family)
    facts.claim(ref, "api_model_id", api_id)
    facts.claim(ref, "official_url", DOCS)
    return ref


def _status(cell: str) -> tuple[str | None, str | None]:
    c = cell.strip()
    if not c:
        return None, None
    low = c.lower()
    if low == "live":
        return "active", None
    if low.startswith("deprecated"):
        dt = parse_datetime(re.sub(r"(?i)^deprecated\s*", "", c).replace("Sept", "Sep")) if len(c) > 10 else None
        return "deprecated", dt.date().isoformat() if dt else None
    if "retired" in low or "removed" in low:
        return "retired", None
    if "preview" in low or "beta" in low:
        return "preview", None
    return None, None


def _modalities(cell: str) -> list[str]:
    out: list[str] = []
    for tok in re.split(r",|/|\band\b|\(", cell.lower()):
        tok = tok.strip().strip(")").strip()
        for word, mod in MODALITIES.items():
            if tok.startswith(word) and mod not in out:
                out.append(mod)
    return out


def _endpoints(cell: str) -> list[str]:
    return [e.strip() for e in re.split(r",", clean_cell(cell)) if e.strip()]


def _table_families(body: str) -> list[str | None]:
    """The h2 heading (model family) preceding each Markdown table, in table order."""
    out: list[str | None] = []
    current: str | None = None
    lines = body.split("\n")
    i = 0
    while i < len(lines):
        m = re.match(r"^##\s+(.+?)\s*$", lines[i])
        if m:
            current = m.group(1).strip()
        if lines[i].lstrip().startswith("|") and i + 1 < len(lines) and re.match(r"^\s*\|?\s*:?-{2,}", lines[i + 1]):
            out.append(current if current and not current.lower().startswith("what can") else None)
            i += 2
            while i < len(lines) and lines[i].lstrip().startswith("|"):
                i += 1
            continue
        i += 1
    return out


_POST = re.compile(r'"slug":\{"_type":"slug","current":"blog/([^"/]+)"\}')


def blog_items(content: bytes) -> list[FeedItem]:
    """Post objects live in the RSC payload: …"date":"2026-09-10T…", … "slug":{"current":"blog/<slug>"}, … "title":"…", …"""
    payload = next_flight_payload(content)
    if not payload:
        return []
    items: dict[str, FeedItem] = {}
    matches = list(_POST.finditer(payload))
    for idx, m in enumerate(matches):
        slug = m.group(1)
        start = matches[idx - 1].end() if idx else max(0, m.start() - 20000)
        before = payload[start:m.start()]
        after = payload[m.end():m.end() + 4000]
        date_m = None
        for date_m in re.finditer(r'"date":"(20\d\d-\d\d-\d\d[^"]*)"', before):
            pass
        title_m = re.search(r'"title":"((?:[^"\\]|\\.)*)"', after)
        if not title_m:
            continue
        title = _unescape(title_m.group(1))
        tags = [t for t in re.findall(r'"tag":\["((?:[^"\\]|\\.)*)"\]', after[:3000])]
        url = f"{BLOG}/{slug}"
        if url not in items:
            items[url] = FeedItem(id=url, url=url, title=title[:300], summary=None, published_at=parse_datetime(date_m.group(1)) if date_m else None,
                                  updated_at=None, categories=[_unescape(t) for t in tags[:5]])
    return sorted(items.values(), key=lambda it: it.published_at.timestamp() if it.published_at else 0, reverse=True)


def _unescape(s: str) -> str:
    import json

    try:
        return json.loads(f'"{s}"').encode("latin-1", errors="ignore").decode("utf-8", errors="ignore") if "Ã" in s or "â" in s else json.loads(f'"{s}"')
    except ValueError:
        return s


CONNECTORS: list[Any] = [CohereConnector]