"""Episode artifact helpers for real BrowserGym rollouts."""

from __future__ import annotations

import base64
import json
import os
from pathlib import Path
from typing import Any, Dict, Mapping


def artifact_root() -> Path:
    return Path(os.getenv("BROWSER_ENV_ARTIFACT_DIR", "artifacts"))


def save_step_artifacts(
    *,
    episode_id: str,
    step: int,
    raw_obs: Mapping[str, Any],
    event: Dict[str, Any],
) -> Dict[str, str]:
    """Persist step metadata and any screenshot-like BrowserGym observation field."""

    step_dir = artifact_root() / "episodes" / episode_id / f"step_{step:03d}"
    step_dir.mkdir(parents=True, exist_ok=True)

    event_path = step_dir / "event.json"
    event_path.write_text(json.dumps(event, indent=2, default=str), encoding="utf-8")

    saved = {"event": str(event_path)}
    screenshot_path = _save_screenshot(raw_obs, step_dir / "screenshot.png")
    if screenshot_path:
        saved["screenshot"] = screenshot_path
    return saved


def _save_screenshot(raw_obs: Mapping[str, Any], path: Path) -> str | None:
    for key in ("screenshot", "image", "page_screenshot", "active_page_screenshot"):
        if key in raw_obs:
            if _write_image(raw_obs[key], path):
                return str(path)
    return None


def _write_image(value: Any, path: Path) -> bool:
    if value is None:
        return False
    if isinstance(value, bytes):
        path.write_bytes(value)
        return True
    if isinstance(value, str):
        try:
            path.write_bytes(base64.b64decode(value))
            return True
        except Exception:
            return False
    try:
        from PIL import Image

        if isinstance(value, Image.Image):
            value.save(path)
            return True
    except Exception:
        pass
    try:
        from PIL import Image

        image = Image.fromarray(value)
        image.save(path)
        return True
    except Exception:
        return False