"""Public leaderboards → append-only `benchmark_results` (tier 2). One connector class per leaderboard domain (the SDK binds one
`source_key` per connector, and each board has its own rate limit and provenance) — module kept as `leaderboards`.

  * aider_leaderboard     raw polyglot_leaderboard.yml from the aider repository → `aider-polyglot` (pass_rate_2 + well-formed rate).
  * swebench_leaderboard  https://www.swebench.com/ embeds `<script id="leaderboard-data">` (5 boards: Verified, Lite, Test, Multimodal,
                          Multilingual) → results keyed by (model, benchmark, config {system, system_org, …}).
  * livebench_leaderboard https://livebench.ai/ is a React app; its bundle lists the release dates and fetches `./table_<date>.csv` +
                          `./categories_<date>.json` (plus a model metadata map with display names / HF links). The connector follows
                          index → bundle → categories → table and emits category averages + the overall mean (LiveBench's own aggregation).
  * artificial_analysis   https://artificialanalysis.ai/leaderboards/models — Next.js RSC flight payload (`self.__next_f.push`) with the full
                          model list: Intelligence Index (config {version}) + component evaluations mapped to registry benchmarks.

Model names on leaderboards are free text → `EntityRef("model", name, aliases=[…])`; the resolver merges by normalized alias when the
model exists and otherwise creates a low-tier entity.
"""
from __future__ import annotations

import csv
import io
import json
import re
from datetime import UTC, date, datetime
from statistics import fmean
from typing import Any

import yaml

from aiatlas.registry import load, org_by_hf, org_ref, organizations
from aiatlas.sdk.connector import BaseConnector, Parsed, RunContext
from aiatlas.sdk.extract.dates import parse_datetime
from aiatlas.sdk.facts import EntityRef, Facts, Target
from aiatlas.sdk.fetch import FetchResult


def benchmark_ref(key: str) -> EntityRef | None:
    b = next((b for b in load("benchmarks") if b["key"] == key), None)
    if not b:
        return None
    return EntityRef(entity_type="benchmark", name=b["name"], identifiers={"registry_benchmark": key}, aliases=list(b.get("aliases", [])), slug_hint=key)


