"""
replace_with_real.py - atomically swap CALIBRATED_PREVIEW figures and tables
with the real outputs from a run of `run_all.py`.

Usage (on the user's machine after a real run):

    python replace_with_real.py output_20260501_120000

The script:
  1. Reads `output_demo/MANIFEST.json` to know what to swap.
  2. For every CALIBRATED_PREVIEW item whose target file exists in the
     supplied real-run directory, copies it over and flips the manifest
     entry to status='REAL', recording the source path and timestamp.
  3. Regenerates `output_demo/.real_data_flag.tex` (a one-line file
     read by main.tex) so that LaTeX captions stop showing the
     "[CALIBRATED PREVIEW]" stamp once all 13 core items are real.
  4. Prints a clear before/after summary.

Idempotent: rerunning is safe. Items already marked REAL are skipped.
"""

from __future__ import annotations

import argparse
import datetime as _dt
import json
import shutil
import sys
from pathlib import Path


HERE = Path(__file__).resolve().parent
DEMO = HERE / "output_demo"
MANIFEST = DEMO / "MANIFEST.json"


def load_manifest() -> dict:
    if not MANIFEST.exists():
        sys.exit(f"ERROR: manifest not found at {MANIFEST}")
    return json.loads(MANIFEST.read_text())


def save_manifest(m: dict) -> None:
    MANIFEST.write_text(json.dumps(m, indent=2))


def update_latex_flag(m: dict) -> None:
    """Write a tiny .tex file consumed by main.tex. The flag is true iff
    every CORE item (fig1..fig13) is REAL. fig14/fig15 are optional."""
    core_items = [k for k in m["items"] if k.startswith(("fig1_", "fig2_", "fig3_",
                                                          "fig4_", "fig5_", "fig6_",
                                                          "fig7_", "fig8_", "fig9_",
                                                          "fig10", "fig11", "fig12",
                                                          "fig13"))]
    all_real = all(m["items"][k]["status"] == "REAL" for k in core_items)
    flag_path = DEMO / ".real_data_flag.tex"
    flag_path.write_text(
        f"% auto-generated by replace_with_real.py at "
        f"{_dt.datetime.now().isoformat(timespec='seconds')}\n"
        f"\\providecommand\\realdataavailable{{}}\n"
        f"\\renewcommand\\realdataavailable{{{'true' if all_real else 'false'}}}\n"
    )
    print(f"  LaTeX flag written: realdataavailable = "
          f"{'true' if all_real else 'false'}  ({flag_path})")


def replace(real_run_dir: Path, *, dry_run: bool = False) -> None:
    if not real_run_dir.is_dir():
        sys.exit(f"ERROR: real-run directory not found: {real_run_dir}")
    real_fig_dir = real_run_dir / "figures"
    real_tab_dir = real_run_dir / "tables"
    if not real_fig_dir.is_dir() or not real_tab_dir.is_dir():
        sys.exit(f"ERROR: real-run directory must contain figures/ and tables/ "
                 f"(got: {sorted(p.name for p in real_run_dir.iterdir())})")

    m = load_manifest()
    n_swapped = 0
    n_skipped_already_real = 0
    n_missing = 0
    print(f"Source: {real_run_dir}")
    print(f"Target: {DEMO}")
    print("-" * 70)

    for key, item in m["items"].items():
        if item["status"] == "REAL":
            n_skipped_already_real += 1
            continue
        # Resolve source paths relative to the real run dir
        src_fig = real_run_dir / item["path"]
        srcs_csv = [real_run_dir / p for p in item.get("csv_paths", [])]
        # All required source files must exist
        if not src_fig.exists() or not all(s.exists() for s in srcs_csv):
            missing = [str(p.relative_to(real_run_dir)) for p in [src_fig, *srcs_csv]
                        if not p.exists()]
            print(f"  SKIP {key:<28}  missing: {', '.join(missing)}")
            n_missing += 1
            continue
        # Atomic swap: copy to .tmp, then rename
        dst_fig = HERE / item["path"]
        dst_fig.parent.mkdir(parents=True, exist_ok=True)
        if not dry_run:
            tmp = dst_fig.with_suffix(dst_fig.suffix + ".tmp")
            shutil.copy2(src_fig, tmp)
            tmp.replace(dst_fig)
            for src_csv, rel in zip(srcs_csv, item.get("csv_paths", [])):
                dst_csv = HERE / rel
                dst_csv.parent.mkdir(parents=True, exist_ok=True)
                tmp = dst_csv.with_suffix(dst_csv.suffix + ".tmp")
                shutil.copy2(src_csv, tmp)
                tmp.replace(dst_csv)
        item["status"] = "REAL"
        item["replaced_from"] = str(real_run_dir)
        item["replaced_at"] = _dt.datetime.now().isoformat(timespec="seconds")
        n_swapped += 1
        print(f"  {'(dry) ' if dry_run else ''}OK   {key:<28}  -> {item['path']}")

    if not dry_run:
        save_manifest(m)
        update_latex_flag(m)

    print("-" * 70)
    print(f"  swapped:   {n_swapped}")
    print(f"  skipped (already real): {n_skipped_already_real}")
    print(f"  missing in source dir:  {n_missing}")
    if n_swapped > 0 and not dry_run:
        print()
        print("  Next: rebuild the paper by running:")
        print(f"       cd paper && pdflatex main.tex && bibtex main && pdflatex main.tex && pdflatex main.tex")


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("real_run_dir", type=str,
                    help="Path to a real run directory produced by run_all.py "
                         "(must contain figures/ and tables/ subdirectories).")
    p.add_argument("--dry-run", action="store_true",
                    help="Show what would be swapped without modifying any files.")
    p.add_argument("--reset-to-preview", action="store_true",
                    help=argparse.SUPPRESS)  # internal: revert manifest to all-preview
    args = p.parse_args()

    if args.reset_to_preview:
        m = load_manifest()
        for k, v in m["items"].items():
            if k != "fig3_symmetry_error_bar" and k != "fig6_batched_speedup":
                # fig14 stays as TODO_LOCAL
                if v.get("status") != "TODO_LOCAL":
                    v["status"] = "CALIBRATED_PREVIEW"
                v.pop("replaced_from", None)
                v.pop("replaced_at", None)
        save_manifest(m)
        update_latex_flag(m)
        print("Manifest reset to CALIBRATED_PREVIEW for non-real items.")
        return

    replace(Path(args.real_run_dir).resolve(), dry_run=args.dry_run)


if __name__ == "__main__":
    main()