"""
orbit_qwen.py — near-real actor turntable using Qwen-Image-Edit.

Unlike orbit_module.py (fake 2.5D depth-card parallax), this actually asks the
generative model to RE-RENDER the subject at each yaw angle. Each view is
anchored to the original front image with a fixed seed so identity, body, hair
and lighting stay consistent while only the viewpoint rotates.

Pipeline:
  1. build a yaw-angle prompt per frame (turntable or swing)
  2. _run_pipeline (Qwen via ComfyUI) → one re-rendered view per angle
  3. bottom-center align onto a common canvas
  4. stitch to a looping MP4

Validated finding (2026-06-25): 2D blending between independently-generated
views (optical-flow morph OR crossfade) always ghosts — the bodies don't
overlap, so any in-between frame shows a double exposure. The cure is DENSITY,
not blending: ~24 crisp keyframes (15° steps) played with NO interpolation at
~12fps reads as a smooth turntable, exactly like classic 3D turntable GIFs.
Interpolation is kept available (interp_factor>1) but defaults OFF.

Reuses edit_api._run_pipeline, so it talks to the same running ComfyUI server.

Usage:
  from orbit_qwen import run_qwen_orbit
  result = run_qwen_orbit("/path/to/front.png", "/out/dir", n_views=12)

CLI: see orbit_qwen_poc.py
"""

import os
import io
import sys
import math
import subprocess
import tempfile

import cv2
import numpy as np
from PIL import Image

# Reuse the real Qwen pipeline from the API service (no server round-trip needed;
# _run_pipeline queues directly to ComfyUI). Import is cheap — only loads the
# workflow JSON; models load lazily and the uvicorn startup hook does not fire.
_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from edit_api import _run_pipeline, _load_output_dir, MAX_AREA  # noqa: E402

__all__ = [
    "yaw_prompt",
    "generate_views",
    "interpolate_views",
    "build_video",
    "run_qwen_orbit",
]


# ---------------------------------------------------------------------------
# 1. Prompt construction
# ---------------------------------------------------------------------------

# Identity lock appended to every angle — this is what keeps it "the same person".
_IDENTITY = (
    "exactly the same woman, identical face, identical body shape and proportions, "
    "same hair, same skin tone, same lighting, photorealistic, sharp focus, "
    "full body visible head to feet, centered, transparent background"
)


def _angle_phrase(deg: float) -> str:
    """
    Natural-language viewpoint for a yaw angle (turntable; subject rotates
    clockwise as deg increases). 0 = facing camera, 180 = facing away.
    """
    d = deg % 360
    # Bucket to the nearest named viewpoint for the clearest model instruction,
    # then add the precise degree as reinforcement.
    if d < 22.5 or d >= 337.5:
        view = "facing the camera directly, front view"
    elif d < 67.5:
        view = "turned slightly to her right, three-quarter front-right view"
    elif d < 112.5:
        view = "full right-side profile, body turned 90 degrees"
    elif d < 157.5:
        view = "three-quarter rear view from behind-right, back partially visible"
    elif d < 202.5:
        view = "facing directly away from the camera, full back view, back of head and back visible"
    elif d < 247.5:
        view = "three-quarter rear view from behind-left, back partially visible"
    elif d < 292.5:
        view = "full left-side profile, body turned 90 degrees"
    else:
        view = "turned slightly to her left, three-quarter front-left view"
    return view


def yaw_prompt(deg: float) -> str:
    """Full prompt for one turntable angle."""
    view = _angle_phrase(deg)
    return (
        f"Rotate the camera around the subject to a {int(deg % 360)} degree turntable angle: "
        f"{view}. The subject stands still in a neutral standing pose; only the viewing "
        f"angle changes, like a 3D turntable. {_IDENTITY}."
    )


def _angles_for(mode: str, n_views: int, sweep_deg: float) -> list:
    """Return the list of yaw angles to render."""
    if mode == "turntable":
        # Full 360, evenly spaced, loops cleanly
        return [360.0 * i / n_views for i in range(n_views)]
    elif mode == "swing":
        # -sweep/2 .. +sweep/2 .. back (front-facing arc only — most reliable)
        half = sweep_deg / 2.0
        fwd = [(-half + sweep_deg * i / (n_views - 1)) for i in range(n_views)]
        # map negatives into 0..360 turntable space (e.g. -45 -> 315)
        return [a % 360 for a in fwd]
    raise ValueError(f"Unknown mode: {mode!r}")


