scripts/_progress.py

"""
Dynamic progress tracker for the quality-check runner.

Maintains in-memory per-task and per-template state, atomically rewrites
PROGRESS.md on every update and appends to a per-task CSV. PROGRESS.md is
designed to be safe to open at any time; it lists running status, ETA, a
worst-first per-template table, and any templates that were skipped
because not enough matching data files exist in the data pool.
"""
from __future__ import annotations

import csv
import json
import os
import shutil
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional


_TASK_CSV_FIELDS = [
    "chart_name",
    "input",
    "ok",
    "elapsed_s",
    "final_svg",
    "final_svg_bytes",
    "chart_svg",
    "chart_svg_fallback_png",
    "n_shapes",
    "n_text",
    "err",
]


def _atomic_write(path: Path, content: str) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(content, encoding="utf-8")
    tmp.replace(path)


class ProgressTracker:
    """Track quality-check progress and keep PROGRESS.md / _tasks.csv current."""

    def __init__(
        self,
        plan: Dict[str, Any],
        output_dir: Path,
        match_csv_path: Optional[Path] = None,
        min_data_for_plan: Optional[int] = None,
    ):
        self.plan = plan
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.progress_md = self.output_dir / "PROGRESS.md"
        self.tasks_csv = self.output_dir / "_tasks.csv"
        self.summary_csv = self.output_dir / "_summary.csv"

        # Build per-template plan map: chart_name -> target (number of jobs).
        self.target_per_tpl: Dict[str, int] = {}
        self.chart_type_per_tpl: Dict[str, str] = {}
        self.engine_per_tpl: Dict[str, str] = {}
        for t in plan["templates"]:
            self.target_per_tpl[t["chart_name"]] = len(t["picked_data_files"])
            self.chart_type_per_tpl[t["chart_name"]] = t["chart_type"]
            self.engine_per_tpl[t["chart_name"]] = t["engine"]

        self.total_tasks = sum(self.target_per_tpl.values())
        self.n_templates = len(self.target_per_tpl)

        # Per-task records accumulated in memory (also written incrementally).
        self.records: List[Dict[str, Any]] = []
        # Per-template rollups (computed on flush, but cached for show).
        self._rollups_cache: Optional[List[Dict[str, Any]]] = None

        self.t_start = time.time()
        self.last_flush = 0.0
        self._csv_fh = None
        self._csv_writer = None
        self._completed_keys: set = set()
        self._is_closed = False

        # Optional context: skipped templates (filtered by min-data threshold).
        self.skipped: List[Dict[str, Any]] = []
        if match_csv_path and min_data_for_plan is not None and match_csv_path.exists():
            self._load_skipped(match_csv_path, min_data_for_plan)

    def _load_skipped(self, match_csv_path: Path, min_data: int) -> None:
        engine = self.plan.get("engine", "")
        included = set(self.target_per_tpl.keys())
        with open(match_csv_path, "r") as fh:
            for r in csv.DictReader(fh):
                if engine and r["engine"] != engine:
                    continue
                if r["chart_name"] in included:
                    continue
                n = int(r["num_compatible_data"])
                if n < min_data:
                    self.skipped.append({
                        "engine": r["engine"],
                        "chart_type": r["chart_type"],
                        "chart_name": r["chart_name"],
                        "num_compatible_data": n,
                    })
        self.skipped.sort(key=lambda r: (r["num_compatible_data"], r["chart_name"]))

    def open_csv(self, resume: bool = False) -> None:
        """Open the tasks CSV. If resume=True and the file exists, load
        previous rows so we can skip them. Otherwise truncate."""
        if resume and self.tasks_csv.exists():
            with open(self.tasks_csv, "r") as fh:
                reader = csv.DictReader(fh)
                for row in reader:
                    self.records.append(row)
                    self._completed_keys.add((row["chart_name"], row["input"]))
            mode = "a"
            write_header = False
        else:
            if self.tasks_csv.exists():
                self.tasks_csv.unlink()
            mode = "w"
            write_header = True
        self._csv_fh = open(self.tasks_csv, mode, newline="", encoding="utf-8")
        self._csv_writer = csv.DictWriter(self._csv_fh, fieldnames=_TASK_CSV_FIELDS)
        if write_header:
            self._csv_writer.writeheader()
            self._csv_fh.flush()

    def already_done(self, chart_name: str, input_basename: str) -> bool:
        return (chart_name, input_basename) in self._completed_keys

    def add(self, result: Dict[str, Any]) -> None:
        """Append a result and update on-disk state."""
        self.records.append(result)
        self._completed_keys.add((result["chart_name"], result["input"]))
        if self._csv_writer:
            # CSV reader will read empty strings as is; coerce non-strings.
            self._csv_writer.writerow({k: result.get(k, "") for k in _TASK_CSV_FIELDS})
            self._csv_fh.flush()
        self._rollups_cache = None
        # Throttle PROGRESS.md updates to once every ~2s to avoid IO churn.
        now = time.time()
        if now - self.last_flush >= 2.0 or len(self.records) == self.total_tasks:
            self.flush()
            self.last_flush = now

    def _rollups(self) -> List[Dict[str, Any]]:
        if self._rollups_cache is not None:
            return self._rollups_cache
        by_tpl: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
        for r in self.records:
            by_tpl[r["chart_name"]].append(r)
        rows = []
        for chart_name, target in self.target_per_tpl.items():
            rs = by_tpl.get(chart_name, [])
            done = len(rs)
            n_ok = sum(1 for r in rs if str(r.get("ok")) == "True")
            n_fail = done - n_ok
            n_fb = sum(1 for r in rs if str(r.get("chart_svg_fallback_png")) == "True")
            n_empty = sum(
                1 for r in rs
                if str(r.get("ok")) == "True"
                and (int(r.get("n_shapes") or 0) + int(r.get("n_text") or 0)) < 8
            )
            mean_size = (
                sum(int(r.get("final_svg_bytes") or 0) for r in rs if str(r.get("ok")) == "True")
                / n_ok if n_ok else 0
            )
            mean_shapes = (
                sum(int(r.get("n_shapes") or 0) for r in rs if str(r.get("ok")) == "True")
                / n_ok if n_ok else 0
            )
            mean_t = (
                sum(float(r.get("elapsed_s") or 0) for r in rs) / done if done else 0
            )
            if done == 0:
                status = "pending"
            elif done < target:
                status = "running"
            elif n_fail > 0:
                status = "done (failures)"
            elif n_fb > 0:
                status = "done (fallback)"
            elif n_empty > 0:
                status = "done (warn)"
            else:
                status = "done"
            rows.append({
                "chart_name": chart_name,
                "chart_type": self.chart_type_per_tpl[chart_name],
                "engine": self.engine_per_tpl[chart_name],
                "target": target,
                "done": done,
                "n_success": n_ok,
                "n_fail": n_fail,
                "n_fallback_png": n_fb,
                "n_empty": n_empty,
                "mean_shapes": round(mean_shapes, 1),
                "mean_size_kb": round(mean_size / 1024, 1),
                "mean_elapsed_s": round(mean_t, 1),
                "status": status,
            })
        # Worst-first sort: failures > fallback > empty > slow > everything else.
        def _sortkey(r):
            return (
                -r["n_fail"],
                -r["n_fallback_png"],
                -r["n_empty"],
                -r["mean_elapsed_s"],
                r["chart_name"],
            )
        rows.sort(key=_sortkey)
        self._rollups_cache = rows
        return rows

    def write_summary_csv(self) -> None:
        rows = self._rollups()
        if not rows:
            return
        # Write to a tmp file then atomic-rename so a concurrent reader
        # (e.g. build_quality_preview run while the driver is still going)
        # never sees a half-written file.
        tmp = self.summary_csv.with_suffix(self.summary_csv.suffix + ".tmp")
        with open(tmp, "w", newline="", encoding="utf-8") as fh:
            w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
            w.writeheader()
            for r in rows:
                w.writerow(r)
        tmp.replace(self.summary_csv)

    def flush(self) -> None:
        rows = self._rollups()
        self.write_summary_csv()
        md = self._render_md(rows)
        _atomic_write(self.progress_md, md)

    def close(self) -> None:
        self._is_closed = True
        self.flush()
        if self._csv_fh:
            self._csv_fh.close()
            self._csv_fh = None
            self._csv_writer = None

    # ---------------------------------------------------------------- markdown
    def _render_md(self, rows: List[Dict[str, Any]]) -> str:
        done = len(self.records)
        total = self.total_tasks
        ok = sum(1 for r in self.records if str(r.get("ok")) == "True")
        fail = done - ok
        fb = sum(1 for r in self.records if str(r.get("chart_svg_fallback_png")) == "True")
        empty = sum(
            1 for r in self.records
            if str(r.get("ok")) == "True"
            and (int(r.get("n_shapes") or 0) + int(r.get("n_text") or 0)) < 8
        )
        elapsed = time.time() - self.t_start
        avg_t = (
            sum(float(r.get("elapsed_s") or 0) for r in self.records) / done
            if done else 0
        )
        # If the tracker was rebuilt from an existing CSV the wall-clock
        # elapsed is bogus (it's just the rebuild time); detect that and
        # fall back to a single-thread approximation so the reported
        # throughput / ETA don't go to infinity.
        if elapsed < max(avg_t, 1.0):
            elapsed_display = "n/a (rebuilt from CSV)"
            rate = 0.0
            eta_s = 0
        else:
            elapsed_display = f"{elapsed/60:.1f} min"
            rate = done / max(elapsed, 1e-6)
            eta_s = (total - done) / max(rate, 1e-6) if done < total else 0

        if done >= total:
            status = "finished"
        elif self._is_closed:
            status = "stopped"
        else:
            status = "running"

        lines: List[str] = []
        lines.append("# Chart Template Quality Check Progress")
        lines.append("")
        lines.append(
            f"**Last updated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} "
            f"&nbsp; &middot; &nbsp; **Status**: `{status}`"
        )
        lines.append("")
        eta_display = "-" if rate == 0 else f"{eta_s/60:.1f} min"
        lines.append(
            f"**Overall**: {done}/{total} tasks "
            f"({(done/total*100 if total else 0):.1f}%) &middot; "
            f"OK {ok} &middot; FAIL **{fail}** &middot; "
            f"fallback_png **{fb}** &middot; empty **{empty}** &middot; "
            f"elapsed {elapsed_display} &middot; "
            f"avg-task {avg_t:.1f}s &middot; "
            f"ETA {eta_display}"
        )
        lines.append("")

        # ---- summary stats by engine ----
        lines.append("## Summary")
        lines.append("")
        lines.append("| metric | value |")
        lines.append("|---|---|")
        lines.append(f"| plan templates | {self.n_templates} |")
        lines.append(f"| total tasks | {total} |")
        lines.append(f"| completed | {done} ({(done/total*100 if total else 0):.1f}%) |")
        lines.append(f"| succeeded | {ok} |")
        lines.append(f"| failed | **{fail}** |")
        lines.append(f"| fallback_png | **{fb}** |")
        lines.append(f"| empty (shapes+text<8) | **{empty}** |")
        lines.append(f"| elapsed | {elapsed_display} |")
        throughput_display = "-" if rate == 0 else f"{rate*60:.1f} tasks/min"
        lines.append(f"| throughput | {throughput_display} |")
        lines.append(f"| avg task time | {avg_t:.1f} s |")
        lines.append("")

        # ---- per-template table ----
        lines.append("## Per-template progress (worst-first)")
        lines.append("")
        lines.append(
            "| chart_name | chart_type | done/target | ok | fail | fb_png | empty | "
            "mean_shapes | mean_size | mean_t (s) | status |"
        )
        lines.append(
            "|---|---|---|---|---|---|---|---|---|---|---|"
        )
        for r in rows:
            status_md = r["status"]
            if r["n_fail"] or r["n_fallback_png"]:
                status_md = f"**{status_md}**"
            lines.append(
                f"| `{r['chart_name']}` "
                f"| {r['chart_type']} "
                f"| {r['done']}/{r['target']} "
                f"| {r['n_success']} "
                f"| {r['n_fail']} "
                f"| {r['n_fallback_png']} "
                f"| {r['n_empty']} "
                f"| {r['mean_shapes']} "
                f"| {r['mean_size_kb']} KB "
                f"| {r['mean_elapsed_s']} "
                f"| {status_md} |"
            )
        lines.append("")

        # ---- skipped templates ----
        if self.skipped:
            lines.append(f"## Templates skipped from this run ({len(self.skipped)} total)")
            lines.append("")
            lines.append(
                "These templates exist in the registry but didn't have enough matching "
                "data files in the pool to be included in this run. Their `chart_name` "
                "is shown along with how many data files in the pool matched their "
                "`requirements`. To cover them, expand the data pool (e.g. adapt more "
                "files from `/data/lizhen/resources/converted/`) and re-run "
                "`scripts/match_templates_to_data.py`."
            )
            lines.append("")
            lines.append("| chart_type | chart_name | matched data files |")
            lines.append("|---|---|---|")
            for r in self.skipped:
                lines.append(
                    f"| {r['chart_type']} | `{r['chart_name']}` "
                    f"| {r['num_compatible_data']} |"
                )
            lines.append("")
        return "\n".join(lines) + "\n"