def _slug(s: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", s.lower())


def org_from_name(facts: Facts, *names: str | None) -> EntityRef | None:
    """Registry organization by hf_org / key / name / alias (punctuation-insensitive: 'Z AI' == 'Z.ai'); models map to the lab, not the holding."""
    for n in names:
        if not n:
            continue
        low = n.strip().lower()
        known = org_by_hf(low)
        key = known["key"] if known else None
        if not key:
            wanted = _slug(low)
            for k, o in organizations().items():
                if wanted in {_slug(k), _slug(o["name"]), *(_slug(a) for a in o.get("aliases", []))}:
                    key = k
                    break
        if key:
            if key == "meta" and "meta-ai" in organizations():
                key = "meta-ai"
            ref = org_ref(key)
            facts.entities.append(ref)
            return ref
    return None


def _dt(v: Any) -> datetime | None:
    if isinstance(v, datetime):
        return v if v.tzinfo else v.replace(tzinfo=UTC)
    if isinstance(v, date):
        return datetime(v.year, v.month, v.day, tzinfo=UTC)
    return parse_datetime(v) if isinstance(v, str) else None


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


# ================================================================================================ Aider polyglot
class AiderLeaderboardConnector(_Leaderboard):
    name = "aider_leaderboard"
    label = "Aider polyglot coding leaderboard"
    description = "polyglot_leaderboard.yml (raw GitHub file of the aider repo): pass rate, well-formed rate, cost, edit format per run."
    source_key = "aider.chat"
    rate_per_min = 10
    expected_min_records = 30
    URL = "https://raw.githubusercontent.com/Aider-AI/aider/main/aider/website/_data/polyglot_leaderboard.yml"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url=self.URL, doc_type="leaderboard", key="aider", min_bytes=1000, priority=1, meta={"content_type": "text/yaml"})]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        bench = benchmark_ref("aider-polyglot")
        if bench is None:
            return facts
        facts.entities.append(bench)
        try:
            rows = yaml.safe_load(res.text) or []
        except yaml.YAMLError:
            return facts
        for row in rows:
            if not isinstance(row, dict) or not row.get("model") or row.get("pass_rate_2") is None:
                continue
            name = str(row["model"]).strip()
            # the `--model` id is NOT an alias: many display-name variants share one API id, which would make every alias lookup
            # ambiguous (the resolver then creates a new entity per run). It is kept in the result config instead.
            ref = facts.entity("model", name[:200])
            evaluated = _dt(row.get("date"))
            config = {k: row.get(k) for k in ("edit_format", "command", "test_cases", "versions", "dirname") if row.get(k) is not None}
            m = re.search(r"--model\s+(\S+)", str(row.get("command") or ""))
            if m:
                config["api_model_id"] = m.group(1)
            config["date"] = evaluated.date().isoformat() if evaluated else None
            if row.get("total_cost") is not None:
                config["total_cost_usd"] = row["total_cost"]
            if row.get("seconds_per_case") is not None:
                config["seconds_per_case"] = row["seconds_per_case"]
            facts.result(model=ref, benchmark=bench, score=float(row["pass_rate_2"]), metric="pass_rate_2", unit="%", config=config, evaluated_at=evaluated,
                         source_url="https://aider.chat/docs/leaderboards/")
            if row.get("percent_cases_well_formed") is not None:
                facts.result(model=ref, benchmark=bench, score=float(row["percent_cases_well_formed"]), metric="percent_cases_well_formed", unit="%",
                             config=config, evaluated_at=evaluated, source_url="https://aider.chat/docs/leaderboards/")
        facts.document_entity = bench
        facts.document_title = "Aider polyglot leaderboard data"
        return facts


# ================================================================================================ SWE-bench
BOARD_KEYS = {"verified": "swe-bench-verified", "lite": "swe-bench-lite", "test": "swe-bench-full", "full": "swe-bench-full",
              "multimodal": "swe-bench-multimodal", "multilingual": "swe-bench-multilingual"}


class SweBenchLeaderboardConnector(_Leaderboard):
    name = "swebench_leaderboard"
    label = "SWE-bench leaderboards (Verified, Lite, Full, Multimodal, Multilingual)"
    description = "Embedded leaderboard JSON of swebench.com: % resolved per (system, model) submission with dates and cost."
    source_key = "swebench.com"
    rate_per_min = 10
    expected_min_records = 50

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url="https://www.swebench.com/", doc_type="leaderboard", key="swebench", min_bytes=20000, priority=1)]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        html = parsed.html
        if not html:
            return facts
        data = html.embedded_json.get("leaderboard-data")
        if not isinstance(data, list):
            return facts
        for board in data:
            key = BOARD_KEYS.get(str(board.get("name") or "").lower())
            bench = benchmark_ref(key) if key else None
            if bench is None:
                continue
            facts.entities.append(bench)
            for r in board.get("results") or []:
                self._row(facts, bench, board.get("name"), r)
        facts.document_title = html.title or "SWE-bench leaderboards"
        return facts

    def _row(self, facts: Facts, bench: EntityRef, board: str, r: dict[str, Any]) -> None:
        model_name = (r.get("model_display") or "").strip()
        resolved = r.get("resolved")
        if not model_name or not isinstance(resolved, (int, float)):
            return
        tags = [t for t in r.get("tags") or [] if isinstance(t, str)]
        model_tags = [t.split(":", 1)[1].strip() for t in tags if t.lower().startswith("model:")]
        org = org_from_name(facts, r.get("model_org"))
        ref = facts.entity("model", model_name[:200], organization=org)     # SWE-bench's "Model:" tag goes to config, not aliases (see aider)
        evaluated = _dt(r.get("date"))
        config = {"board": board, "system": r.get("agent") or r.get("name"), "system_org": r.get("agent_org"), "open_source_system": r.get("os_system"),
                  "checked_by_swebench": r.get("checked"), "submission": r.get("folder"), "date": evaluated.date().isoformat() if evaluated else None}
        if model_tags:
            config["model_tag"] = model_tags[0]
        if r.get("reasoning_effort"):
            config["reasoning_effort"] = r["reasoning_effort"]
        if isinstance(r.get("cost"), (int, float)):
            config["total_cost_usd"] = r["cost"]
        if isinstance(r.get("instance_cost"), (int, float)):
            config["cost_per_instance_usd"] = r["instance_cost"]
        facts.result(model=ref, benchmark=bench, score=float(resolved), metric="resolved", unit="%", config={k: v for k, v in config.items() if v is not None},
                     evaluated_at=evaluated, source_url="https://www.swebench.com/")


