"""
orbit_module.py — 2.5D actor orbit preview via depth-card parallax.

Pipeline:
  1. load actor image  — use provided path directly (selection is caller's responsibility)
  2. create_depth_map  — fake depth from alpha mask distance transform
  3. find_bg_plate     — static background for compositing (original for nobg; blurred for opaque)
  4. render_orbit      — per-frame: parallax-warp actor, composite over static bg
  5. save_orbit_output — write RGBA PNGs + MP4

The "orbit" illusion requires a static reference (background plate).
Without it the viewer has nothing to anchor on and it reads as a side-slide.

Usage:
  from orbit_module import run_orbit_pipeline
  result = run_orbit_pipeline(image_path, output_dir)

CLI: see orbit_poc.py
"""

import os
import math
import shutil
import subprocess
import tempfile

import cv2
import numpy as np
from scipy.ndimage import distance_transform_edt

__all__ = [
    "create_depth_map",
    "render_orbit_frame",
    "render_orbit",
    "save_orbit_output",
    "run_orbit_pipeline",
]

# ---------------------------------------------------------------------------
# Image loading helpers
# ---------------------------------------------------------------------------

def _load_rgba(path: str) -> np.ndarray:
    """Load any image as RGBA uint8 H×W×4."""
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f"Cannot read image: {path}")
    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGRA)
    elif img.shape[2] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
    return cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)


def _has_real_alpha(rgba: np.ndarray) -> bool:
    """True if the image contains meaningful transparency (not just all-255)."""
    alpha = rgba[:, :, 3]
    transparent_pct = float((alpha < 32).mean())
    return transparent_pct > 0.05  # >5% transparent pixels = real alpha


# ---------------------------------------------------------------------------
# Background plate
# ---------------------------------------------------------------------------

def _find_original_for_nobg(actor_path: str) -> str | None:
    """
    Given a .nobg.png sidecar path, find the original opaque image.
    e.g. foo.nobg.png → foo.png or foo.jpg
    """
    root, _ = os.path.splitext(actor_path)
    if not root.endswith(".nobg"):
        return None
    base = root[: -len(".nobg")]
    for ext in (".png", ".jpg", ".jpeg", ".webp"):
        p = base + ext
        if os.path.exists(p):
            return p
    return None


def _make_bg_plate(actor_rgba: np.ndarray) -> np.ndarray:
    """
    Build a static background plate for opaque images (no nobg available).

    Strategy: blur the source image with a large kernel. The blurred copy stays
    fixed while the sharp actor layer shifts — creates subtle depth separation.
    """
    H, W = actor_rgba.shape[:2]
    # Large blur: simulates out-of-focus background
    blurred_rgb = cv2.GaussianBlur(actor_rgba[:, :, :3], (0, 0), max(H, W) * 0.04)
    plate = np.dstack([blurred_rgb, np.full((H, W), 255, dtype=np.uint8)])
    return plate.astype(np.uint8)


def get_bg_plate(actor_path: str, actor_rgba: np.ndarray) -> np.ndarray:
    """
    Return the background plate for the orbit:
    - .nobg.png: load the matching original opaque image (best result)
    - Opaque image: return blurred copy (subtle but functional)
    Plate is resized to match actor_rgba dimensions.
    """
    H, W = actor_rgba.shape[:2]

    orig = _find_original_for_nobg(actor_path)
    if orig:
        bg = _load_rgba(orig)
        if bg.shape[:2] != (H, W):
            bg = cv2.resize(bg, (W, H), interpolation=cv2.INTER_AREA)
        bg[:, :, 3] = 255  # ensure fully opaque background
        return bg

    return _make_bg_plate(actor_rgba)


# ---------------------------------------------------------------------------
# Depth map
# ---------------------------------------------------------------------------

