"""Video recording with the per-frame HUD overlay.

Also owns the warm-up-frame drop: the first frames after a reset render with the wrong camera
pose and unresolved textures, which made reviewers describe scenes (missing containers) that
never existed in the episode.
"""
from __future__ import annotations

import os

import numpy as np

WARMUP_FRAMES = 2


class Recorder:
    def __init__(self, env, title="", episode=-1, fps=14):
        self.env = env
        self.title = title
        self.episode = episode
        self.fps = fps
        self.frames = []
        self.phase = "start"
        self.result = ""
        self.lines_fn = None          # optional task-supplied extra HUD lines

    def capture(self):
        img = self.env.render()
        if img is None:
            return
        from PIL import Image, ImageDraw
        im = Image.fromarray(np.asarray(img)[..., :3].copy())
        d = ImageDraw.Draw(im)
        lines = [f"=== {self.title} ===" + (f"  EP {self.episode}" if self.episode >= 0 else "")]
        if self.result:
            lines.append(f"RESULT: {self.result}")
        lines.append(f"ACTION: {self.phase}")
        if self.lines_fn is not None:
            try:
                lines += list(self.lines_fn())
            except Exception:
                pass
        d.rectangle([0, 0, 470, 18*len(lines)+6], fill=(0, 0, 0))
        y = 3
        for ln in lines:
            d.text((6, y), ln, fill=(255, 235, 60)); y += 18
        self.frames.append(np.array(im))

    def save(self, path):
        import imageio.v2 as imageio
        frames = self.frames[WARMUP_FRAMES:] if len(self.frames) > WARMUP_FRAMES+4 else self.frames
        if not frames:
            print("[rec] no frames to write"); return None
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        imageio.mimsave(path, frames, fps=self.fps)
        print(f"[rec] video -> {path} ({len(frames)} frames)", flush=True)
        return path