"""Shared helpers for lab connectors: announcement events from feeds/listings, key/value tables, model-name normalisation."""
from __future__ import annotations

import re
from datetime import datetime
from typing import Any

from aiatlas.connectors._identity import family_ref
from aiatlas.ontology.licenses import normalize_license
from aiatlas.ontology.openness import derive_openness, openness_dimensions
from aiatlas.ontology.taxonomy import normalize_modalities, normalize_status
from aiatlas.sdk.extract.dates import parse_datetime
from aiatlas.sdk.extract.feeds import FeedItem
from aiatlas.sdk.facts import EntityRef, Facts, Target

RELEASE_WORDS = re.compile(r"\b(introducing|announcing|launch(?:ing|es|ed)?|releas(?:e|es|ing|ed)|now available|new model|preview|deprecat|pricing|price)\b", re.IGNORECASE)
MODEL_WORDS = re.compile(r"\b(model|gpt|claude|gemini|llama|mistral|qwen|deepseek|grok|command|phi|nemotron|gemma|sonnet|opus|haiku|o\d|codex|sora|veo|imagen|whisper|"
                         r"embedding|reasoning|agent)\w*", re.IGNORECASE)


def announcement_events(facts: Facts, org: EntityRef, items: list[FeedItem], *, source_name: str, follow: bool = True, max_follow: int = 30,
                        needs_llm: bool = True, importance_default: int = 1) -> int:
    """One ANNOUNCEMENT event per feed item (deduped by URL) and, optionally, a follow-up fetch of the article for LLM extraction."""
    n = 0
    for i, it in enumerate(items):
        if not it.url or not it.title:
            continue
        text = f"{it.title} {it.summary or ''}"
        is_release = bool(RELEASE_WORDS.search(text) and MODEL_WORDS.search(text))
        importance = 2 if is_release else importance_default
        facts.event("ANNOUNCEMENT", "release" if is_release else "company", f"{org.name}: {it.title}", entity=org, importance=importance,
                    effective_at=it.published_at, dedupe_key=f"ANNOUNCEMENT:{it.url}", source_url=it.url,
                    meta={"source": source_name, "categories": it.categories[:5], "summary": (it.summary or "")[:300], "is_release": is_release})
        if follow and i < max_follow:
            facts.follow(it.url, doc_type="news", needs_llm=needs_llm and is_release, priority=1 if is_release else 3,
                         meta={"llm_task": "release_announcement", "published_at": it.published_at.isoformat() if it.published_at else None, "title": it.title})
        n += 1
    return n


def kv_tables(tables: list[dict[str, Any]]) -> dict[str, str]:
    """Merge 2-column key/value tables into one dict (lower-cased keys)."""
    out: dict[str, str] = {}
    for t in tables:
        rows = t.get("rows") or []
        headers = t.get("headers") or []
        if headers and len(headers) == 2 and rows:
            for r in rows:
                if len(r) >= 2 and r[0]:
                    out.setdefault(r[0].strip().lower(), r[1].strip())
        elif rows and all(len(r) == 2 for r in rows):
            for r in rows:
                if r[0]:
                    out.setdefault(r[0].strip().lower(), r[1].strip())
    return out


def transpose_feature_table(table: dict[str, Any]) -> dict[str, dict[str, str]]:
    """Feature-comparison table (first column = feature, other columns = models) → {model: {feature: value}}."""
    headers = table.get("headers") or []
    if len(headers) < 2:
        return {}
    models = [clean_cell(h) for h in headers[1:]]
    out: dict[str, dict[str, str]] = {m: {} for m in models if m}
    for row in table.get("rows") or []:
        if not row:
            continue
        feature = clean_cell(row[0]).lower()
        for m, val in zip(models, row[1:], strict=False):
            if m:
                out[m][feature] = clean_cell(val)
    return out


_LINK = re.compile(r"\[([^\]]*)\]\([^)]*\)")
_CODE = re.compile(r"`([^`]*)`")


