#!/usr/bin/env python3
"""
Perf-stats backup helpers.

EvalScope writes the per-benchmark cumulative perf summary into
``report.json`` under ``perf_metrics.summary`` and the raw per-sample records
into ``predictions/*.jsonl``. Both are overwritten on every run, so a
checkpoint restart resets cumulative stats and loses prior samples.

This module maintains **two durable files** under ``<output_dir>/`` so the
breakpoint-resume issue can be reconstructed:

  1. ``perf_stats_backup/<benchmark>__<model>.json``
     - The latest complete ``perf_metrics.summary`` captured right after
       each ``run_task()`` finishes.
  2. ``predictions_archive/<benchmark>__<model>.jsonl``
     - All per-sample records ever seen, deduplicated by sample ``index``.

``bash.collect_results`` prefers these archive files when aggregating metrics.
"""

import json
import shutil
from datetime import datetime
from pathlib import Path


PERF_STATS_BACKUP_DIRNAME = 'perf_stats_backup'
PREDICTIONS_ARCHIVE_DIRNAME = 'predictions_archive'
ACTIVE_TIME_DIRNAME = 'active_time'


def _safe(s: str) -> str:
    return s.replace('/', '_').replace('\\', '_').replace(' ', '_')


def get_backup_paths(output_dir: Path, benchmark: str, model_name: str):
    """Return the (perf_stats_backup, predictions_archive) paths for a benchmark/model."""
    safe = _safe(model_name)
    backup_dir = Path(output_dir) / PERF_STATS_BACKUP_DIRNAME
    archive_dir = Path(output_dir) / PREDICTIONS_ARCHIVE_DIRNAME
    perf_stats_path = backup_dir / f'{_safe(benchmark)}__{safe}.json'
    predictions_path = archive_dir / f'{_safe(benchmark)}__{safe}.jsonl'
    return perf_stats_path, predictions_path


def backup_perf_stats(output_dir: Path, benchmark: str, model_name: str,
                       report_json: Path) -> Path:
    """Copy ``perf_metrics.summary`` from a finished report to the durable backup.

    Only writes if ``report_json`` exists and contains a ``perf_metrics.summary``.
    Returns the backup path.
    """
    if not report_json or not report_json.exists():
        return None
    try:
        data = json.loads(report_json.read_text(encoding='utf-8'))
    except Exception:
        return None
    summary = (data.get('perf_metrics') or {}).get('summary')
    if not summary:
        return None

    perf_stats_path, _ = get_backup_paths(Path(output_dir), benchmark, model_name)
    perf_stats_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        'benchmark': benchmark,
        'model': model_name,
        'updated_at': datetime.now().isoformat(timespec='seconds'),
        'n_samples': summary.get('n_samples'),
        'summary': summary,
    }
    # Don't overwrite a strictly larger backup with a smaller one — that
    # would be the symptom of a reset run and we want to keep the larger.
    if perf_stats_path.exists():
        try:
            old = json.loads(perf_stats_path.read_text(encoding='utf-8'))
            old_n = old.get('n_samples') or 0
            new_n = summary.get('n_samples') or 0
            if new_n < old_n:
                # Don't clobber the larger historical backup.
                return perf_stats_path
        except Exception:
            pass
    perf_stats_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8')
    return perf_stats_path


