"""Load per-unit ``progress.json`` files from GCS for the local /train dashboard."""

from __future__ import annotations

import json
import re
from typing import Any

DIR_RE = re.compile(r"^(.+)_seed(\d+)_fold(\d+)$")


def parse_gs_uri(uri: str) -> tuple[str, str]:
    u = uri.strip().rstrip("/")
    if not u.startswith("gs://"):
        raise ValueError("EXCEL_CHUNKER_GCS_EXPERIMENTS_URI must start with gs://")
    rest = u[5:]
    slash = rest.find("/")
    if slash == -1:
        raise ValueError("URI must include a prefix path after the bucket name")
    bucket = rest[:slash]
    prefix = rest[slash + 1 :].lstrip("/")
    if not prefix:
        raise ValueError("URI prefix must not be empty (point at .../experiments)")
    if not prefix.endswith("/"):
        prefix = prefix + "/"
    return bucket, prefix


def fetch_suite_from_gcs(uri: str) -> dict[str, Any]:
    """Return a dict shaped like ``SuiteProgress.to_api_dict()`` from blob ``*/progress.json``."""
    try:
        from google.cloud import storage
    except ImportError as e:
        raise RuntimeError(
            "Install google-cloud-storage (e.g. pip install google-cloud-storage)"
        ) from e

    bucket_name, prefix = parse_gs_uri(uri)
    client = storage.Client()
    bucket = client.bucket(bucket_name)
    units: list[dict[str, Any]] = []
    for blob in bucket.list_blobs(prefix=prefix):
        if not blob.name.endswith("/progress.json"):
            continue
        try:
            raw = blob.download_as_text(timeout=60)
            row = json.loads(raw)
        except Exception:
            continue
        if not isinstance(row, dict):
            continue
        parent = blob.name.rsplit("/", 2)[-2] if "/" in blob.name else ""
        m = DIR_RE.match(parent)
        try:
            seed_v = int(row["seed"]) if row.get("seed") is not None else (int(m.group(2)) if m else 0)
            fold_v = int(row["fold"]) if row.get("fold") is not None else (int(m.group(3)) if m else 0)
        except (TypeError, ValueError):
            seed_v, fold_v = 0, 0
        units.append(
            {
                "arch": row.get("arch") or (m.group(1) if m else ""),
                "seed": seed_v,
                "fold": fold_v,
                "run_tag": row.get("run_tag") or parent,
                "status": row.get("status") or "unknown",
                "epoch": int(row.get("epoch") or 0),
                "epochs": int(row.get("epochs") or 0),
                "loss": row.get("loss"),
                "val_acc": row.get("val_acc"),
                "best_acc": row.get("best_acc"),
                "message": row.get("message") or "",
                "error": row.get("error"),
                "updated_at": row.get("updated_at"),
            }
        )

    def _sort_key(u: dict) -> tuple:
        return (u.get("seed", 0), u.get("fold", 0), u.get("arch", ""), u.get("run_tag", ""))

    # Stable fallback order (seed / fold / arch).
    units.sort(key=_sort_key)
    # Newest worker activity first so a fresh Vertex run is visible at the top.
    units.sort(key=lambda u: u.get("updated_at") or "", reverse=True)
    done_like = {"done", "skipped", "error"}
    completed = sum(1 for u in units if u.get("status") in done_like)
    running = any(u.get("status") == "running" for u in units)
    return {
        "running": running,
        "started_at": None,
        "finished_at": None,
        "total_units": len(units),
        "completed_units": completed,
        "units": units,
        "log_tail": [],
        "source": "gcs",
        "gcs_uri": uri,
    }