"""Full COMBO pipeline: train -> export -> upload to HuggingFace Hub.

Evaluation happens during training (per-epoch) and results for the best epoch
are saved to results_best.json automatically.

Runs each step as a subprocess, passing the same config arguments throughout.
"""

import argparse
import subprocess
import sys
import time
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

from combo.config import load_config


STEPS = ["train", "export", "upload"]


def run_step(name: str, cmd: list[str], dry_run: bool = False) -> None:
    """Run a pipeline step as a subprocess."""
    header = f"{'=' * 60}\n  PIPELINE STEP: {name.upper()}\n{'=' * 60}"
    print(f"\n{header}")
    print(f"Command: {' '.join(cmd)}\n")

    if dry_run:
        print("  [dry-run] Skipping execution.\n")
        return

    start = time.time()
    result = subprocess.run(cmd)
    elapsed = time.time() - start

    if result.returncode != 0:
        print(f"\nPIPELINE FAILED at step '{name}' (exit code {result.returncode})")
        print(f"  Elapsed: {elapsed:.1f}s")
        sys.exit(result.returncode)

    print(f"\nStep '{name}' completed in {elapsed:.1f}s")


def main():
    parser = argparse.ArgumentParser(
        description="Run the full COMBO pipeline: train -> export -> upload.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--config", type=Path, default=Path("config/base.yaml"),
        help="Path to base configuration file (default: config/base.yaml).",
    )
    parser.add_argument(
        "--task-config", type=Path, required=True,
        help="Path to task-specific configuration override.",
    )
    parser.add_argument(
        "--start-from", choices=STEPS, default="train",
        help="Start the pipeline from this step (default: train).",
    )
    parser.add_argument(
        "--stop-after", choices=STEPS, default="upload",
        help="Stop the pipeline after this step (default: upload).",
    )
    parser.add_argument(
        "--checkpoint", type=str, default="checkpoint_best.pt",
        help="Checkpoint filename for export step (default: checkpoint_best.pt).",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Print commands without executing them.",
    )
    args = parser.parse_args()

    # Determine which steps to run
    start_idx = STEPS.index(args.start_from)
    stop_idx = STEPS.index(args.stop_after)
    steps_to_run = STEPS[start_idx : stop_idx + 1]

    # Load config to display summary
    config = load_config(args.config, args.task_config)

    print("=" * 60)
    print("  COMBO FULL PIPELINE")
    print("=" * 60)
    print(f"  Config:       {args.config}")
    print(f"  Task config:  {args.task_config}")
    print(f"  Model name:   {config.model_name}")
    print(f"  Repo ID:      {config.repo_id}")
    print(f"  Language:      {config.data.language}")
    print(f"  Treebanks:    {', '.join(config.data.treebanks)}")
    print(f"  Base model:   {config.model.base_model}")
    print(f"  Output dir:   {config.language_output_dir}")
    print(f"  Steps:        {' -> '.join(steps_to_run)}")
    if args.dry_run:
        print(f"  Mode:         DRY RUN")
    print()

    # Build shared config args
    config_args = ["--config", str(args.config), "--task-config", str(args.task_config)]

    # --- Step 1: Train ---
    if "train" in steps_to_run:
        cmd = [sys.executable, "-m", "combo.cli.train"] + config_args
        run_step("train", cmd, dry_run=args.dry_run)

    # --- Step 2: Export (+ optional upload) ---
    if "export" in steps_to_run or "upload" in steps_to_run:
        push = "upload" in steps_to_run
        cmd = [
            sys.executable, "-m", "combo.cli.export",
            *config_args,
            "--checkpoint", args.checkpoint,
        ]
        if push:
            cmd.append("--push-to-hub")
        step_name = "export + upload" if push else "export"
        run_step(step_name, cmd, dry_run=args.dry_run)

    # --- Summary ---
    print("\n" + "=" * 60)
    print("  PIPELINE COMPLETE")
    print("=" * 60)
    if "upload" in steps_to_run and not args.dry_run:
        print(f"  Model uploaded to: https://huggingface.co/{config.repo_id}")
    print()


if __name__ == "__main__":
    main()
