#!/usr/bin/env python3
"""End-to-end offline test for bash/run.py + collect_results + perf_backup.

This script does NOT call any model API. It uses an existing benchmark
directory (aime24/seed_42) and walks through every code path that the real
``run.py`` would touch:

1. backup_perf_stats — confirm a perf summary snapshot is written
2. archive_predictions — confirm per-sample records are deduplicated into
   the archive
3. restore_from_backup — confirm a fresh work_dir can be rehydrated from
   the backup files
4. collect_results aggregation — confirm correct metrics are produced
5. Whitelist mode — confirm only the requested benchmarks appear in the
   summary
6. hle_low alias — confirm hle_low directory maps to canonical "hle"
7. Reset recovery — confirm that even after wiping the predictions file
   and the report summary, the archive+backup still recover the metrics

Run from the project root:

    python bash/test_e2e.py
"""

import json
import shutil
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import collect_results as cr  # noqa: E402
import perf_backup as pb  # noqa: E402


PROJECT_ROOT = Path('/data1/sora/evalscope')
SOURCE_OUTPUT = PROJECT_ROOT / 'output'


def banner(msg):
    print('\n' + '=' * 60)
    print(f' {msg}')
    print('=' * 60)


def assert_close(actual, expected, name, tol=1e-3):
    if abs(actual - expected) > tol:
        raise AssertionError(f'{name}: expected ~{expected}, got {actual}')
    print(f'  OK  {name}={actual:.4f}')


