from __future__ import annotations

import argparse
import shutil
from datetime import datetime
from pathlib import Path

from src.runtime import configure_runtime_env

configure_runtime_env()

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback, CallbackList, CheckpointCallback
from stable_baselines3.common.save_util import load_from_zip_file
from stable_baselines3.common.vec_env import DummyVecEnv, VecMonitor
from torch.utils.tensorboard import SummaryWriter

from src.env import make_env_fn
from src.utils import (
    collect_torch_runtime_diagnostics,
    ensure_artifact_paths,
    resolve_torch_device,
    seed_everything,
    write_latest_pointer,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train PPO agent on NES Super Mario Bros.")
    parser.add_argument("--env-id", type=str, default="SuperMarioBros-1-1-v0")
    parser.add_argument("--movement", type=str, default="right_only", choices=["right_only", "simple"])
    parser.add_argument("--total-timesteps", type=int, default=1_000_000)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--n-envs", type=int, default=4)
    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("--learning-rate", type=float, default=2.5e-4)
    parser.add_argument("--n-steps", type=int, default=128)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--n-epochs", type=int, default=4)
    parser.add_argument("--gamma", type=float, default=0.99)
    parser.add_argument("--gae-lambda", type=float, default=0.95)
    parser.add_argument("--ent-coef", type=float, default=0.01)
    parser.add_argument("--clip-range", type=float, default=0.2)

    parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint frequency in env steps.")
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        help="Torch device: auto | cpu | mps | cuda | cuda:0 ...",
    )
    parser.add_argument(
        "--allow-device-fallback",
        action="store_true",
        help="Allow explicit --device requests (cuda/mps) to fallback to CPU when unavailable.",
    )
    parser.add_argument("--clip-reward", action="store_true", help="Enable reward clipping to {-1,0,1}.")
    parser.add_argument("--run-name", type=str, default="", help="Optional custom run name.")
    parser.add_argument(
        "--init-model-path",
        type=str,
        default="",
        help="Optional .zip model path to initialize weights before training.",
    )
    parser.add_argument(
        "--allow-partial-init",
        action="store_true",
        help=(
            "Allow partial init when checkpoint is from a different action space "
            "(e.g. right_only -> simple). Compatible policy layers will be loaded, "
            "incompatible layers (like action head) are skipped."
        ),
    )
    parser.add_argument("--progress-bar", action="store_true", help="Enable SB3 progress bar (requires tqdm/rich).")
    return parser.parse_args()


def load_partial_policy_weights(model: PPO, init_model_path: Path, device: str) -> tuple[int, int]:
    _, params, _ = load_from_zip_file(str(init_model_path), device=device)
    if params is None or "policy" not in params:
        raise RuntimeError(f"No policy parameters found in checkpoint: {init_model_path}")

    src_state = params["policy"]
    dst_state = model.policy.state_dict()

    compatible_state = {}
    skipped_count = 0
    for key, value in src_state.items():
        if key in dst_state and dst_state[key].shape == value.shape:
            compatible_state[key] = value
        else:
            skipped_count += 1

    if not compatible_state:
        raise RuntimeError("No compatible policy tensors found for partial init.")

    model.policy.load_state_dict(compatible_state, strict=False)
    return len(compatible_state), skipped_count


