from __future__ import annotations

import argparse
from datetime import datetime
from pathlib import Path

import imageio.v2 as imageio
from src.policy_utils import select_action
from src.runtime import configure_runtime_env

configure_runtime_env()

from stable_baselines3 import PPO

from src.env import get_action_set, make_mario_env, reset_compat, step_compat
from src.utils import ensure_artifact_paths, latest_model_path, seed_everything


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Record Mario agent rollout to mp4 (headless).")
    parser.add_argument("--model-path", type=str, default="", help="Path to .zip model. Empty means latest model.")
    parser.add_argument("--env-id", type=str, default="SuperMarioBros-1-1-v0")
    parser.add_argument("--movement", type=str, default="auto", choices=["auto", "right_only", "simple"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--frame-skip", type=int, default=4)
    parser.add_argument("--reward-mode", type=str, default="raw", choices=["raw", "clip", "progress"])
    parser.add_argument("--progress-scale", type=float, default=0.02)
    parser.add_argument("--death-penalty", type=float, default=-50.0)
    parser.add_argument("--flag-bonus", type=float, default=100.0)
    parser.add_argument("--stall-penalty", type=float, default=0.05)
    parser.add_argument("--stall-steps", type=int, default=40)
    parser.add_argument("--backward-penalty-scale", type=float, default=0.01)
    parser.add_argument("--milestone-interval", type=int, default=32)
    parser.add_argument("--milestone-bonus", type=float, default=1.0)
    parser.add_argument("--no-progress-terminate-steps", type=int, default=120)
    parser.add_argument("--no-progress-terminate-penalty", type=float, default=20.0)
    parser.add_argument(
        "--time-penalty",
        type=float,
        default=0.0,
        help="Per-step living cost added to reward (<=0.0). 0.0 disables.",
    )
    parser.add_argument(
        "--hard-stuck-steps",
        type=int,
        default=0,
        help="Consecutive near-stationary steps before truncation. 0 disables.",
    )
    parser.add_argument(
        "--hard-stuck-epsilon",
        type=float,
        default=1.0,
        help="Treat |x_pos delta| < epsilon as no movement for hard-stuck detection.",
    )
    parser.add_argument(
        "--hard-stuck-penalty",
        type=float,
        default=5.0,
        help="Extra penalty applied when hard-stuck truncation triggers.",
    )
    parser.add_argument("--clip-reward", action="store_true")
    parser.add_argument("--fps", type=int, default=30)
    parser.add_argument("--duration-sec", type=int, default=10)
    parser.add_argument("--max-steps", type=int, default=5_000)
    parser.add_argument("--stochastic", action="store_true")
    parser.add_argument(
        "--epsilon",
        type=float,
        default=0.0,
        help="Epsilon-greedy probability for deterministic policy. Ignored when --stochastic is set.",
    )
    parser.add_argument(
        "--epsilon-random-mode",
        type=str,
        default="uniform",
        choices=["uniform", "policy"],
        help="Random action source for epsilon-greedy: uniform=action_space.sample, policy=model stochastic sample.",
    )
    parser.add_argument("--output", type=str, default="", help="Output mp4 path.")
    return parser.parse_args()


def _action_count_for_movement(movement: str) -> int:
    return len(get_action_set(movement))


def _model_action_count(model: PPO):
    if hasattr(model, "action_space") and hasattr(model.action_space, "n"):
        return int(model.action_space.n)
    action_net = getattr(getattr(model, "policy", None), "action_net", None)
    out_features = getattr(action_net, "out_features", None)
    if out_features is not None:
        return int(out_features)
    return None


def resolve_movement(movement_arg: str, model: PPO) -> str:
    model_action_n = _model_action_count(model)

    if movement_arg == "auto":
        if model_action_n is not None:
            for candidate in ("right_only", "simple"):
                if _action_count_for_movement(candidate) == model_action_n:
                    return candidate
        return "right_only"

    if model_action_n is not None:
        expected_n = _action_count_for_movement(movement_arg)
        if expected_n != model_action_n:
            raise ValueError(
                f"movement='{movement_arg}' has {expected_n} actions, but model expects {model_action_n}. "
                "Use --movement auto or pass the matching movement."
            )
    return movement_arg


def resolve_model(user_path: str) -> Path:
    if user_path:
        p = Path(user_path).expanduser().resolve()
        if not p.exists():
            raise FileNotFoundError(f"Model not found: {p}")
        return p

    paths = ensure_artifact_paths()
    latest = latest_model_path(paths.models)
    if latest is None:
        raise FileNotFoundError("No model found under artifacts/models. Please run training first.")
    return latest


def resolve_output_path(user_output: str) -> Path:
    if user_output:
        return Path(user_output).expanduser().resolve()

    paths = ensure_artifact_paths()
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    return (paths.videos / f"mario_replay_{ts}.mp4").resolve()


def save_frames_fallback(frames, output_path: Path) -> Path:
    frame_dir = output_path.with_suffix("")
    frame_dir.mkdir(parents=True, exist_ok=True)
    for i, frame in enumerate(frames):
        imageio.imwrite(frame_dir / f"frame_{i:06d}.png", frame)
    return frame_dir


def main() -> None:
    args = parse_args()
    if args.epsilon < 0.0 or args.epsilon > 1.0:
        raise ValueError(f"--epsilon must be in [0, 1], got {args.epsilon}")
    seed_everything(args.seed)
    reward_mode = "clip" if args.clip_reward and args.reward_mode == "raw" else args.reward_mode

    model_path = resolve_model(args.model_path)
    output_path = resolve_output_path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    model = PPO.load(str(model_path))
    print(f"[video] model={model_path}")
    movement = resolve_movement(args.movement, model)
    print(
        f"[video] movement={movement} reward_mode={reward_mode} "
        f"time_penalty={args.time_penalty} hard_stuck_steps={args.hard_stuck_steps} "
        f"hard_stuck_epsilon={args.hard_stuck_epsilon} hard_stuck_penalty={args.hard_stuck_penalty} "
        f"epsilon={args.epsilon} epsilon_random_mode={args.epsilon_random_mode}"
    )
    env = make_mario_env(
        env_id=args.env_id,
        seed=args.seed,
        movement=movement,
        reward_mode=reward_mode,
        clip_reward=args.clip_reward,
        frame_skip=args.frame_skip,
        render_mode="rgb_array",
        progress_scale=args.progress_scale,
        death_penalty=args.death_penalty,
        flag_bonus=args.flag_bonus,
        stall_penalty=args.stall_penalty,
        stall_steps=args.stall_steps,
        backward_penalty_scale=args.backward_penalty_scale,
        milestone_interval=args.milestone_interval,
        milestone_bonus=args.milestone_bonus,
        no_progress_terminate_steps=args.no_progress_terminate_steps,
        no_progress_terminate_penalty=args.no_progress_terminate_penalty,
        time_penalty=args.time_penalty,
        hard_stuck_steps=args.hard_stuck_steps,
        hard_stuck_epsilon=args.hard_stuck_epsilon,
        hard_stuck_penalty=args.hard_stuck_penalty,
    )

    obs, _ = reset_compat(env, seed=args.seed)
    frames = []

    first_frame = env.render()
    if first_frame is not None:
        # nes-py may reuse the same frame buffer; copy to avoid aliasing all frames.
        frames.append(first_frame.copy())

    target_frames = max(1, args.fps * args.duration_sec)
    done = False
    step_count = 0
    reward_sum = 0.0
    episode_count = 1

    while len(frames) < target_frames and step_count < args.max_steps:
        action = select_action(
            model=model,
            obs=obs,
            deterministic=not args.stochastic,
            epsilon=args.epsilon,
            epsilon_random_mode=args.epsilon_random_mode,
            env=env,
        )
        obs, reward, terminated, truncated, info = step_compat(env, action)
        reward_sum += float(reward)
        step_count += 1
        done = terminated or truncated

        frame = env.render()
        if frame is not None:
            frames.append(frame.copy())

        if done:
            obs, _ = reset_compat(env, seed=args.seed + episode_count)
            episode_count += 1
            frame = env.render()
            if frame is not None and len(frames) < target_frames:
                frames.append(frame.copy())
            done = False

    if not frames:
        raise RuntimeError("No frames captured. Check render_mode support and environment setup.")

    try:
        writer = imageio.get_writer(str(output_path), fps=args.fps, codec="libx264", quality=8)
        for frame in frames:
            writer.append_data(frame)
        writer.close()
        print(f"[video] Saved mp4: {output_path}")
    except Exception as exc:
        frame_dir = save_frames_fallback(frames, output_path)
        print(f"[warn] mp4 write failed: {exc}")
        print(f"[fallback] Saved frame sequence: {frame_dir}")
        print(
            "[fallback] Convert frames to mp4 with: "
            f"ffmpeg -framerate {args.fps} -i {frame_dir}/frame_%06d.png -c:v libx264 -pix_fmt yuv420p {output_path}"
        )

    print(
        f"[stats] frames={len(frames)} approx_sec={len(frames)/max(args.fps,1):.2f} "
        f"steps={step_count} reward_sum={reward_sum:.2f} episodes={episode_count}"
    )
    env.close()


if __name__ == "__main__":
    main()