def clean_cell(s: str) -> str:
    import html as _html

    s = _html.unescape(s or "")
    s = _LINK.sub(r"\1", s)
    s = _CODE.sub(r"\1", s)
    s = re.sub(r"\*\*|__", "", s)
    s = re.sub(r"<[^>]+>", " ", s)
    return re.sub(r"\s+", " ", s).strip()


def link_in_cell(s: str) -> str | None:
    m = re.search(r"\]\((https?://[^)\s]+)\)", s or "")
    return m.group(1) if m else None


def parse_retirement(text: str) -> tuple[str | None, bool]:
    """'Not sooner than September 1, 2027' → ('2027-09-01', True=tentative)."""
    if not text or text.strip().lower() in ("n/a", "—", "-", ""):
        return None, False
    tentative = "not sooner" in text.lower() or "tentative" in text.lower()
    m = re.search(r"([A-Z][a-z]+ \d{1,2}, \d{4}|\d{4}-\d{2}-\d{2})", text)
    if not m:
        return None, tentative
    dt = parse_datetime(m.group(1))
    return (dt.date().isoformat() if dt else None), tentative


def money(cell: str) -> float | None:
    """'$10 / MTok', '$0.25 / MTok1' (footnote), '$12.50 per 1M tokens' → 10.0 …"""
    if not cell:
        return None
    m = re.search(r"\$\s*(\d+(?:\.\d+)?)", cell.replace(",", ""))
    if not m:
        return None
    val = float(m.group(1))
    low = cell.lower()
    if re.search(r"/\s*1?k\b|per\s*1?k\b|1,000 tokens", low):
        return val * 1000
    return val


def tokens(cell: str) -> int | None:
    from aiatlas.sdk.extract.numbers import parse_context_length

    return parse_context_length(cell or "")


def month_year(cell: str) -> str | None:
    """'Jun 2026' → '2026-06'; 'Reliable knowledge cutoff' cells."""
    if not cell:
        return None
    dt = parse_datetime(cell.strip())
    if dt and re.search(r"[A-Za-z]{3,9}\.?,? (\d{1,2},? )?\d{4}", cell):
        return f"{dt.year:04d}-{dt.month:02d}"
    if re.fullmatch(r"\d{4}-\d{2}(-\d{2})?", cell.strip()):
        return cell.strip()
    return None


def model_ref(facts: Facts, name: str, org: EntityRef, *, api_id: str | None = None, provider_key: str | None = None, family: str | None = None,
              aliases: list[str] | None = None, identifiers: dict[str, str] | None = None) -> EntityRef:
    """Official model EntityRef (identity = the lab's API id when given). `family` = the lab's own family label ("Claude", "Gemini") — when
    absent the ontology infers the versioned family ("Llama 3.1", "Qwen3") from the name; both become a `model_family` hint."""
    ids: dict[str, str] = dict(identifiers or {})
    if api_id and provider_key:
        ids[f"{provider_key}_model_id"] = api_id
    fam = EntityRef(entity_type="model_family", name=family, organization=org, identifiers={"family_key": f"{re.sub(r'[^a-z0-9.]+', '-', family.lower()).strip('-')}@{org.slug_hint}"}) if family else family_ref(name, org)
    ref = facts.entity("model", name, identifiers=ids, organization=org, aliases=aliases or [], family=fam, identity_confidence="high" if ids else "medium")
    facts.relate(org, "develops", ref)
    if family:
        facts.claim(ref, "family", family)
    elif fam is not None:
        facts.claim(ref, "family", fam.name)
    return ref