# ================================================================================================ LiveBench
RELEASES = re.compile(r'\[("20\d\d-\d\d-\d\d"(?:,"20\d\d-\d\d-\d\d")*)\]')
MODEL_META = re.compile(r'"([\w.:/-]+)":\{((?:[^{}]|\{[^{}]*\})*?displayName:"(?:[^{}]|\{[^{}]*\})*)\}')
META_FIELD = re.compile(r'(url|huggingface|organization|displayName):"([^"]*)"')


class LiveBenchLeaderboardConnector(_Leaderboard):
    name = "livebench_leaderboard"
    label = "LiveBench leaderboard"
    description = "Latest LiveBench release table (CSV + categories JSON discovered from the site bundle): category averages and overall mean."
    source_key = "livebench.ai"
    rate_per_min = 10
    expected_min_records = 20
    BASE = "https://livebench.ai/"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url=self.BASE, doc_type="listing", key="livebench_index", min_bytes=200, priority=1)]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        key = target.key or ""
        if key == "livebench_index" and parsed.html:
            script = next((s.attributes.get("src") for s in parsed.html.css("script[src]") if "main." in (s.attributes.get("src") or "")), None)
            if script:
                facts.follow(self.BASE + script.lstrip("./"), doc_type="script", key="livebench_js", min_bytes=10000, priority=1, meta={"content_type": "text/javascript"})
        elif key == "livebench_js":
            js = res.text
            dates = sorted({d for m in RELEASES.finditer(js) for d in json.loads("[" + m.group(1) + "]")})
            if not dates:
                ctx.log.warning("livebench bundle: no release list found")
                return facts
            latest = self.config.get("release") or dates[-1]
            models = {}
            for m in MODEL_META.finditer(js):
                fields = dict(META_FIELD.findall(m.group(2)))
                models[m.group(1)] = {"url": fields.get("url"), "huggingface": fields.get("huggingface"), "organization": fields.get("organization"),
                                      "display_name": fields.get("displayName")}
            tag = latest.replace("-", "_")
            facts.follow(f"{self.BASE}categories_{tag}.json", doc_type="data", key="livebench_categories", min_bytes=20, priority=1,
                         meta={"release": latest, "models": models, "content_type": "application/json"})
        elif key == "livebench_categories":
            cats = parsed.json if parsed.kind == "json" else None
            release = target.meta.get("release")
            if isinstance(cats, dict) and release:
                facts.follow(f"{self.BASE}table_{release.replace('-', '_')}.csv", doc_type="leaderboard", key="livebench_table", min_bytes=200, priority=1,
                             meta={"release": release, "categories": cats, "models": target.meta.get("models") or {}, "content_type": "text/csv"})
        elif key == "livebench_table":
            self._table(facts, target, res.text)
        return facts

    def _table(self, facts: Facts, target: Target, text: str) -> None:
        bench = benchmark_ref("livebench")
        if bench is None:
            return
        facts.entities.append(bench)
        release = target.meta.get("release")
        cats: dict[str, list[str]] = target.meta.get("categories") or {}
        meta: dict[str, dict[str, Any]] = target.meta.get("models") or {}
        evaluated = _dt(release)
        for row in csv.DictReader(io.StringIO(text)):
            mid = (row.get("model") or "").strip()
            if not mid:
                continue
            info = meta.get(mid) or {}
            name = info.get("display_name") or mid
            # no `hf_repo` identifier here: several LiveBench ids (effort variants, MLX conversions) share one display name and would
            # glue unrelated hub repositories onto a single entity; the link is kept as a claim on the leaderboard row's model
            org = org_from_name(facts, info.get("organization"))
            ref = facts.entity("model", name[:200], organization=org, aliases=[mid])
            hf = info.get("huggingface") or ""
            m = re.search(r"huggingface\.co/([\w.-]+/[\w.-]+)", hf)
            if m:
                facts.claim(ref, "livebench_hf_link", f"https://huggingface.co/{m.group(1)}")
            scores = {k: float(v) for k, v in row.items() if k and k != "model" and v not in (None, "", "-") and _is_number(v)}
            cat_avgs: dict[str, float] = {}
            for cat, tasks in cats.items():
                vals = [scores[t] for t in tasks if t in scores]
                if vals:
                    cat_avgs[cat] = round(fmean(vals), 3)
            if not cat_avgs:
                continue
            base_cfg = {"release": release, "livebench_model_id": mid, "aggregation": "mean of category averages; category = mean of its subtasks"}
            facts.result(model=ref, benchmark=bench, score=round(fmean(cat_avgs.values()), 3), metric="global_average", unit="%", config=base_cfg,
                         evaluated_at=evaluated, source_url=self.BASE)
            for cat, avg in cat_avgs.items():
                facts.result(model=ref, benchmark=bench, score=avg, metric=f"category:{cat}", unit="%", config={"release": release, "livebench_model_id": mid, "subtasks": cats[cat]},
                             evaluated_at=evaluated, source_url=self.BASE)
        facts.document_entity = bench
        facts.document_title = f"LiveBench {release}"