class EpisodeEndLoggingCallback(BaseCallback):
    """Log per-episode terminal diagnostics to stdout and TensorBoard."""

    REASON_TO_CODE = {"death": 0, "no_progress": 1, "timeout": 2, "clear": 3, "hard_stuck": 4}

    def __init__(self, n_envs: int, tb_log_dir: Path):
        super().__init__(verbose=0)
        self.n_envs = n_envs
        self.tb_log_dir = tb_log_dir
        self.tb_writer: SummaryWriter | None = None
        self.episode_count = 0
        self.episode_max_x_pos = [0.0 for _ in range(max(n_envs, 1))]

    @staticmethod
    def _resolve_done_reason(info: dict) -> str:
        done_reason = str(info.get("done_reason", "")).strip().lower()
        if done_reason in EpisodeEndLoggingCallback.REASON_TO_CODE:
            return done_reason
        if bool(info.get("flag_get", False)):
            return "clear"
        if bool(info.get("terminated_by_hard_stuck", False)):
            return "hard_stuck"
        if bool(info.get("terminated_by_stall", False)):
            return "no_progress"
        if bool(info.get("TimeLimit.truncated", False)):
            return "timeout"
        try:
            if float(info.get("time", 1.0)) <= 0.0:
                return "timeout"
        except (TypeError, ValueError):
            pass
        return "death"

    def _on_training_start(self) -> None:
        self.tb_log_dir.mkdir(parents=True, exist_ok=True)
        self.tb_writer = SummaryWriter(log_dir=str(self.tb_log_dir))

    def _on_step(self) -> bool:
        infos = self.locals.get("infos")
        dones = self.locals.get("dones")
        if infos is None or dones is None:
            return True

        for env_idx, (info, done) in enumerate(zip(infos, dones)):
            if env_idx >= len(self.episode_max_x_pos):
                continue
            x_pos = float(info.get("x_pos", 0.0))
            if x_pos > self.episode_max_x_pos[env_idx]:
                self.episode_max_x_pos[env_idx] = x_pos

            if not bool(done):
                continue

            self.episode_count += 1
            max_x_pos = self.episode_max_x_pos[env_idx]
            flag_get = 1.0 if bool(info.get("flag_get", False)) else 0.0
            done_reason = self._resolve_done_reason(info)
            done_reason_code = float(self.REASON_TO_CODE[done_reason])
            episode_step = self.episode_count

            self.logger.record_mean("rollout/episode_max_x_pos", max_x_pos)
            self.logger.record_mean("rollout/flag_get", flag_get)
            self.logger.record_mean(f"rollout/done_reason_{done_reason}", 1.0)

            if self.tb_writer is not None:
                self.tb_writer.add_scalar("episode_end/episode_max_x_pos", max_x_pos, episode_step)
                self.tb_writer.add_scalar("episode_end/flag_get", flag_get, episode_step)
                self.tb_writer.add_scalar("episode_end/done_reason_code", done_reason_code, episode_step)
                self.tb_writer.add_scalar("episode_end/done_reason_death", 1.0 if done_reason == "death" else 0.0, episode_step)
                self.tb_writer.add_scalar(
                    "episode_end/done_reason_no_progress", 1.0 if done_reason == "no_progress" else 0.0, episode_step
                )
                self.tb_writer.add_scalar("episode_end/done_reason_timeout", 1.0 if done_reason == "timeout" else 0.0, episode_step)
                self.tb_writer.add_scalar("episode_end/done_reason_clear", 1.0 if done_reason == "clear" else 0.0, episode_step)
                self.tb_writer.add_scalar(
                    "episode_end/done_reason_hard_stuck", 1.0 if done_reason == "hard_stuck" else 0.0, episode_step
                )
                self.tb_writer.flush()

            print(
                f"[episode_end] ep={episode_step} env={env_idx} reason={done_reason} "
                f"max_x={max_x_pos:.1f} flag_get={bool(flag_get)}"
            )
            self.episode_max_x_pos[env_idx] = 0.0

        return True

    def _on_training_end(self) -> None:
        if self.tb_writer is not None:
            self.tb_writer.close()
            self.tb_writer = None


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    reward_mode = "clip" if args.clip_reward and args.reward_mode == "raw" else args.reward_mode

    paths = ensure_artifact_paths()
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    env_slug = args.env_id.replace("/", "_")
    run_name = args.run_name.strip() or f"ppo_{env_slug}_{timestamp}"

    run_model_dir = paths.models / run_name
    run_log_dir = paths.logs / run_name
    run_tb_dir = run_log_dir / "tb"
    run_model_dir.mkdir(parents=True, exist_ok=True)
    run_tb_dir.mkdir(parents=True, exist_ok=True)

    requested_device = args.device.lower().strip()
    device, device_msg = resolve_torch_device(args.device)
    print(f"[device] {device} | {device_msg}")
    print(f"[device_diag] {collect_torch_runtime_diagnostics()}")

    explicit_device_requested = requested_device not in {"auto", "cpu"}
    if explicit_device_requested and device == "cpu" and not args.allow_device_fallback:
        raise RuntimeError(
            f"Requested --device {args.device!r}, but it resolved to CPU. {device_msg}\n"
            "Either fix your CUDA/MPS environment, or re-run with --allow-device-fallback "
            "to continue on CPU intentionally."
        )

    print(f"[run] {run_name}")
    print(
        f"[config] env_id={args.env_id}, movement={args.movement}, reward_mode={reward_mode}, "
        f"clip_reward={args.clip_reward}, n_envs={args.n_envs}, "
        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}"
    )

    env_fns = [
        make_env_fn(
            env_id=args.env_id,
            seed=args.seed + i,
            movement=args.movement,
            reward_mode=reward_mode,
            clip_reward=args.clip_reward,
            frame_skip=args.frame_skip,
            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,
        )
        for i in range(args.n_envs)
    ]
    vec_env = DummyVecEnv(env_fns)
    vec_env = VecMonitor(vec_env, filename=str(run_log_dir / "monitor.csv"))

    model = PPO(
        policy="CnnPolicy",
        env=vec_env,
        learning_rate=args.learning_rate,
        n_steps=args.n_steps,
        batch_size=args.batch_size,
        n_epochs=args.n_epochs,
        gamma=args.gamma,
        gae_lambda=args.gae_lambda,
        ent_coef=args.ent_coef,
        clip_range=args.clip_range,
        tensorboard_log=str(run_tb_dir),
        seed=args.seed,
        verbose=1,
        device=device,
    )

    if args.init_model_path.strip():
        init_model_path = Path(args.init_model_path).expanduser().resolve()
        if not init_model_path.exists():
            raise FileNotFoundError(f"Init model not found: {init_model_path}")
        try:
            model.set_parameters(str(init_model_path), exact_match=False, device=device)
            print(f"[init] Loaded initial weights from: {init_model_path}")
        except RuntimeError as exc:
            if not args.allow_partial_init:
                raise RuntimeError(
                    f"{exc}\n"
                    "Hint: checkpoint and current env may use different action spaces. "
                    "Try one of:\n"
                    "1) keep the same --movement as the checkpoint;\n"
                    "2) remove --init-model-path and train from scratch;\n"
                    "3) add --allow-partial-init to load compatible layers only."
                ) from exc
            loaded_count, skipped_count = load_partial_policy_weights(model, init_model_path, device=device)
            print(
                f"[init] Partial init from {init_model_path}: "
                f"loaded_tensors={loaded_count}, skipped_tensors={skipped_count}"
            )

    callback = CheckpointCallback(
        save_freq=max(args.save_freq // max(args.n_envs, 1), 1),
        save_path=str(run_model_dir),
        name_prefix="ppo_mario_ckpt",
        save_replay_buffer=False,
        save_vecnormalize=False,
    )
    episode_end_logging_callback = EpisodeEndLoggingCallback(
        n_envs=args.n_envs,
        tb_log_dir=run_tb_dir / "episode_end",
    )

    try:
        model.learn(
            total_timesteps=args.total_timesteps,
            callback=CallbackList([callback, episode_end_logging_callback]),
            tb_log_name="ppo",
            progress_bar=args.progress_bar,
        )
    finally:
        vec_env.close()

    final_model = run_model_dir / "final_model"
    model.save(str(final_model))

    latest_model = paths.models / "latest_model.zip"
    shutil.copy2(str(final_model) + ".zip", latest_model)
    pointer = write_latest_pointer(paths.models, latest_model)

    print(f"[done] Final model: {final_model}.zip")
    print(f"[done] Latest model pointer: {pointer}")


if __name__ == "__main__":
    main()
