#!/usr/bin/env python3
"""One-click pipeline: prepare -> train -> export -> evaluate -> plot loss."""

import argparse
import subprocess
import sys
from pathlib import Path
import json

from platform_utils import safe_path, is_windows


def run(cmd):
    print("running:", " ".join(cmd))
    # 使用safe_path确保跨平台兼容性
    cmd = [safe_path(arg) for arg in cmd]

    # 在Windows上可能需要特殊处理
    if is_windows():
        subprocess.run(cmd, check=True, shell=False)
    else:
        subprocess.run(cmd, check=True)


def parse_args():
    parser = argparse.ArgumentParser(description="Run full HAI pipeline.")
    base_dir = Path(__file__).resolve().parent
    parser.add_argument(
        "--config",
        default=str(base_dir / "config.json"),
        help="Path to training config JSON",
    )
    parser.add_argument(
        "--device",
        default="auto",
        help="cpu, cuda, or auto (used for export_samples.py)",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    base_dir = Path(__file__).resolve().parent
    config_path = Path(args.config)
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)
    timesteps = cfg.get("timesteps", 200)
    seq_len = cfg.get("sample_seq_len", cfg.get("seq_len", 64))
    batch_size = cfg.get("sample_batch_size", cfg.get("batch_size", 2))
    clip_k = cfg.get("clip_k", 5.0)
    run([sys.executable, str(base_dir / "prepare_data.py")])
    run([sys.executable, str(base_dir / "train.py"), "--config", args.config, "--device", args.device])
    run(
        [
            sys.executable,
            str(base_dir / "export_samples.py"),
            "--include-time",
            "--device",
            args.device,
            "--config",
            str(config_path),
            "--timesteps",
            str(timesteps),
            "--seq-len",
            str(seq_len),
            "--batch-size",
            str(batch_size),
            "--clip-k",
            str(clip_k),
            "--use-ema",
        ]
    )
    run([sys.executable, str(base_dir / "evaluate_generated.py")])
    run([sys.executable, str(base_dir / "plot_loss.py")])


if __name__ == "__main__":
    main()