# ---------------------------------------------------------------------------------------------- canonical vocabularies for lab docs
# source label (lower-cased, punctuation-insensitive) → canonical capability slug
CAPABILITY_MAP: dict[str, str] = {
    "function calling": "function_calling", "function_calling": "function_calling", "tool calling": "function_calling", "tool_calling": "function_calling",
    "tool use": "function_calling", "tool_use": "function_calling", "tools": "function_calling", "native tool use": "function_calling",
    "structured outputs": "structured_output", "structured_outputs": "structured_output", "structured output": "structured_output", "json mode": "structured_output",
    "json_mode": "structured_output", "json output": "structured_output", "response_format": "structured_output",
    "reasoning": "reasoning", "thinking": "reasoning", "extended thinking": "reasoning", "extended_thinking": "reasoning", "adaptive thinking": "reasoning",
    "vision": "vision", "image understanding": "vision", "image input": "vision", "image_input": "vision", "image inputs": "vision",
    "audio understanding": "audio_input", "audio input": "audio_input", "audio_input": "audio_input", "speech input": "audio_input",
    "audio generation": "audio_output", "audio output": "audio_output", "audio_output": "audio_output", "speech generation": "audio_output", "text to speech": "audio_output",
    "image generation": "image_generation", "image_generation": "image_generation", "video generation": "video_generation", "video_generation": "video_generation",
    "code execution": "code_execution", "code_execution": "code_execution", "code interpreter": "code_execution", "code_interpreter": "code_execution",
    "grounding with google search": "search_grounding", "search grounding": "search_grounding", "search_grounding": "search_grounding", "web search": "search_grounding",
    "web_search": "search_grounding", "grounding with google maps": "maps_grounding",
    "caching": "caching", "prompt caching": "caching", "prompt_caching": "caching", "context caching": "caching", "context_caching": "caching",
    "batch api": "batch", "batch": "batch", "batch_api": "batch", "batch mode": "batch",
    "fine tuning": "fine_tuning", "fine-tuning": "fine_tuning", "fine_tuning": "fine_tuning", "tuning": "fine_tuning", "supervised fine-tuning": "fine_tuning",
    "streaming": "streaming", "live api": "live_api", "live_api": "live_api", "realtime": "realtime", "computer use": "computer_use", "computer_use": "computer_use",
    "file search": "file_search", "file_search": "file_search", "url context": "url_context", "url_context": "url_context", "mcp": "mcp", "distillation": "distillation",
    "predicted outputs": "predicted_outputs", "predicted_outputs": "predicted_outputs", "embeddings": "embeddings", "flex inference": "flex_inference",
    "priority inference": "priority_inference", "citations": "citations", "pdf support": "document_input", "document understanding": "document_input", "files api": "document_input",
    "web fetch": "web_fetch", "memory": "memory", "agent skills": "agent_skills", "evals": "evals",
}


def normalize_capability(raw: str) -> str:
    s = re.sub(r"\s+", " ", str(raw).strip().lower())
    return CAPABILITY_MAP.get(s) or CAPABILITY_MAP.get(s.replace("_", " ")) or CAPABILITY_MAP.get(s.replace("-", " ")) or re.sub(r"[^a-z0-9]+", "_", s).strip("_")


def normalize_capabilities(values: list[str] | None) -> list[str]:
    """Source labels / slugs → sorted canonical capability slugs (unknown labels are kept as slugs, never dropped)."""
    if not values:
        return []
    return sorted({normalize_capability(v) for v in values if str(v).strip()})


def claim_license(facts: Facts, ref: EntityRef, raw: str | None, *, weights_available: bool | None = None) -> str | None:
    """`license` = ontology key (raw label kept in `license_raw`); with `weights_available` the openness category is derived, never asserted."""
    if not raw:
        return None
    key = normalize_license(raw)
    facts.claim(ref, "license_raw", raw)
    facts.claim(ref, "license", key or raw)
    if weights_available is not None:
        if weights_available:
            facts.claim(ref, "weights_available", True)
        dims = openness_dimensions(weights_available=weights_available, license_key=key)
        facts.claim(ref, "openness", derive_openness(dims, license_key=key))
    return key


def claim_status(facts: Facts, ref: EntityRef, raw: str | None) -> str | None:
    if not raw:
        return None
    canon = normalize_status(raw)
    facts.claim(ref, "status", canon or raw)
    if canon and canon != str(raw).strip().lower():
        facts.claim(ref, "status_raw", raw)
    return canon