def archive_predictions(output_dir: Path, benchmark: str, model_name: str,
                         predictions_dir: Path) -> Path:
    """Append the latest predictions into the durable archive.

    Reads every ``*.jsonl`` under ``predictions_dir``, parses each line, and
    appends each sample to the archive file (deduplicated by ``index``). If
    the same ``index`` was already archived with a different perf payload,
    the newer record wins.

    Returns the archive path.
    """
    _, archive_path = get_backup_paths(Path(output_dir), benchmark, model_name)
    archive_path.parent.mkdir(parents=True, exist_ok=True)

    if not predictions_dir or not predictions_dir.exists():
        return archive_path

    # Load existing archive (index -> record)
    existing = {}
    if archive_path.exists():
        with archive_path.open('r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                idx = obj.get('index')
                if idx is not None:
                    existing[idx] = obj

    # Add new predictions
    new_count = 0
    updated_count = 0
    for jsonl in sorted(predictions_dir.rglob('*.jsonl')):
        with jsonl.open('r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                idx = obj.get('index')
                if idx is None:
                    continue
                if idx in existing:
                    updated_count += 1
                else:
                    new_count += 1
                existing[idx] = obj

    # Write back atomically
    tmp_path = archive_path.with_suffix('.jsonl.tmp')
    with tmp_path.open('w', encoding='utf-8') as f:
        for idx in sorted(existing.keys()):
            f.write(json.dumps(existing[idx], ensure_ascii=False))
            f.write('\n')
    shutil.move(str(tmp_path), str(archive_path))

    return archive_path


def restore_from_backup(output_dir: Path, benchmark: str, model_name: str,
                        target_predictions_dir: Path, target_report_json: Path):
    """If both ``perf_stats_backup`` and ``predictions_archive`` exist for a
    benchmark/model, materialise them into ``target_predictions_dir`` and
    ``target_report_json`` before the next run starts, so evalscope resumes
    from the larger historical state instead of overwriting it.

    Returns ``True`` if a backup was restored, ``False`` otherwise.
    """
    perf_stats_path, archive_path = get_backup_paths(Path(output_dir), benchmark, model_name)
    if not perf_stats_path.exists() and not archive_path.exists():
        return False

    # Restore the archive into the predictions directory so evalscope's
    # use_cache mechanism sees the previously-completed samples.
    if archive_path.exists():
        target_predictions_dir.mkdir(parents=True, exist_ok=True)
        target_jsonl = target_predictions_dir / archive_path.name
        shutil.copy2(archive_path, target_jsonl)

    # Restore the perf summary into the report file so the cumulative stats
    # survive. We do NOT overwrite score/metrics — only the perf section.
    if perf_stats_path.exists() and target_report_json.exists():
        try:
            payload = json.loads(perf_stats_path.read_text(encoding='utf-8'))
            summary = payload.get('summary')
            if not summary:
                return True
            report = json.loads(target_report_json.read_text(encoding='utf-8'))
            report.setdefault('perf_metrics', {})['summary'] = summary
            target_report_json.write_text(
                json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8'
            )
        except Exception:
            pass

    return True


# ============================================================
# Active run-time tracking (excludes idle gaps between resumes)
# ============================================================


def get_active_time_path(output_dir: Path, benchmark: str, model_name: str) -> Path:
    """Return the path to the durable active-time record for a benchmark/model."""
    safe = _safe(model_name)
    active_dir = Path(output_dir) / ACTIVE_TIME_DIRNAME
    return active_dir / f'{_safe(benchmark)}__{safe}.json'


def record_active_time(output_dir: Path, benchmark: str, model_name: str,
                       seconds: float) -> Path:
    """Accumulate ``seconds`` of active run time for a benchmark/model.

    This is intended to measure only the time ``run_task()`` was actually
    executing, excluding idle gaps caused by manual interruption/resume.
    """
    path = get_active_time_path(Path(output_dir), benchmark, model_name)
    path.parent.mkdir(parents=True, exist_ok=True)

    data = {'total_seconds': 0.0}
    if path.exists():
        try:
            data = json.loads(path.read_text(encoding='utf-8'))
        except Exception:
            pass

    data['total_seconds'] = float(data.get('total_seconds', 0.0)) + float(seconds)
    data['updated_at'] = datetime.now().isoformat(timespec='seconds')
    data['benchmark'] = benchmark
    data['model'] = model_name

    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
    return path


def load_active_time(output_dir: Path, benchmark: str, model_name: str) -> float:
    """Return accumulated active run time in seconds, or 0.0 if no record."""
    path = get_active_time_path(Path(output_dir), benchmark, model_name)
    if not path.exists():
        return 0.0
    try:
        data = json.loads(path.read_text(encoding='utf-8'))
        return float(data.get('total_seconds', 0.0))
    except Exception:
        return 0.0