"""Safe subprocess-isolated PPO evaluation and MuJoCo rendering.

This module prevents native OpenGL conflicts by keeping Stable-Baselines3
model inference and MuJoCo frame rendering in separate Python processes.
It targets Gymnasium MuJoCo environments with continuous Box actions,
Stable-Baselines3 PPO checkpoints, and VecNormalize statistics.
"""

__version__ = "1.0.0"
__all__ = ["evaluate_and_render_ppo_mujoco"]

# Q3. 재사용 가능한 PPO/MuJoCo 분리 평가·렌더링 도구
#
# 대상:
# - Gymnasium MuJoCo 연속 행동(Box) 환경
# - Stable-Baselines3 PPO 체크포인트
# - VecNormalize 통계
#
# 핵심:
# - SB3 프로세스는 평가와 행동 생성만 수행
# - MuJoCo 프로세스는 렌더링만 수행
# - 동일 seed/행동의 보상 일치를 자동 검증

import json
import math
import os
import subprocess
import sys
from pathlib import Path


GENERIC_EVALUATION_WORKER = r"""
import json
from pathlib import Path

import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.vec_env import VecNormalize


request = json.loads(input())

env_id = request["env_id"]
model_path = Path(request["model_path"])
stats_path = Path(request["stats_path"])
device = request["device"]

n_eval_episodes = int(
    request["n_eval_episodes"]
)
n_random_episodes = int(
    request["n_random_episodes"]
)

eval_seed = int(request["eval_seed"])
random_seed = int(request["random_seed"])
replay_seed = int(request["replay_seed"])

max_replay_steps = request[
    "max_replay_steps"
]
if max_replay_steps is not None:
    max_replay_steps = int(max_replay_steps)

model = PPO.load(
    model_path,
    device=device,
)


# ------------------------------------------------------------
# 1. 정규화 통계를 적용한 PPO 평가
# ------------------------------------------------------------

eval_raw_env = make_vec_env(
    env_id,
    n_envs=1,
    seed=eval_seed,
)
eval_env = VecNormalize.load(
    stats_path,
    eval_raw_env,
)
eval_env.training = False
eval_env.norm_reward = False

try:
    mean_reward, std_reward = evaluate_policy(
        model,
        eval_env,
        n_eval_episodes=n_eval_episodes,
        deterministic=True,
        warn=False,
    )
finally:
    eval_env.close()


# ------------------------------------------------------------
# 2. 랜덤 정책 기준선
# ------------------------------------------------------------

random_env = gym.make(env_id)
random_rewards = []

try:
    for episode_index in range(
        n_random_episodes
    ):
        episode_seed = (
            random_seed + episode_index
        )
        obs, info = random_env.reset(
            seed=episode_seed
        )
        random_env.action_space.seed(
            episode_seed
        )

        episode_reward = 0.0
        terminated = False
        truncated = False

        while not (terminated or truncated):
            action = (
                random_env.action_space.sample()
            )
            (
                obs,
                reward,
                terminated,
                truncated,
                info,
            ) = random_env.step(action)
            episode_reward += float(reward)

        random_rewards.append(episode_reward)
finally:
    random_env.close()

random_mean_reward = float(
    np.mean(random_rewards)
)
random_std_reward = float(
    np.std(random_rewards)
)


# ------------------------------------------------------------
# 3. 결정론적 재생 행동 생성
# ------------------------------------------------------------

replay_raw_env = make_vec_env(
    env_id,
    n_envs=1,
    seed=replay_seed,
)
replay_env = VecNormalize.load(
    stats_path,
    replay_raw_env,
)
replay_env.training = False
replay_env.norm_reward = False

actions = []
reference_reward = 0.0

try:
    obs = replay_env.reset()

    environment_limit = int(
        replay_raw_env.get_attr(
            "spec"
        )[0].max_episode_steps
        or 1000
    )
    replay_limit = environment_limit

    if max_replay_steps is not None:
        replay_limit = min(
            replay_limit,
            max_replay_steps,
        )

    for step_index in range(replay_limit):
        action, state = model.predict(
            obs,
            deterministic=True,
        )

        if not hasattr(
            replay_env.action_space,
            "low",
        ):
            raise TypeError(
                "이 도구는 연속 행동(Box) 환경을 "
                "대상으로 합니다."
            )

        env_action = np.clip(
            action[0],
            replay_env.action_space.low,
            replay_env.action_space.high,
        ).astype(np.float64)

        actions.append(env_action.tolist())

        obs, rewards, dones, infos = (
            replay_env.step(
                env_action[None, :]
            )
        )
        reference_reward += float(rewards[0])

        if bool(dones[0]):
            break
finally:
    replay_env.close()

if not actions:
    raise RuntimeError(
        "재생할 행동을 생성하지 못했습니다."
    )

improvement = (
    float(mean_reward) - random_mean_reward
)

result = {
    "env_id": env_id,
    "algorithm": "PPO",
    "mean_reward": float(mean_reward),
    "std_reward": float(std_reward),
    "n_eval_episodes": n_eval_episodes,
    "eval_seed": eval_seed,
    "random_mean_reward": random_mean_reward,
    "random_std_reward": random_std_reward,
    "n_random_episodes": n_random_episodes,
    "random_seed": random_seed,
    "improvement_over_random": improvement,
    "training_success": improvement > 0.0,
    "replay_seed": replay_seed,
    "reference_reward": reference_reward,
    "actions": actions,
}

print(
    json.dumps(
        result,
        ensure_ascii=False,
    ),
    flush=True,
)
"""