def _is_number(v: str) -> bool:
    try:
        float(v)
        return True
    except ValueError:
        return False


# ================================================================================================ Artificial Analysis
AA_EVALS = {"gpqa": ("gpqa", "accuracy", "GPQA Diamond"), "hle": ("humanitys-last-exam", "accuracy", None), "aime25": ("aime-2025", "accuracy", None),
            "livecodebench": ("livecodebench", "pass@1", None), "scicode": ("scicode", "accuracy", None), "ifbench": ("ifbench", "accuracy", None),
            "mmmuPro": ("mmmu-pro", "accuracy", None), "tau2": ("tau2-bench", "pass^1", None), "terminalbenchV40": ("terminal-bench", "accuracy", "v4.0"),
            "terminalbenchV21": ("terminal-bench", "accuracy", "v2.1"), "terminalbenchHard": ("terminal-bench", "accuracy", "hard")}
PUSH = re.compile(r'self\.__next_f\.push\(\[1,"(.*?)"\]\)\s*</script>', re.DOTALL)


class ArtificialAnalysisConnector(_Leaderboard):
    name = "artificial_analysis"
    label = "Artificial Analysis — Intelligence Index and component evaluations"
    description = "Model leaderboard page (RSC payload): Intelligence Index with version, component evaluations, openness and release dates."
    source_key = "artificialanalysis.ai"
    rate_per_min = 6
    expected_min_records = 100
    URL = "https://artificialanalysis.ai/leaderboards/models"

    async def discover(self, ctx: RunContext) -> list[Target]:
        return [Target(url=self.URL, doc_type="leaderboard", key="aa_models", min_bytes=50000, priority=1)]

    async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:
        facts = Facts()
        models, version = parse_rsc_models(res.text)
        if not models:
            ctx.log.warning("artificial analysis: no model objects in the RSC payload")
            return facts
        index = benchmark_ref("artificial-analysis-intelligence-index")
        assert index
        facts.entities.append(index)
        benches = {k: benchmark_ref(k) for k, *_ in AA_EVALS.values()}
        for slug, m in models.items():
            name = (m.get("name") or slug).strip()
            creator = m.get("creator") or {}
            creator_name = creator.get("name") or m.get("modelCreatorName")
            org = org_from_name(facts, creator.get("slug"), creator_name)
            if org is None and creator_name:
                org = facts.entity("company", creator_name, identifiers={"artificial_analysis_creator": creator.get("slug") or _slug(creator_name)})
            ref = facts.entity("model", name[:200], identifiers={"artificial_analysis": slug}, organization=org,
                               aliases=[a for a in {slug, m.get("shortName")} if a and a != name])
            if m.get("releaseDate"):
                facts.claim(ref, "release_date", str(m["releaseDate"])[:10])
            if isinstance(m.get("isOpenWeights"), bool):
                facts.claim(ref, "openness", "open-weights" if m["isOpenWeights"] else "proprietary")
            if isinstance(m.get("isReasoning"), bool):
                facts.claim(ref, "reasoning", m["isReasoning"])
            if m.get("deprecated") is True:
                facts.claim(ref, "status", "deprecated")
            if isinstance(m.get("contextWindowTokens"), int) and m["contextWindowTokens"] > 0:
                facts.claim(ref, "context_length", m["contextWindowTokens"], unit="tokens")
            if isinstance(m.get("medianOutputTokensPerSecond"), (int, float)):
                facts.claim(ref, "metric.aa_median_output_tokens_per_second", round(float(m["medianOutputTokensPerSecond"]), 1))
            ii = m.get("intelligenceIndex")
            if isinstance(ii, (int, float)):
                cfg = {"version": version, "estimated": bool(m.get("intelligenceIndexIsEstimated")), "aa_slug": slug}
                facts.result(model=ref, benchmark=index, score=round(float(ii), 2), metric="index", unit="", config=cfg, source_url=self.URL)
            for field, (bkey, metric, variant) in AA_EVALS.items():
                v = m.get(field)
                bench = benches.get(bkey)
                if bench is None or not isinstance(v, (int, float)):
                    continue
                if bench not in facts.entities:
                    facts.entities.append(bench)
                cfg = {"evaluator": "Artificial Analysis", "index_version": version, "aa_slug": slug}
                if variant:
                    cfg["variant"] = variant
                facts.result(model=ref, benchmark=bench, score=round(float(v) * 100, 2), metric=metric, unit="%", config=cfg, source_url=self.URL)
        facts.document_entity = index
        facts.document_title = f"Artificial Analysis models — Intelligence Index v{version}" if version else "Artificial Analysis models"
        return facts


