"""Refresh the coverage Space from the live MIC fleet.

Run from a clone of this Space:

    python tools/update.py --tci ../transformers-ci --mic ../transformers-mic [--python PY] [--upload]

--tci     transformers-ci checkout (its contracts.coverage builds the fleet coverage)
--mic     transformers-mic checkout (sweep_sources.json, and tmic for the shrink pages)
--python  interpreter for the shrink pages; its transformers must know every fleet model type
--rerender  render every shrink page, not only those whose fixture changed (after a tmic shrink_view change)
--upload  push data.json and shrink/ to huggingface-mic/mic-coverage
"""
import argparse
import datetime
import json
import sys
from collections import Counter
from pathlib import Path

from huggingface_hub import HfApi

HERE = Path(__file__).resolve().parent
SPACE = HERE.parent
REPO = "huggingface-mic/mic-coverage"
GOAL_SHARE = 0.3

parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
parser.add_argument("--tci", type=Path, required=True)
parser.add_argument("--mic", type=Path, required=True)
parser.add_argument("--python", default=sys.executable)
parser.add_argument("--rerender", action="store_true")
parser.add_argument("--upload", action="store_true")
args = parser.parse_args()

sys.path[:0] = [str(args.tci / "src"), str(HERE)]
from transformersci.contracts import coverage  # noqa: E402

import gated  # noqa: E402
import nockpt  # noqa: E402
import results  # noqa: E402
import shrink  # noqa: E402

cov = coverage.build()
data = json.loads((SPACE / "data.json").read_text())
covered, shrinkable = set(cov["covered"]), set(cov["sweep"]["shrinkable"])
for row in data["rows"]:
    if row[0] in covered:
        row[1] = "Done"
    elif row[0] in shrinkable and row[1] != "Done":
        row[1] = "Reachable"
nockpt.apply(data, args.mic / "experiments/results/sweep_sources.json")
overrides = json.loads((HERE / "overrides.json").read_text())
for row in data["rows"]:
    if row[0] in overrides and row[0] not in covered:
        row[1], row[2] = overrides[row[0]]

fixtures = coverage.fleet_fixtures(cov["fleet"]["revision"])[1]
pages = SPACE / "shrink"
todo = [f"{repo}@{revision}" for repo, revision in fixtures] if args.rerender else shrink.stale(fixtures, pages)
print(f"rendering {len(todo)} shrink page(s)")
for error in shrink.render(todo, pages, args.mic, args.python):
    print("  failed", error)
data["results"], data["meta"]["results"] = results.fetch()
shrink.link_back(pages, data["results"])
fleets = {}
for name in ("prototypes", "control"):  # control wins for a repo in both
    base = f"https://huggingface.co/datasets/{coverage.fleet.REPOSITORY}/resolve/{cov['fleet']['revision']}"
    for row in coverage.fleet._rows("contracts", coverage._get(f"{base}/{name}/contracts.csv")):
        fleets[row["mic_repo"]] = name
data["shrink"] = shrink.links(fixtures, pages, fleets)
data["gated"] = gated.gating(data["rows"], fixtures)
data["meta"].update(
    generated=datetime.date.today().isoformat(), fleet=cov["fleet"]["revision"], sweep=cov["sweep"]["revision"]
)
(SPACE / "data.json").write_text(json.dumps(data))

rows = sorted(data["rows"], key=lambda r: -r[3])
goal = rows[: round(len(rows) * GOAL_SHARE)]
print(cov["contracts"], "contracts; all:", dict(Counter(r[1] for r in data["rows"])), "goal:", dict(Counter(r[1] for r in goal)))
if args.upload:
    commit = HfApi().upload_folder(
        repo_id=REPO, repo_type="space", folder_path=SPACE, allow_patterns=["data.json", "shrink/*"],
        commit_message=f"Coverage: fleet {cov['fleet']['revision'][:8]}, {cov['contracts']} contracts",
    )
    print(commit.commit_url)