def main():
    banner('Setting up isolated test output dir')
    test_root = PROJECT_ROOT / 'output_e2e_test'
    if test_root.exists():
        shutil.rmtree(test_root)
    test_root.mkdir(parents=True)

    # Copy the canonical aime24 data into the isolated dir
    src = SOURCE_OUTPUT / 'aime24' / 'seed_42'
    dst = test_root / 'aime24' / 'seed_42'
    shutil.copytree(src, dst)

    # And copy hle_low as hle_low under the test dir (so we exercise the
    # alias code path).
    hle_src = SOURCE_OUTPUT / 'hle_low' / 'seed_42'
    hle_dst = test_root / 'hle_low' / 'seed_42'
    if hle_src.exists():
        shutil.copytree(hle_src, hle_dst)

    benchmark = 'aime24'
    model_name = 'DeepSeek-V4-Flash-Int8'

    banner('1) backup_perf_stats — snapshot summary')
    report_json = dst / 'reports' / model_name / f'{benchmark}.json'
    backup_path = pb.backup_perf_stats(test_root, benchmark, model_name, report_json)
    assert backup_path.exists(), 'backup file was not created'
    payload = json.loads(backup_path.read_text())
    assert payload['n_samples'] == 30, f'expected 30 samples, got {payload["n_samples"]}'
    print(f'  OK  backup written: {backup_path}, n_samples={payload["n_samples"]}')

    banner('2) archive_predictions — deduplicate into archive')
    preds = dst / 'predictions' / model_name
    archive_path = pb.archive_predictions(test_root, benchmark, model_name, preds)
    assert archive_path.exists(), 'archive file was not created'
    n_lines = sum(1 for line in archive_path.read_text().splitlines() if line.strip())
    assert n_lines == 30, f'expected 30 archived samples, got {n_lines}'
    print(f'  OK  archive written: {archive_path}, samples={n_lines}')

    banner('3) restore_from_backup — rehydrate a fresh work_dir')
    fresh_dir = test_root / 'fresh_work' / 'seed_42'
    fresh_report = fresh_dir / 'reports' / model_name / f'{benchmark}.json'
    fresh_preds = fresh_dir / 'predictions' / model_name
    # Pre-create empty structure to mimic what eval_task would do
    fresh_report.parent.mkdir(parents=True, exist_ok=True)
    fresh_report.write_text(json.dumps({
        'name': f'{model_name}@{benchmark}',
        'perf_metrics': {'summary': {'n_samples': 1}},
    }))
    restored = pb.restore_from_backup(
        test_root, benchmark, model_name, fresh_preds, fresh_report,
    )
    assert restored, 'restore_from_backup returned False'
    restored_report = json.loads(fresh_report.read_text())
    assert restored_report['perf_metrics']['summary']['n_samples'] == 30, \
        f'restored n_samples={restored_report["perf_metrics"]["summary"]["n_samples"]}'
    print('  OK  restored report n_samples=30 (matches backup)')

    banner('4) collect_results — aggregate metrics')
    csv, xlsx = cr.eval_benchmark([benchmark], test_root, model_name)
    import pandas as pd
    df = pd.read_csv(csv)
    aime24 = df[df['Benchmark'] == 'aime24'].iloc[0]
    # Single seed_42 run, so values match the report directly (no multi-run
    # averaging). The full-summary CSV averages across run_1..run_11.
    assert_close(aime24['得分'], 0.6333, 'aime24 score (single seed)')
    assert_close(aime24['实测时间(h)'], 0.1469, 'aime24 duration (h)')
    assert_close(aime24['总样本数'], 30, 'aime24 sample count')
    assert_close(aime24['延迟_mean(s)'], 35.24971, 'aime24 latency mean')
    assert_close(aime24['TTFT_mean(s)'], 0.26301, 'aime24 TTFT mean')
    assert_close(aime24['TPOT_mean(s)'], 0.02165, 'aime24 TPOT mean')
    assert_close(aime24['输入tokens_mean'], 119.33, 'aime24 input_tokens_mean')
    assert_close(aime24['输出tokens_mean'], 1609.8, 'aime24 output_tokens_mean')

    banner('5) Whitelist mode — only the requested benchmarks appear')
    csv2, _ = cr.eval_benchmark([benchmark], test_root, model_name)
    df2 = pd.read_csv(csv2)
    assert df2['Benchmark'].dropna().tolist() == ['aime24'], \
        f'whitelist did not limit benchmarks: {df2["Benchmark"].tolist()}'
    print('  OK  summary only contains aime24 + total')

    banner('6) hle_low → hle alias')
    if hle_src.exists():
        csv3, _ = cr.eval_benchmark(['hle'], test_root, model_name)
        df3 = pd.read_csv(csv3)
        assert 'hle' in df3['Benchmark'].tolist(), \
            f'hle canonical name missing from: {df3["Benchmark"].tolist()}'
        assert 'hle_low' not in df3['Benchmark'].tolist(), \
            f'hle_low should be aliased away: {df3["Benchmark"].tolist()}'
        print('  OK  hle_low directory maps to canonical "hle"')

    banner('7) Reset recovery — wipe report+predictions, expect archive to restore')
    # Wipe the current predictions
    shutil.rmtree(preds)
    # Reset report.json
    data = json.loads(report_json.read_text())
    data['perf_metrics']['summary']['n_samples'] = 1
    data['perf_metrics']['summary']['latency']['mean'] = 0.0
    report_json.write_text(json.dumps(data, indent=2))

    csv4, _ = cr.eval_benchmark([benchmark], test_root, model_name)
    df4 = pd.read_csv(csv4)
    aime24_after = df4[df4['Benchmark'] == 'aime24'].iloc[0]
    assert_close(aime24_after['得分'], 0.6333, 'aime24 score (post-reset)')
    assert_close(aime24_after['总样本数'], 30, 'aime24 sample count (post-reset)')
    assert_close(aime24_after['延迟_mean(s)'], 35.24971, 'aime24 latency mean (post-reset)')
    assert_close(aime24_after['TTFT_mean(s)'], 0.26301, 'aime24 TTFT mean (post-reset)')
    print('  OK  archive fully recovered metrics after predictions wipe')

    banner('8) Backup monotonicity — new backup must not overwrite a larger one')
    # We have a backup with n_samples=30. Now write a fresh report with
    # n_samples=2 and call backup_perf_stats again — the backup should NOT
    # be overwritten.
    data = json.loads(report_json.read_text())
    data['perf_metrics']['summary']['n_samples'] = 2
    report_json.write_text(json.dumps(data, indent=2))
    pb.backup_perf_stats(test_root, benchmark, model_name, report_json)
    payload2 = json.loads(backup_path.read_text())
    assert payload2['n_samples'] == 30, \
        f'backup was clobbered: n_samples now {payload2["n_samples"]}'
    print('  OK  backup preserved n_samples=30 even after a reset run')

    banner('All checks passed')
    print(f'Test artifacts under {test_root} (kept for inspection)')


if __name__ == '__main__':
    main()