#!/usr/bin/env python3
"""
record_video.py

Runs a trained MaskablePPO model and records the rollout to MP4.

Frames are created from the observation tensor using a simple rendering
helper. If later you add pixel-based observations, drop-in replacement
will automatically produce real environment frames.

Output:
    ./videos/utdg_run_<timestamp>.mp4
"""

import os
import sys
import time
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import numpy as np
import gymnasium as gym
import imageio

from sb3_contrib import MaskablePPO
from sb3_contrib.common.wrappers import ActionMasker
from stable_baselines3.common.vec_env import DummyVecEnv

from utdg_env.utils.hydra_loader import load_config, pretty_print_cfg
import utdg_env


# ---------------------------------------------------------
# Frame Rendering Helper
# ---------------------------------------------------------
def render_frame(obs: dict, width: int = 512, height: int = 512):
    """
    Creates a visualization from dict observation.
    Produces a grayscale image showing key game state
    (gold, enemies, slots, etc.) as a simple visual trace.

    Not a real screenshot of the Godot game.
    Replace later with pixel obs or screen capture if desired.
    """
    img = np.zeros((height, width, 3), dtype=np.uint8)

    # Normalize and draw values as horizontal bars
    def bar(val, max_val, y, color):
        w = int(width * (val / max_val))
        img[y:y+20, :w] = color

    bar(int(obs["gold"][0]), 300, 20, (255, 215, 0))
    bar(int(obs["enemy_count"][0]), 50, 60, (255, 0, 0))
    bar(int(obs["tower_count"][0]), 20, 100, (0, 255, 0))
    bar(int(obs["base_health"][0]), 20, 140, (0, 128, 255))

    return img


# ---------------------------------------------------------
# Mask callback
# ---------------------------------------------------------
def mask_fn(env):
    while hasattr(env, "env"):
        env = env.env
    return env.get_action_mask()


# ---------------------------------------------------------
# Main Recording Routine
# ---------------------------------------------------------
def main():
    cfg = load_config()
    print("\n========== Loaded Config ==========")
    print(pretty_print_cfg(cfg))
    print("===================================\n")

    video_dir = "videos"
    os.makedirs(video_dir, exist_ok=True)

    video_path = os.path.join(
        video_dir, f"utdg_run_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
    )
    writer = imageio.get_writer(video_path, fps=10)

    # Build env
    env = DummyVecEnv([
        lambda: ActionMasker(gym.make("UTDGEnv-v0", sb3_mode=False), mask_fn)
    ])

    # Load trained policy
    model_path = os.path.join("models", "model_policy.zip")
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file missing → {model_path}")

    print(f"[record] Loading model from: {model_path}")
    model = MaskablePPO.load(model_path, env=env)
    print("[record] Model loaded.")

    obs = env.reset()
    done = False
    timestep = 0

    print(f"[record] Recording to: {video_path}")

    while not done and timestep < 2000:  # hard cap safety
        timestep += 1

        # Predict an action
        action, _ = model.predict(obs, deterministic=True)

        # Step
        obs, reward, done, info = env.step(action)

        # Convert to single env obs
        frame = render_frame({k: v[0] for k, v in obs.items()})
        writer.append_data(frame)

        print(f"[step {timestep}] action={action[0]}, reward={reward[0]}, done={done[0]}")

        if done[0]:
            print("[record] Episode finished.")
            break

    writer.close()
    env.close()

    print(f"\n[record] Saved video → {video_path}")
    print("[record] Done.")


if __name__ == "__main__":
    main()