# ---------------------------------------------------------------------------
# 2. View generation (Qwen)
# ---------------------------------------------------------------------------

def _autocrop_alpha(pil: Image.Image, pad: int = 8) -> Image.Image:
    """Crop to the alpha bounding box (+pad) so every view is framed on the body."""
    if pil.mode != "RGBA":
        return pil
    alpha = np.array(pil)[:, :, 3]
    ys, xs = np.where(alpha > 16)
    if len(xs) == 0:
        return pil
    x0, x1 = max(0, xs.min() - pad), min(pil.width, xs.max() + pad)
    y0, y1 = max(0, ys.min() - pad), min(pil.height, ys.max() + pad)
    return pil.crop((x0, y0, x1, y1))


def generate_views(
    image_path: str,
    output_dir: str,
    n_views: int = 12,
    seed: int = 42,
    mode: str = "turntable",
    sweep_deg: float = 180.0,
    anchor: str = "original",
    max_area: int = 0,
    steps: int = 8,
    on_progress=None,
) -> list:
    """
    Render one Qwen view per yaw angle.

    anchor='original' — every view edits the SAME front image (stable identity)
    anchor='chain'    — each view edits the previous result (smoother transitions,
                        but identity can drift over a full turn)

    Returns list of dicts: {deg, path, pil}.
    """
    os.makedirs(output_dir, exist_ok=True)
    views_dir = os.path.join(output_dir, "views")
    os.makedirs(views_dir, exist_ok=True)

    base_pil = Image.open(image_path).convert("RGB")
    angles = _angles_for(mode, n_views, sweep_deg)

    results = []
    prev_pil = None
    for i, deg in enumerate(angles):
        src_pil = base_pil if anchor == "original" or prev_pil is None else prev_pil
        prompt = yaw_prompt(deg)
        if on_progress:
            on_progress(i, len(angles), deg)

        png = _run_pipeline(
            src_pil, prompt, seed,
            max_area or MAX_AREA,
            steps=steps,
        )
        view_pil = Image.open(io.BytesIO(png)).convert("RGBA")
        view_pil = _autocrop_alpha(view_pil)

        path = os.path.join(views_dir, f"view_{i:03d}_{int(deg):03d}deg.png")
        view_pil.save(path)
        results.append({"deg": deg, "path": path, "pil": view_pil})

        if anchor == "chain":
            # Feed an RGB version forward (pipeline wants RGB anyway)
            prev_pil = view_pil.convert("RGB")

    return results


# ---------------------------------------------------------------------------
# 3. Smoothing — canvas-align + optical-flow interpolation
# ---------------------------------------------------------------------------

def _to_common_canvas(views: list, pad_frac: float = 0.12) -> list:
    """
    Place every view on one fixed-size RGBA canvas, bottom-centered (feet anchored),
    so the body doesn't jump frame-to-frame. Returns list of HxWx4 uint8 arrays.
    """
    H = max(v["pil"].height for v in views)
    W = max(v["pil"].width for v in views)
    padH, padW = int(H * pad_frac), int(W * pad_frac)
    CH, CW = H + 2 * padH, W + 2 * padW

    out = []
    for v in views:
        p = v["pil"]
        canvas = Image.new("RGBA", (CW, CH), (0, 0, 0, 0))
        # bottom-centered: feet sit on a common baseline
        x = (CW - p.width) // 2
        y = CH - padH - p.height
        canvas.paste(p, (x, y), p)
        out.append(np.array(canvas))
    return out