def parse_rsc_models(html: str) -> tuple[dict[str, dict[str, Any]], str | None]:
    """Concatenate the Next.js flight chunks, parse each `id:json` line and collect every object with `slug` + `intelligenceIndex`."""
    chunks = []
    for p in PUSH.findall(html):
        try:
            chunks.append(json.loads('"' + p + '"'))
        except json.JSONDecodeError:
            continue
    blob = "".join(chunks)
    version_m = re.search(r"Intelligence Index v(\d+(?:\.\d+)?)", blob)
    models: dict[str, dict[str, Any]] = {}
    for line in blob.split("\n"):
        _, sep, payload = line.partition(":")
        if not sep or "intelligenceIndex" not in payload:
            continue
        try:
            data = json.loads(payload)
        except json.JSONDecodeError:
            continue
        stack = [data]
        while stack:
            o = stack.pop()
            if isinstance(o, dict):
                if "intelligenceIndex" in o and isinstance(o.get("slug"), str):
                    models.setdefault(o["slug"], o)
                else:
                    stack.extend(o.values())
            elif isinstance(o, list):
                stack.extend(o)
    return models, (version_m.group(1) if version_m else None)


CONNECTORS = [AiderLeaderboardConnector, SweBenchLeaderboardConnector, LiveBenchLeaderboardConnector, ArtificialAnalysisConnector]