"""Run a minimal evaluation loop for the RAM Pong policy."""

from __future__ import annotations

import argparse
import csv
import json
import sys
from datetime import datetime, timezone
from pathlib import Path

import envpool
import numpy as np
from PIL import Image, ImageDraw


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from pong_policy import ACTION_NAMES, NOOP, RamPolicy
from pong_state import PongState, detect_state, ram_ball_to_screen, ram_paddle_y_to_screen
from scripts.make_report import build_report


TRIAL_FIELDS = (
    "trial_index",
    "timestamp",
    "policy_name",
    "seed",
    "episodes",
    "score_mean",
    "score_min",
    "score_max",
    "env_steps",
    "notes",
)


EPISODE_FIELDS = (
    "timestamp",
    "policy_name",
    "seed",
    "episode",
    "score",
    "steps",
    "done",
    "rgb_video_path",
    "obs_video_path",
)

def reset_env(env: object) -> tuple[np.ndarray, dict]:
    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        obs, info = reset_out
    else:
        obs, info = reset_out, {}
    return np.asarray(obs), info


def step_env(env: object, action: int) -> tuple[np.ndarray, float, bool, dict]:
    step_out = env.step(np.array([action], dtype=np.int32))
    if len(step_out) == 5:
        obs, reward, terminated, truncated, info = step_out
        done = np.logical_or(terminated, truncated)
    else:
        obs, reward, done, info = step_out
    return (
        np.asarray(obs),
        float(np.asarray(reward).reshape(-1)[0]),
        bool(np.asarray(done).reshape(-1)[0]),
        info,
    )


def obs_to_frame(obs: np.ndarray, scale: int) -> Image.Image:
    frame = np.asarray(obs).reshape(1, 4, 84, 84)[0, -1]
    image = Image.fromarray(frame, mode="L")
    if scale != 1:
        image = image.resize((84 * scale, 84 * scale), Image.Resampling.NEAREST)
    return image.convert("P")


def rgb_to_frame(frame: np.ndarray, scale: int, overlay: dict | None = None) -> Image.Image:
    image = Image.fromarray(np.asarray(frame).reshape(-1, 210, 160, 3)[0], mode="RGB")
    if scale != 1:
        image = image.resize((160 * scale, 210 * scale), Image.Resampling.NEAREST)
    if overlay is not None:
        draw_rgb_overlay(image, scale, overlay)
    return image


def clamp(value: int, low: int, high: int) -> int:
    return max(low, min(high, value))


def draw_rgb_overlay(image: Image.Image, scale: int, overlay: dict) -> None:
    state = overlay.get("state")
    if not isinstance(state, PongState):
        return

    draw = ImageDraw.Draw(image)
    width, height = image.size

    screen_ball_x, screen_ball_y = ram_ball_to_screen(state.ball_x, state.ball_y)
    ball_x = clamp(round(screen_ball_x * scale), 0, width - 1)
    ball_y = clamp(round(screen_ball_y * scale), 0, height - 1)
    self_x = clamp(143 * scale, 0, width - 1)
    opponent_x = clamp(16 * scale, 0, width - 1)
    self_y = clamp(round(ram_paddle_y_to_screen(state.self_paddle_y) * scale), 0, height - 1)
    opponent_y = clamp(round(ram_paddle_y_to_screen(state.opponent_paddle_y) * scale), 0, height - 1)
    paddle_half = 8 * scale
    cross = 3 * scale

    if state.ball_active:
        draw.line((ball_x - cross, ball_y, ball_x + cross, ball_y), fill=(255, 230, 0), width=max(1, scale))
        draw.line((ball_x, ball_y - cross, ball_x, ball_y + cross), fill=(255, 230, 0), width=max(1, scale))
    else:
        draw.rectangle((ball_x - cross, ball_y - cross, ball_x + cross, ball_y + cross), outline=(255, 180, 0))

    draw.line(
        (self_x, self_y - paddle_half, self_x, self_y + paddle_half),
        fill=(0, 220, 255),
        width=max(2, 2 * scale),
    )
    draw.line(
        (opponent_x, opponent_y - paddle_half, opponent_x, opponent_y + paddle_half),
        fill=(255, 80, 220),
        width=max(2, 2 * scale),
    )

    action = int(overlay.get("action", NOOP))
    label = (
        f"s={overlay.get('step', 0)} "
        f"a={ACTION_NAMES.get(action, str(action))} "
        f"r={overlay.get('reward', 0):.0f} "
        f"score={overlay.get('score', 0):.0f} "
        f"ball=({state.ball_x},{state.ball_y}) "
        f"pad={state.self_paddle_y}"
    )
    text_box = (2, 2, min(width - 1, 154 * scale), 12 * scale)
    draw.rectangle(text_box, fill=(0, 0, 0))
    draw.text((4, 3), label, fill=(255, 255, 255))