def create_depth_map(image_rgba: np.ndarray) -> np.ndarray:
    """
    Float32 H×W depth in [0,1].  1 = closest (subject centre), 0 = far/background.

    Uses the alpha mask distance transform:
    - For transparent-bg images: EDT of the foreground mask (subject body)
    - For opaque images: EDT from image edges (assumes subject is centred)

    Power-law shaping (^0.5) keeps the gradient gradual near the centre.
    """
    alpha = image_rgba[:, :, 3]
    mask = (alpha > 32).astype(np.uint8)

    if mask.sum() == 0:
        return np.zeros(alpha.shape, dtype=np.float32)

    dist = distance_transform_edt(mask).astype(np.float32)
    max_d = dist.max()
    if max_d > 0:
        dist /= max_d
    return np.sqrt(dist)


# ---------------------------------------------------------------------------
# Orbit rendering
# ---------------------------------------------------------------------------

def _alpha_composite(fg: np.ndarray, bg: np.ndarray) -> np.ndarray:
    """Alpha-composite RGBA fg over RGBA bg. Returns RGBA uint8."""
    a = fg[:, :, 3:4].astype(np.float32) / 255.0
    out_rgb = fg[:, :, :3].astype(np.float32) * a + bg[:, :, :3].astype(np.float32) * (1.0 - a)
    out_a   = fg[:, :, 3:4].astype(np.float32) + bg[:, :, 3:4].astype(np.float32) * (1.0 - a)
    return np.dstack([out_rgb.clip(0, 255), out_a.clip(0, 255)]).astype(np.uint8)


def render_orbit_frame(
    actor_rgba: np.ndarray,
    depth: np.ndarray,
    theta: float,
    parallax_strength: float = 0.08,
    bg_rgba: np.ndarray | None = None,
) -> np.ndarray:
    """
    Swing-mode frame: depth-based parallax shift only.
    Closer pixels (depth≈1) shift more than far pixels (depth≈0).
    Composited over static bg plate for perceivable depth.
    Returns RGBA uint8 H×W×4.
    """
    H, W = actor_rgba.shape[:2]

    shift_x = depth * (W * parallax_strength * math.sin(theta))
    shift_y = depth * (H * parallax_strength * 0.03 * -math.cos(theta))

    yc, xc = np.mgrid[0:H, 0:W].astype(np.float32)
    map_x = (xc - shift_x).astype(np.float32)
    map_y = (yc - shift_y).astype(np.float32)

    bgra = cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA)
    warped_bgra = cv2.remap(bgra, map_x, map_y,
                            interpolation=cv2.INTER_LINEAR,
                            borderMode=cv2.BORDER_CONSTANT,
                            borderValue=(0, 0, 0, 0))
    warped = cv2.cvtColor(warped_bgra, cv2.COLOR_BGRA2RGBA)

    if bg_rgba is None:
        return warped
    return _alpha_composite(warped, bg_rgba)


def _perspective_card_frame(
    bgra_src: np.ndarray,
    theta: float,
) -> np.ndarray:
    """
    Orbit-mode frame: simulate a flat card rotating around its vertical axis.

    - Squishes width by |cos(θ)|, centred on canvas
    - Mirrors source for the back half (cos < 0) so the "back" is visible
    - Returns BGRA with transparent borders (ready for bg composite)

    At θ=0°  → full width, unmirrored (front face)
    At θ=90° → hair-thin line
    At θ=180°→ full width, mirrored (back face)
    """
    H, W = bgra_src.shape[:2]
    cos_t = math.cos(theta)
    compress = abs(cos_t)
    is_back = cos_t < 0

    if compress < 0.025:
        # Near-edge-on: return a one-pixel-wide vertical strip to avoid singularity
        out = np.zeros_like(bgra_src)
        mid = W // 2
        out[:, mid:mid+1] = bgra_src[:, mid:mid+1]
        return out

    new_w = max(int(round(W * compress)), 2)
    x0 = (W - new_w) // 2
    x1 = x0 + new_w

    # Source corners: mirror left↔right for the back face
    if is_back:
        src = np.float32([[W, 0], [0, 0], [0, H], [W, H]])
    else:
        src = np.float32([[0, 0], [W, 0], [W, H], [0, H]])

    dst = np.float32([[x0, 0], [x1, 0], [x1, H], [x0, H]])

    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(bgra_src, M, (W, H),
                               flags=cv2.INTER_LINEAR,
                               borderMode=cv2.BORDER_CONSTANT,
                               borderValue=(0, 0, 0, 0))