def claim_api_aliases(facts: Facts, ref: EntityRef, aliases: list[str] | str | None) -> None:
    """`api_aliases` is always a list; `api_alias` (the first one) is kept for compatibility."""
    items = [aliases] if isinstance(aliases, str) else list(aliases or [])
    items = [a.strip() for a in items if isinstance(a, str) and a.strip()]
    if not items:
        return
    facts.claim(ref, "api_aliases", sorted(dict.fromkeys(items)))
    facts.claim(ref, "api_alias", items[0])


def first_target_with(targets: list[Target], key: str) -> Target | None:
    return next((t for t in targets if t.key == key), None)


def iso_date(dt: datetime | None) -> str | None:
    return dt.date().isoformat() if dt else None


_FLIGHT_CHUNK = re.compile(r'self\.__next_f\.push\(\[1,"((?:[^"\\]|\\.)*)"\]\)', re.DOTALL)


def next_flight_payload(content: bytes | str) -> str:
    """Concatenate the React Server Components payload a Next.js (App Router) page streams through `self.__next_f.push([1, "…"])`.
    The chunks are JS string literals, so they decode as JSON strings. Returns '' when the page is not a Next.js RSC page."""
    import json

    html = content.decode("utf-8", errors="replace") if isinstance(content, bytes) else content
    parts: list[str] = []
    for m in _FLIGHT_CHUNK.finditer(html):
        try:
            parts.append(json.loads(f'"{m.group(1)}"'))
        except ValueError:
            continue
    return "".join(parts)


def json_after(text: str, marker: str) -> Any | None:
    """Parse the JSON object/array that starts right after `marker` (brace-matched, string-aware)."""
    import json

    i = text.find(marker)
    if i < 0:
        return None
    j = i + len(marker)
    while j < len(text) and text[j] not in "{[":
        j += 1
    if j >= len(text):
        return None
    opener = text[j]
    closer = "}" if opener == "{" else "]"
    depth = 0
    in_str = False
    k = j
    while k < len(text):
        ch = text[k]
        if in_str:
            if ch == "\\":
                k += 1
            elif ch == '"':
                in_str = False
        elif ch == '"':
            in_str = True
        elif ch == opener:
            depth += 1
        elif ch == closer:
            depth -= 1
            if depth == 0:
                break
        k += 1
    try:
        return json.loads(text[j:k + 1])
    except ValueError:
        return None


def slug_of(url: str) -> str:
    """Last non-empty path segment of a URL, without extension: '/api/docs/models/gpt-4.1.md' → 'gpt-4.1'."""
    from urllib.parse import urlparse

    path = urlparse(url).path.rstrip("/")
    seg = path.rsplit("/", 1)[-1]
    return re.sub(r"\.(md|mdx|html?)$", "", seg)


def claim_modalities(facts: Facts, ref: EntityRef, inputs: Any = None, outputs: Any = None) -> tuple[list[str], list[str]]:
    """Canonical modality lists (`pdf` → document, `Images` → image); `modalities` = union. Only what the page states."""
    mi, mo = normalize_modalities(inputs), normalize_modalities(outputs)
    if mi:
        facts.claim(ref, "modalities_input", mi)
    if mo:
        facts.claim(ref, "modalities_output", mo)
    if mi or mo:
        facts.claim(ref, "modalities", sorted(set(mi) | set(mo)))
    if "image" in mi:
        facts.claim(ref, "vision", True)
    if "audio" in mi:
        facts.claim(ref, "audio", True)
    return mi, mo


__all__ = [
    "CAPABILITY_MAP",
    "MODEL_WORDS",
    "RELEASE_WORDS",
    "announcement_events",
    "claim_api_aliases",
    "claim_license",
    "claim_modalities",
    "claim_status",
    "clean_cell",
    "first_target_with",
    "iso_date",
    "json_after",
    "kv_tables",
    "link_in_cell",
    "model_ref",
    "money",
    "month_year",
    "next_flight_payload",
    "normalize_capabilities",
    "normalize_capability",
    "parse_retirement",
    "slug_of",
    "tokens",
    "transpose_feature_table",
]