GENERIC_RENDER_WORKER = r"""
import json
import os
from pathlib import Path

import gymnasium as gym
import numpy as np


request = json.loads(input())

env_id = request["env_id"]
replay_seed = int(request["replay_seed"])
actions = request["actions"]

width = int(request["width"])
height = int(request["height"])

video_path_value = request["video_path"]
video_path = (
    Path(video_path_value)
    if video_path_value is not None
    else None
)

render_env = gym.make(
    env_id,
    render_mode="rgb_array",
    width=width,
    height=height,
)

frames = []
episode_reward = 0.0
steps = 0
fps = int(
    render_env.metadata.get(
        "render_fps",
        30,
    )
)

try:
    obs, info = render_env.reset(
        seed=replay_seed
    )

    first_frame = render_env.render()
    if first_frame is None:
        raise RuntimeError(
            "초기 RGB 프레임 생성 실패"
        )

    frames.append(
        np.asarray(
            first_frame,
            dtype=np.uint8,
        )
    )

    for step_index, action_values in enumerate(
        actions
    ):
        action = np.asarray(
            action_values,
            dtype=np.float64,
        )
        action = np.clip(
            action,
            render_env.action_space.low,
            render_env.action_space.high,
        )

        (
            obs,
            reward,
            terminated,
            truncated,
            info,
        ) = render_env.step(action)

        episode_reward += float(reward)
        steps = step_index + 1

        frame = render_env.render()
        if frame is None:
            raise RuntimeError(
                f"{steps}번째 RGB 프레임 생성 실패"
            )

        frames.append(
            np.asarray(
                frame,
                dtype=np.uint8,
            )
        )

        if terminated or truncated:
            break
finally:
    render_env.close()

if steps != len(actions):
    raise RuntimeError(
        "행동 배열을 끝까지 재생하지 못했습니다."
    )

video_bytes = None

if video_path is not None:
    import imageio.v2 as imageio

    video_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    imageio.mimsave(
        video_path,
        frames,
        fps=fps,
        codec="libx264",
        quality=8,
        macro_block_size=16,
    )

    if (
        not video_path.is_file()
        or video_path.stat().st_size == 0
    ):
        raise RuntimeError(
            "비디오 파일 생성 실패"
        )

    video_bytes = video_path.stat().st_size

result = {
    "render_backend": os.environ["MUJOCO_GL"],
    "rendered_reward": episode_reward,
    "steps": steps,
    "frames": len(frames),
    "fps": fps,
    "width": width,
    "height": height,
    "video_path": (
        str(video_path)
        if video_path is not None
        else None
    ),
    "video_bytes": video_bytes,
}

print(
    json.dumps(
        result,
        ensure_ascii=False,
    ),
    flush=True,
)
"""


def _run_json_worker(
    *,
    code,
    payload,
    worker_env,
    label,
    timeout_seconds,
):
    result = subprocess.run(
        [
            sys.executable,
            "-u",
            "-X",
            "faulthandler",
            "-c",
            code,
        ],
        env=worker_env,
        input=json.dumps(
            payload,
            ensure_ascii=False,
        ),
        text=True,
        capture_output=True,
        timeout=timeout_seconds,
    )

    if result.returncode != 0:
        raise RuntimeError(
            f"{label} 실패 "
            f"(return code {result.returncode})\n"
            f"stdout tail:\n"
            f"{result.stdout[-3000:]}\n"
            f"stderr tail:\n"
            f"{result.stderr[-6000:]}"
        )

    try:
        return json.loads(
            result.stdout.strip().splitlines()[-1]
        )
    except (
        IndexError,
        json.JSONDecodeError,
    ) as exc:
        raise RuntimeError(
            f"{label} 결과 JSON 파싱 실패"
        ) from exc