def render_orbit(
    actor_rgba: np.ndarray,
    depth: np.ndarray,
    n_frames: int = 36,
    parallax_strength: float = 0.08,
    mode: str = "swing",
    max_angle_deg: float = 35.0,
    bg_rgba: np.ndarray | None = None,
) -> list:
    """
    Render all orbit frames.

    mode='swing'  — sinusoidal ±max_angle_deg depth-parallax, loops cleanly
    mode='orbit'  — full 360° perspective card rotation (compress + mirror)

    Returns list of RGBA uint8 frames.
    """
    if mode == "swing":
        max_rad = math.radians(max_angle_deg)
        angles = [max_rad * math.sin(2 * math.pi * i / n_frames) for i in range(n_frames)]
        return [render_orbit_frame(actor_rgba, depth, theta, parallax_strength, bg_rgba)
                for theta in angles]

    elif mode == "orbit":
        # Don't use the photo bg plate — the transparent areas should stay transparent
        # so the perspective compression (card getting thin at 90°, mirrored at 180°)
        # is clearly visible. Solid bg is added at MP4 write time.
        angles = [2 * math.pi * i / n_frames for i in range(n_frames)]
        bgra_src = cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA)
        return [
            cv2.cvtColor(_perspective_card_frame(bgra_src, theta), cv2.COLOR_BGRA2RGBA)
            for theta in angles
        ]

    else:
        raise ValueError(f"Unknown mode: {mode!r}")


# ---------------------------------------------------------------------------
# Output saving
# ---------------------------------------------------------------------------

def _composite_over_solid(frame_rgba: np.ndarray, bg: tuple = (18, 18, 18)) -> np.ndarray:
    """Alpha-composite RGBA over a solid colour; return BGR uint8 for ffmpeg."""
    rgb  = frame_rgba[:, :, :3].astype(np.float32)
    a    = frame_rgba[:, :, 3:4].astype(np.float32) / 255.0
    bg_f = np.array(bg, dtype=np.float32)
    out  = (rgb * a + bg_f * (1.0 - a)).clip(0, 255).astype(np.uint8)
    return cv2.cvtColor(out, cv2.COLOR_RGB2BGR)


def save_orbit_output(
    frames: list,
    output_dir: str,
    fps: int = 24,
    bg_color: tuple = (18, 18, 18),
) -> dict:
    """
    Write orbit_frames/frame_NNN.png (RGBA) and orbit_preview.mp4.
    Returns dict with paths.
    """
    frames_dir = os.path.join(output_dir, "orbit_frames")
    os.makedirs(frames_dir, exist_ok=True)

    frame_paths = []
    for i, frame in enumerate(frames):
        path = os.path.join(frames_dir, f"frame_{i:03d}.png")
        cv2.imwrite(path, cv2.cvtColor(frame, cv2.COLOR_RGBA2BGRA))
        frame_paths.append(path)

    video_path = os.path.join(output_dir, "orbit_preview.mp4")
    _frames_to_mp4(frames, video_path, fps=fps, bg_color=bg_color)

    return {
        "frames_dir": frames_dir,
        "n_frames": len(frames),
        "video_path": video_path,
        "frame_paths": frame_paths,
    }