def save_gif(frames: list[Image.Image], path: Path, duration_ms: int) -> None:
    if not frames:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    encoded_frames = [frame_with_marker(frame, index) for index, frame in enumerate(frames)]
    encoded_frames[0].save(
        path,
        save_all=True,
        append_images=encoded_frames[1:],
        duration=duration_ms,
        loop=0,
        optimize=False,
    )


def frame_with_marker(frame: Image.Image, index: int) -> Image.Image:
    marked = frame.copy()
    marker = index % 256
    if marked.mode == "RGB":
        marked.putpixel((0, 0), (marker, marker, marker))
    else:
        marked.putpixel((0, 0), marker)
    return marked


def run_episode(
    seed: int,
    max_steps: int,
    rgb_video_path: Path | None,
    obs_video_path: Path | None,
    video_every: int,
    obs_video_scale: int,
    rgb_video_scale: int,
    video_duration_ms: int,
    annotate_rgb: bool,
) -> tuple[float, int, bool]:
    env_kwargs = {}
    if rgb_video_path is not None:
        env_kwargs["render_mode"] = "rgb_array"
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed, **env_kwargs)
    try:
        policy = RamPolicy()
        obs, info = reset_env(env)
        policy.reset()
        score = 0.0
        obs_frames = []
        rgb_frames = []
        state = detect_state(info)
        if obs_video_path is not None:
            obs_frames.append(obs_to_frame(obs, obs_video_scale))
        if rgb_video_path is not None:
            overlay = {"step": 0, "action": NOOP, "reward": 0.0, "score": score, "state": state}
            rgb_frames.append(rgb_to_frame(env.render(), rgb_video_scale, overlay if annotate_rgb else None))

        for step in range(1, max_steps + 1):
            state = detect_state(info)
            action = policy.select_action(info)
            obs, reward, done, info = step_env(env, action)
            score += reward
            if step % video_every == 0:
                if obs_video_path is not None:
                    obs_frames.append(obs_to_frame(obs, obs_video_scale))
                if rgb_video_path is not None:
                    overlay = {
                        "step": step,
                        "action": action,
                        "reward": reward,
                        "score": score,
                        "state": state,
                    }
                    rgb_frames.append(
                        rgb_to_frame(env.render(), rgb_video_scale, overlay if annotate_rgb else None)
                    )
            if done:
                if obs_video_path is not None:
                    save_gif(obs_frames, obs_video_path, video_duration_ms)
                if rgb_video_path is not None:
                    save_gif(rgb_frames, rgb_video_path, video_duration_ms)
                return score, step, True
        if obs_video_path is not None:
            save_gif(obs_frames, obs_video_path, video_duration_ms)
        if rgb_video_path is not None:
            save_gif(rgb_frames, rgb_video_path, video_duration_ms)
        return score, max_steps, False
    finally:
        env.close()


def append_trial(path: Path, trial: dict) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(trial, separators=(",", ":")))
        handle.write("\n")


def read_trials(path: Path) -> list[dict]:
    if not path.exists():
        return []

    trials = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                trials.append(json.loads(line))
    return trials


def write_summary(path: Path, trials: list[dict]) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=TRIAL_FIELDS)
        writer.writeheader()
        for trial in trials:
            writer.writerow({field: trial.get(field, "") for field in TRIAL_FIELDS})