def _flow_morph_rgb(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
    """
    Optical-flow morph between two SOLID RGB frames (3-channel) at fraction t.
    Operates on composited-over-bg images so there is no alpha halo/ghost.
    Warps a→mid and b→mid, then blends.
    """
    ag = cv2.cvtColor(a, cv2.COLOR_RGB2GRAY)
    bg = cv2.cvtColor(b, cv2.COLOR_RGB2GRAY)
    flow_ab = cv2.calcOpticalFlowFarneback(ag, bg, None, 0.5, 5, 31, 5, 7, 1.5, 0)
    flow_ba = cv2.calcOpticalFlowFarneback(bg, ag, None, 0.5, 5, 31, 5, 7, 1.5, 0)

    H, W = ag.shape
    yc, xc = np.mgrid[0:H, 0:W].astype(np.float32)
    wa = cv2.remap(a, (xc + flow_ab[..., 0] * t), (yc + flow_ab[..., 1] * t),
                   cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    wb = cv2.remap(b, (xc + flow_ba[..., 0] * (1 - t)), (yc + flow_ba[..., 1] * (1 - t)),
                   cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    return (wa.astype(np.float32) * (1 - t) + wb.astype(np.float32) * t).clip(0, 255).astype(np.uint8)


def interpolate_views(
    views: list,
    factor: int = 4,
    loop: bool = True,
    smooth: bool = True,
    bg: tuple = (18, 18, 18),
) -> list:
    """
    Expand keyframes into a smooth sequence.

    Keyframes are first composited over the solid bg, so all blending happens
    in opaque RGB space — this removes the transparent-alpha ghosting that
    plagued earlier flow morphs.

    factor — intermediate frames per keyframe pair (1 = keyframes only)
    loop   — also blend last→first (seamless turntable)
    smooth — optical-flow morph (True) vs simple crossfade (False)

    Returns list of HxWx3 uint8 RGB frames.
    """
    canvases = _to_common_canvas(views)
    bg_arr = np.array(bg, dtype=np.float32)

    def _flatten(rgba):
        a = rgba[:, :, 3:4].astype(np.float32) / 255.0
        return (rgba[:, :, :3].astype(np.float32) * a + bg_arr * (1 - a)).clip(0, 255).astype(np.uint8)

    solid = [_flatten(c) for c in canvases]
    if factor <= 1:
        return solid

    n = len(solid)
    pairs = n if loop else n - 1
    frames = []
    for i in range(pairs):
        a, b = solid[i], solid[(i + 1) % n]
        frames.append(a)
        for k in range(1, factor):
            t = k / factor
            if smooth:
                frames.append(_flow_morph_rgb(a, b, t))
            else:
                frames.append((a.astype(np.float32) * (1 - t) +
                               b.astype(np.float32) * t).astype(np.uint8))
    if not loop:
        frames.append(solid[-1])
    return frames


# ---------------------------------------------------------------------------
# 4. Video
# ---------------------------------------------------------------------------

def _composite_solid(frame: np.ndarray, bg=(18, 18, 18)) -> np.ndarray:
    """Accept RGB (already flattened) or RGBA; return BGR for ffmpeg."""
    if frame.shape[2] == 3:
        return cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    rgb = frame[:, :, :3].astype(np.float32)
    a = frame[:, :, 3:4].astype(np.float32) / 255.0
    bg_f = np.array(bg, dtype=np.float32)
    out = (rgb * a + bg_f * (1 - a)).clip(0, 255).astype(np.uint8)
    return cv2.cvtColor(out, cv2.COLOR_RGB2BGR)


def build_video(frames: list, output_path: str, fps: int = 24, bg=(18, 18, 18)) -> None:
    if not frames:
        return
    with tempfile.TemporaryDirectory(prefix="orbit_qwen_") as tmp:
        for i, fr in enumerate(frames):
            cv2.imwrite(os.path.join(tmp, f"f_{i:04d}.jpg"),
                        _composite_solid(fr, bg), [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(tmp, "f_%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:]}")


# ---------------------------------------------------------------------------
# 5. Orchestration
# ---------------------------------------------------------------------------

def run_qwen_orbit(
    image_path: str,
    output_dir: str,
    n_views: int = 24,
    seed: int = 42,
    mode: str = "turntable",
    sweep_deg: float = 180.0,
    anchor: str = "original",
    interp_factor: int = 1,
    smooth: bool = False,
    fps: int = 12,
    max_area: int = 0,
    steps: int = 8,
    on_progress=None,
) -> dict:
    """
    Full near-real turntable: generate Qwen views → align → MP4.

    Defaults reflect the validated recipe: 24 crisp keyframes, NO blending,
    12fps. Raise interp_factor only if you accept morph ghosting.

    Returns dict: views (list), n_views, n_frames, video_path, views_dir.
    """
    os.makedirs(output_dir, exist_ok=True)

    views = generate_views(
        image_path, output_dir,
        n_views=n_views, seed=seed, mode=mode, sweep_deg=sweep_deg,
        anchor=anchor, max_area=max_area, steps=steps, on_progress=on_progress,
    )

    loop = (mode == "turntable")
    frames = interpolate_views(views, factor=interp_factor, loop=loop, smooth=smooth)

    video_path = os.path.join(output_dir, "turntable.mp4")
    build_video(frames, video_path, fps=fps)

    return {
        "views": [{"deg": v["deg"], "path": v["path"]} for v in views],
        "n_views": len(views),
        "n_frames": len(frames),
        "video_path": video_path,
        "views_dir": os.path.join(output_dir, "views"),
    }