def _frames_to_mp4(
    frames: list, output_path: str, fps: int = 24, bg_color: tuple = (18, 18, 18)
) -> None:
    """Composite frames over solid bg, write MP4 via ffmpeg."""
    if not frames:
        return
    with tempfile.TemporaryDirectory(prefix="orbit_mp4_") as tmpdir:
        for i, frame in enumerate(frames):
            bgr = _composite_over_solid(frame, bg_color)
            cv2.imwrite(
                os.path.join(tmpdir, f"frame_{i:04d}.jpg"), bgr,
                [cv2.IMWRITE_JPEG_QUALITY, 95],
            )
        H, W = frames[0].shape[:2]
        W2, H2 = W - (W % 2), H - (H % 2)
        cmd = [
            "ffmpeg", "-y",
            "-framerate", str(fps),
            "-i", os.path.join(tmpdir, "frame_%04d.jpg"),
            "-vf", f"crop={W2}:{H2}:0:0",
            "-c:v", "libx264", "-pix_fmt", "yuv420p",
            "-crf", "18", "-movflags", "+faststart",
            output_path,
        ]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            raise RuntimeError(f"ffmpeg failed: {r.stderr[-600:]}")


# ---------------------------------------------------------------------------
# Debug helpers
# ---------------------------------------------------------------------------

def _save_debug(actor_rgba, actor_path, bg_rgba, debug_dir):
    os.makedirs(debug_dir, exist_ok=True)
    if os.path.exists(actor_path):
        shutil.copy2(actor_path, os.path.join(debug_dir, "selected_frame.png"))
    cv2.imwrite(os.path.join(debug_dir, "actor_rgba.png"),
                cv2.cvtColor(actor_rgba, cv2.COLOR_RGBA2BGRA))
    cv2.imwrite(os.path.join(debug_dir, "mask.png"), actor_rgba[:, :, 3])
    if bg_rgba is not None:
        cv2.imwrite(os.path.join(debug_dir, "bg_plate.png"),
                    cv2.cvtColor(bg_rgba, cv2.COLOR_RGBA2BGRA))


def _save_depth_debug(depth, debug_dir):
    os.makedirs(debug_dir, exist_ok=True)
    d8 = (depth * 255).astype(np.uint8)
    cv2.imwrite(os.path.join(debug_dir, "depth.png"), d8)
    cv2.imwrite(os.path.join(debug_dir, "depth_colorized.png"),
                cv2.applyColorMap(d8, cv2.COLORMAP_MAGMA))


# ---------------------------------------------------------------------------
# Full pipeline
# ---------------------------------------------------------------------------

def run_orbit_pipeline(
    image_path: str,
    output_dir: str,
    n_frames: int = 36,
    parallax_strength: float = 0.08,
    mode: str = "swing",
    fps: int = 24,
    max_angle_deg: float = 35.0,
    debug: bool = True,
) -> dict:
    """
    Full pipeline: load → bg-plate → depth → render → save.

    image_path: the specific image to orbit (caller selects; no sharpness heuristic)
    Returns dict: actor_path, frames_dir, video_path, n_frames, debug_dir, has_alpha, has_bg
    """
    os.makedirs(output_dir, exist_ok=True)
    debug_dir = os.path.join(output_dir, "debug")

    # 1. Load actor — prefer nobg sidecar for cleaner depth
    actor_path = image_path
    root, _ = os.path.splitext(image_path)
    nobg_candidate = root + ".nobg.png"
    if not root.endswith(".nobg") and os.path.exists(nobg_candidate):
        actor_path = nobg_candidate

    actor_rgba = _load_rgba(actor_path)
    has_alpha = _has_real_alpha(actor_rgba)

    # 2. Background plate (static reference — essential for perceivable depth)
    bg_rgba = get_bg_plate(actor_path, actor_rgba)
    has_bg = bg_rgba is not None

    if debug:
        _save_debug(actor_rgba, actor_path, bg_rgba, debug_dir)

    # 3. Depth map
    depth = create_depth_map(actor_rgba)

    if debug:
        _save_depth_debug(depth, debug_dir)

    # 4. Render
    frames = render_orbit(
        actor_rgba, depth,
        n_frames=n_frames,
        parallax_strength=parallax_strength,
        mode=mode,
        max_angle_deg=max_angle_deg,
        bg_rgba=bg_rgba,
    )

    # 5. Save
    result = save_orbit_output(frames, output_dir, fps=fps)
    result.update({
        "actor_path": actor_path,
        "debug_dir": debug_dir,
        "has_alpha": has_alpha,
        "has_bg": has_bg,
    })
    return result