def append_episode(path: Path, episode: dict) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(episode, separators=(",", ":")))
        handle.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate the minimal RAM Pong policy.")
    parser.add_argument("--episodes", type=int, default=1)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--max-steps", type=int, default=10_000)
    parser.add_argument("--policy-name", default="minimal_ram_v0")
    parser.add_argument("--notes", default="")
    parser.add_argument("--trials-path", type=Path, default=Path("trials.jsonl"))
    parser.add_argument("--summary-path", type=Path, default=Path("summary.csv"))
    parser.add_argument("--episodes-path", type=Path, default=Path("episodes.jsonl"))
    parser.add_argument("--report-path", type=Path, default=Path("report.html"))
    parser.add_argument("--video-dir", type=Path, default=Path("videos"))
    parser.add_argument("--no-video", action="store_true")
    parser.add_argument("--video-every", type=int, default=1)
    parser.add_argument("--obs-video-scale", type=int, default=4)
    parser.add_argument("--rgb-video-scale", type=int, default=2)
    parser.add_argument("--video-duration-ms", type=int, default=50)
    parser.add_argument("--no-rgb-overlay", action="store_true")
    args = parser.parse_args()
    if args.video_every < 1:
        raise ValueError("--video-every must be at least 1")
    if args.obs_video_scale < 1:
        raise ValueError("--obs-video-scale must be at least 1")
    if args.rgb_video_scale < 1:
        raise ValueError("--rgb-video-scale must be at least 1")

    print(f"envpool: {envpool.__version__}")
    print(f"episodes: {args.episodes}")
    print(f"seed: {args.seed}")
    print(f"max steps: {args.max_steps}")
    print("episode,score,steps,done,rgb_video_path,obs_video_path")

    scores = []
    total_steps = 0
    timestamp = datetime.now(timezone.utc).isoformat()
    trial_index = len(read_trials(args.trials_path)) + 1
    for episode in range(args.episodes):
        episode_seed = args.seed + episode
        rgb_video_path = None
        obs_video_path = None
        if not args.no_video:
            safe_timestamp = timestamp.replace(":", "").replace("+", "Z")
            base_name = f"{args.policy_name}_seed{episode_seed}_ep{episode + 1}_{safe_timestamp}"
            rgb_video_path = args.video_dir / f"{base_name}_rgb.gif"
            obs_video_path = args.video_dir / f"{base_name}_obs.gif"

        score, steps, done = run_episode(
            seed=episode_seed,
            max_steps=args.max_steps,
            rgb_video_path=rgb_video_path,
            obs_video_path=obs_video_path,
            video_every=args.video_every,
            obs_video_scale=args.obs_video_scale,
            rgb_video_scale=args.rgb_video_scale,
            video_duration_ms=args.video_duration_ms,
            annotate_rgb=not args.no_rgb_overlay,
        )
        scores.append(score)
        total_steps += steps
        rgb_video_label = str(rgb_video_path) if rgb_video_path is not None else ""
        obs_video_label = str(obs_video_path) if obs_video_path is not None else ""
        print(f"{episode + 1},{score},{steps},{done},{rgb_video_label},{obs_video_label}")
        episode_record = {
            "trial_index": trial_index,
            "timestamp": timestamp,
            "policy_name": args.policy_name,
            "seed": episode_seed,
            "episode": episode + 1,
            "score": score,
            "steps": steps,
            "done": done,
            "rgb_video_path": rgb_video_label,
            "obs_video_path": obs_video_label,
        }
        if rgb_video_label:
            episode_record["video_path"] = rgb_video_label
        append_episode(args.episodes_path, episode_record)

    trial = {
        "trial_index": trial_index,
        "timestamp": timestamp,
        "policy_name": args.policy_name,
        "seed": args.seed,
        "episodes": args.episodes,
        "score_mean": float(np.mean(scores)),
        "score_min": float(np.min(scores)),
        "score_max": float(np.max(scores)),
        "env_steps": total_steps,
        "notes": args.notes,
    }
    append_trial(args.trials_path, trial)
    write_summary(args.summary_path, read_trials(args.trials_path))
    build_report(
        trials_path=args.trials_path,
        summary_path=args.summary_path,
        episodes_path=args.episodes_path,
        output_path=args.report_path,
    )

    print()
    print(f"trial appended: {args.trials_path}")
    print(f"episodes appended: {args.episodes_path}")
    print(f"summary written: {args.summary_path}")
    print(f"report written: {args.report_path}")
    print(
        "scores "
        f"mean={trial['score_mean']} "
        f"min={trial['score_min']} "
        f"max={trial['score_max']} "
        f"env_steps={trial['env_steps']}"
    )


if __name__ == "__main__":
    main()
