#!/usr/bin/env python3
"""Build the dashboard's public JSON payload from a W&B project."""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
from collections import defaultdict
from datetime import datetime, timezone

_SCORE_PREFIX = "miner/"
_SCORE_SUFFIX = "/score"
_LEN_SUFFIX = "/completion_len"
_ORIG_SCORE = "original/score"
_ORIG_LEN = "original/completion_len"


def _miner_hotkey_from_score_key(key: str) -> str | None:
    if key.startswith(_SCORE_PREFIX) and key.endswith(_SCORE_SUFFIX):
        hk = key[len(_SCORE_PREFIX) : -len(_SCORE_SUFFIX)]
        return hk or None
    return None


def _as_float(value) -> float | None:
    try:
        f = float(value)
    except (TypeError, ValueError):
        return None
    return f if math.isfinite(f) else None


def extract_rankings(summary: dict) -> list[dict]:
    """Return the current miner ranking from a run summary."""
    out = []
    for key, value in summary.items():
        hk = _miner_hotkey_from_score_key(key)
        if hk is None:
            continue
        score = _as_float(value)
        if score is None:
            continue
        out.append({"hotkey": hk, "score": score})
    out.sort(key=lambda r: r["score"], reverse=True)
    return out


def extract_history(run) -> dict:
    """Return aligned per-epoch miner and base-model series for one run."""
    rows_by_epoch: dict[float, dict] = {}
    miner_keys: set[str] = set()
    saw_orig_score = saw_orig_len = False

    for row in run.scan_history():
        epoch = _as_float(row.get("epoch"))
        if epoch is None:
            continue
        bucket = rows_by_epoch.setdefault(epoch, {})
        for key, value in row.items():
            hk = _miner_hotkey_from_score_key(key)
            if hk is not None:
                bucket[("miner", hk, "score")] = _as_float(value)
                miner_keys.add(hk)
            elif key.startswith(_SCORE_PREFIX) and key.endswith(_LEN_SUFFIX):
                hk2 = key[len(_SCORE_PREFIX) : -len(_LEN_SUFFIX)]
                if hk2:
                    bucket[("miner", hk2, "completion_len")] = _as_float(value)
                    miner_keys.add(hk2)
            elif key == _ORIG_SCORE:
                bucket[("orig", "score")] = _as_float(value)
                saw_orig_score = True
            elif key == _ORIG_LEN:
                bucket[("orig", "completion_len")] = _as_float(value)
                saw_orig_len = True

    epochs = sorted(rows_by_epoch)
    miners: dict[str, dict] = {}
    for hk in miner_keys:
        miners[hk] = {
            "score": [rows_by_epoch[e].get(("miner", hk, "score")) for e in epochs],
            "completion_len": [rows_by_epoch[e].get(("miner", hk, "completion_len")) for e in epochs],
        }
    original = {}
    if saw_orig_score:
        original["score"] = [rows_by_epoch[e].get(("orig", "score")) for e in epochs]
    if saw_orig_len:
        original["completion_len"] = [rows_by_epoch[e].get(("orig", "completion_len")) for e in epochs]

    return {"epochs": [int(e) if e == int(e) else e for e in epochs], "miners": miners, "original": original}


def build_aggregate(validators: list[dict]) -> dict:
    """Average miner scores across validators that ranked each miner."""
    score_sums: dict[str, float] = defaultdict(float)
    score_counts: dict[str, int] = defaultdict(int)
    for v in validators:
        for r in v["rankings"]:
            score_sums[r["hotkey"]] += r["score"]
            score_counts[r["hotkey"]] += 1

    rankings = [
        {"hotkey": hk, "score": score_sums[hk] / score_counts[hk], "validators": score_counts[hk]}
        for hk in score_sums
    ]
    rankings.sort(key=lambda r: r["score"], reverse=True)

    history = {"epochs": [], "miners": {}, "original": {}}
    if rankings:
        top = rankings[0]["hotkey"]
        history = _mean_history(validators, top)
    return {"rankings": rankings, "history": history}


def _mean_history(validators: list[dict], top_hotkey: str) -> dict:
    """Average per-epoch history for the leading miner and base model."""
    epochs = sorted({e for v in validators for e in v["history"].get("epochs", [])})

    def mean_series(getter):
        out = []
        for e in epochs:
            vals = []
            for v in validators:
                hist = v["history"]
                if e not in hist.get("epochs", []):
                    continue
                pos = hist["epochs"].index(e)
                val = getter(hist, pos)
                if val is not None and val == val:
                    vals.append(val)
            out.append(sum(vals) / len(vals) if vals else None)
        return out

    miner_score = mean_series(lambda h, p: (h["miners"].get(top_hotkey, {}).get("score") or [None] * (p + 1))[p] if h["miners"].get(top_hotkey) else None)
    miner_len = mean_series(lambda h, p: (h["miners"].get(top_hotkey, {}).get("completion_len") or [None] * (p + 1))[p] if h["miners"].get(top_hotkey) else None)
    orig_score = mean_series(lambda h, p: (h["original"].get("score") or [None] * (p + 1))[p] if h["original"].get("score") else None)
    orig_len = mean_series(lambda h, p: (h["original"].get("completion_len") or [None] * (p + 1))[p] if h["original"].get("completion_len") else None)

    original = {}
    if any(x is not None for x in orig_score):
        original["score"] = orig_score
    if any(x is not None for x in orig_len):
        original["completion_len"] = orig_len
    return {
        "epochs": epochs,
        "miners": {top_hotkey: {"score": miner_score, "completion_len": miner_len}},
        "original": original,
    }


def build(project: str, entity: str | None, api=None) -> dict:
    if api is None:
        import wandb

        api = wandb.Api()
    path = f"{entity}/{project}" if entity else project

    validators: list[dict] = []
    seen: set[str] = set()
    for run in api.runs(path, order="-created_at"):
        hotkey = dict(run.config).get("validator_hotkey")
        if not isinstance(hotkey, str) or not hotkey or hotkey in seen:
            continue
        seen.add(hotkey)
        validators.append(
            {
                "hotkey": hotkey,
                "rankings": extract_rankings(dict(run.summary)),
                "history": extract_history(run),
            }
        )

    return {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "validators": validators,
        "aggregate": build_aggregate(validators),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate the static dashboard's data.json from wandb.")
    parser.add_argument("--wandb-project", required=True, help="Shared wandb project every validator logs into")
    parser.add_argument("--wandb-entity", default=None, help="Wandb entity/team (default: the API key's default entity)")
    parser.add_argument("--out", default="data.json", help="Output path for the generated JSON")
    args = parser.parse_args(argv)

    key = os.environ.get("WANDB_KEY") or os.environ.get("WANDB_API_KEY")
    if not key:
        print("error: set WANDB_KEY (or WANDB_API_KEY) in the environment", file=sys.stderr)
        return 2
    os.environ["WANDB_API_KEY"] = key

    data = build(args.wandb_project, args.wandb_entity)
    with open(args.out, "w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2)
        fh.write("\n")
    print(f"wrote {args.out}: {len(data['validators'])} validators, "
          f"{len(data['aggregate']['rankings'])} ranked miners")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())