#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from benchmark_progress_mark_local import (
    DEFAULT_INFERENCE_MODE,
    DEFAULT_MAX_FRAMES,
    DEFAULT_MIN_FRAMES,
    DEFAULT_MODEL_PATH,
    DEFAULT_PREFIX_BATCH_SIZE,
    DEFAULT_PREFIX_SAMPLE_FRAMES,
    RobometerLocalRunner,
    build_frame_retry_schedule,
    is_cuda_oom_error,
    load_all_video_frames,
    sample_video_frames_with_indices,
    write_overview_plot,
)


def save_results(output_dir: Path, payload: dict) -> None:
    with (output_dir / "results.json").open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
        f.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run Robometer on a single video with the same benchmark-style inference logic, without progress marks."
    )
    parser.add_argument("--video", required=True, help="Path to the input video")
    parser.add_argument("--task", required=True, help="Task instruction")
    parser.add_argument("--model-path", default=str(DEFAULT_MODEL_PATH))
    parser.add_argument("--fps", type=float, default=3.0)
    parser.add_argument("--max-frames", type=int, default=DEFAULT_MAX_FRAMES)
    parser.add_argument("--min-frames", type=int, default=DEFAULT_MIN_FRAMES)
    parser.add_argument(
        "--inference-mode",
        choices=["frame_steps", "whole"],
        default=DEFAULT_INFERENCE_MODE,
        help="frame_steps matches benchmark behavior; whole does one forward pass on the full sampled trajectory",
    )
    parser.add_argument("--prefix-sample-frames", type=int, default=DEFAULT_PREFIX_SAMPLE_FRAMES)
    parser.add_argument("--prefix-batch-size", type=int, default=DEFAULT_PREFIX_BATCH_SIZE)
    parser.add_argument(
        "--adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_true",
        help="On CUDA OOM in whole mode, retry with smaller frame budgets",
    )
    parser.add_argument(
        "--no-adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_false",
        help="Disable frame-budget retry on CUDA OOM",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path(__file__).parent / "outputs_benchmark_style",
        help="Root directory for outputs",
    )
    parser.set_defaults(adaptive_max_frames=True)
    args = parser.parse_args()

    video_path = Path(args.video).expanduser().resolve()
    if not video_path.exists():
        raise FileNotFoundError(f"Video not found: {video_path}")

    output_dir = args.output_root / video_path.stem
    output_dir.mkdir(parents=True, exist_ok=True)

    runner = RobometerLocalRunner(model_path=str(args.model_path))

    print(f"[RUN] {video_path.name}")
    print(f"[TASK] {args.task}")
    print(f"[VIDEO] {video_path}")
    print(f"[OUT] {output_dir}")

    all_frames, native_fps = load_all_video_frames(video_path)
    total_frames = len(all_frames)
    retry_schedule = (
        build_frame_retry_schedule(args.max_frames, args.min_frames, bool(args.adaptive_max_frames))
        if args.inference_mode == "whole"
        else [int(args.max_frames)]
    )

    progress_pred = None
    success_probs = None
    sampled_indices = None
    used_max_frames = retry_schedule[0]

    for attempt_idx, frame_budget in enumerate(retry_schedule, start=1):
        frames, sampled_indices = sample_video_frames_with_indices(
            all_frames,
            native_fps=native_fps,
            fps=float(args.fps),
            max_frames=int(frame_budget),
            required_frames=[],
        )
        print(
            f"Loaded {total_frames} total frames; sampled {len(frames)} frames at fps={float(args.fps):g} "
            f"(max_frames={int(frame_budget)}, try {attempt_idx}/{len(retry_schedule)})"
        )
        try:
            progress_pred, success_probs = runner.compute_rewards_per_frame(
                video_frames=frames,
                task=args.task,
                inference_mode=args.inference_mode,
                prefix_sample_frames=int(args.prefix_sample_frames),
                prefix_batch_size=int(args.prefix_batch_size),
            )
            used_max_frames = int(frame_budget)
            break
        except RuntimeError as exc:
            if args.inference_mode != "whole" or not is_cuda_oom_error(exc) or attempt_idx == len(retry_schedule):
                raise
            next_budget = retry_schedule[attempt_idx]
            print(
                f"[OOM] whole inference hit CUDA OOM at max_frames={int(frame_budget)}; "
                f"retrying with max_frames={int(next_budget)}"
            )
            runner.reload_model()

    if progress_pred is None or success_probs is None or sampled_indices is None:
        raise RuntimeError("Failed to compute Robometer outputs.")
    if progress_pred.size == 0:
        raise RuntimeError("Robometer returned empty progress predictions.")
    if progress_pred.size != len(sampled_indices):
        raise RuntimeError(
            f"Progress length mismatch: got {progress_pred.size} predictions for {len(sampled_indices)} sampled frames"
        )

    np.save(str(output_dir / "progress.npy"), progress_pred)
    np.save(str(output_dir / "success_probs.npy"), success_probs)

    payload = {
        "video": str(video_path),
        "instruction": args.task,
        "num_total_frames": total_frames,
        "native_fps": native_fps,
        "sample_fps": float(args.fps),
        "num_sampled_frames": len(sampled_indices),
        "sampled_original_frame_indices": sampled_indices,
        "max_frames_used": used_max_frames,
        "inference_mode": args.inference_mode,
        "prefix_sample_frames": int(args.prefix_sample_frames),
        "prefix_batch_size": int(args.prefix_batch_size),
        "progress_min": float(np.min(progress_pred)),
        "progress_max": float(np.max(progress_pred)),
        "progress_mean": float(np.mean(progress_pred)),
    }
    save_results(output_dir, payload)
    write_overview_plot(output_dir, sampled_indices, progress_pred, [])
    print(f"[OK] Saved results to {output_dir}")


if __name__ == "__main__":
    main()