def evaluate_and_render_ppo_mujoco(
    *,
    env_id,
    model_path,
    vecnormalize_path,
    video_path=None,
    device="cpu",
    n_eval_episodes=30,
    n_random_episodes=30,
    eval_seed=1000,
    random_seed=2000,
    replay_seed=42,
    max_replay_steps=None,
    width=640,
    height=480,
    backends=("egl", "osmesa"),
    timeout_seconds=180,
):
    """
    PPO/VecNormalize 모델을 평가하고 MuJoCo 영상을
    별도 프로세스에서 안전하게 렌더링합니다.

    video_path=None이면 프레임 생성까지만 검사하고
    파일은 만들지 않습니다.
    """

    model_path = Path(model_path)
    vecnormalize_path = Path(
        vecnormalize_path
    )

    for path in (
        model_path,
        vecnormalize_path,
    ):
        if not path.is_file():
            raise FileNotFoundError(path)

    if n_eval_episodes < 1:
        raise ValueError(
            "n_eval_episodes는 1 이상이어야 합니다."
        )

    if n_random_episodes < 1:
        raise ValueError(
            "n_random_episodes는 1 이상이어야 합니다."
        )

    if not backends:
        raise ValueError(
            "렌더링 백엔드가 하나 이상 필요합니다."
        )

    evaluation_payload = {
        "env_id": env_id,
        "model_path": str(model_path),
        "stats_path": str(vecnormalize_path),
        "device": device,
        "n_eval_episodes": n_eval_episodes,
        "n_random_episodes": (
            n_random_episodes
        ),
        "eval_seed": eval_seed,
        "random_seed": random_seed,
        "replay_seed": replay_seed,
        "max_replay_steps": (
            max_replay_steps
        ),
    }

    evaluation_env = os.environ.copy()
    evaluation_env.pop("MUJOCO_GL", None)
    evaluation_env.pop(
        "PYOPENGL_PLATFORM",
        None,
    )

    evaluation_result = _run_json_worker(
        code=GENERIC_EVALUATION_WORKER,
        payload=evaluation_payload,
        worker_env=evaluation_env,
        label="SB3 평가·행동 생성",
        timeout_seconds=timeout_seconds,
    )

    actions = evaluation_result.pop(
        "actions"
    )

    render_payload = {
        "env_id": env_id,
        "replay_seed": replay_seed,
        "actions": actions,
        "width": width,
        "height": height,
        "video_path": (
            str(Path(video_path))
            if video_path is not None
            else None
        ),
    }

    render_errors = {}
    render_result = None

    for backend in backends:
        worker_env = os.environ.copy()
        worker_env["MUJOCO_GL"] = backend
        worker_env[
            "PYOPENGL_PLATFORM"
        ] = backend
        worker_env.pop("DISPLAY", None)

        try:
            render_result = _run_json_worker(
                code=GENERIC_RENDER_WORKER,
                payload=render_payload,
                worker_env=worker_env,
                label=(
                    "MuJoCo 렌더링 "
                    f"({backend})"
                ),
                timeout_seconds=(
                    timeout_seconds
                ),
            )
            break
        except RuntimeError as exc:
            render_errors[backend] = str(exc)

    if render_result is None:
        raise RuntimeError(
            "모든 렌더링 백엔드가 실패했습니다:\n"
            + json.dumps(
                render_errors,
                indent=2,
                ensure_ascii=False,
            )
        )

    reference_reward = float(
        evaluation_result[
            "reference_reward"
        ]
    )
    rendered_reward = float(
        render_result["rendered_reward"]
    )

    replay_verified = math.isclose(
        reference_reward,
        rendered_reward,
        rel_tol=1e-7,
        abs_tol=1e-7,
    )

    if not replay_verified:
        raise RuntimeError(
            "동일 행동 재생 보상 불일치: "
            f"reference={reference_reward}, "
            f"rendered={rendered_reward}"
        )

    if (
        render_result["frames"]
        != render_result["steps"] + 1
    ):
        raise RuntimeError(
            "프레임 수가 steps + 1과 "
            "일치하지 않습니다."
        )

    return {
        **evaluation_result,
        **render_result,
        "replay_reward_verified": True,
    }