"""
frame-compositor — GPU video compositor (ZeroGPU)
Bilinear quad warp matching JS editor preview exactly.
Optimizations: static layer composite cache, threaded audio+encode,
precomputed frame ranges, pinned memory transfers, parallel layer pre-render.
"""

import json
import os
import queue
import tempfile
import threading
import urllib.request
import urllib.parse
import re
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import List, Dict, Tuple, Optional

import subprocess
import gradio as gr
import spaces
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont
import av
from torchcodec.decoders import VideoDecoder

VIDEO_WIDTH  = 540
VIDEO_HEIGHT = 960
FPS          = 30
BATCH        = 32
DEVICE       = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ── Font cache ────────────────────────────────────────────────────────────────

_FONT_CACHE: Dict[str, str] = {}


def _download_font(family: str, weight: int = 400, cache_dir: str = "/tmp/fonts") -> Optional[str]:
    os.makedirs(cache_dir, exist_ok=True)
    key = f"{family}_{weight}"
    if key in _FONT_CACHE:
        return _FONT_CACHE[key]

    slug = family.replace(" ", "_")
    path = os.path.join(cache_dir, f"{slug}_{weight}.ttf")
    if os.path.exists(path):
        _FONT_CACHE[key] = path
        return path

    try:
        url = f"https://fonts.googleapis.com/css2?family={family.replace(' ', '+')}:ital,wght@0,{weight};1,{weight}&display=swap"
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=10) as r:
            css = r.read().decode("utf-8")

        ttf_url = None
        for block in re.split(r"@font-face\s*\{", css)[1:]:
            if re.search(rf"font-weight:\s*{weight}", block):
                m = re.search(r"url\((https://[^)]+\.(?:ttf|otf))\)", block)
                if m:
                    ttf_url = m.group(1)
                    break
        if not ttf_url:
            m = re.search(r"url\((https://[^)]+\.(?:ttf|otf))\)", css)
            if m:
                ttf_url = m.group(1)

        if not ttf_url:
            print(f"[font] No TTF/OTF URL found for {family} w{weight}", flush=True)
            return None

        with urllib.request.urlopen(ttf_url, timeout=10) as r:
            with open(path, "wb") as f:
                f.write(r.read())

        _FONT_CACHE[key] = path
        print(f"[font] Downloaded {family} w{weight} → {path}", flush=True)
        return path
    except Exception as e:
        print(f"[font] Failed for {family}: {e}", flush=True)
        return None


# ── Layer pre-rendering ───────────────────────────────────────────────────────

def _wrap_text(text: str, font: ImageFont.FreeTypeFont, max_w: int) -> List[str]:
    lines = []
    for paragraph in text.split("\n"):
        words = paragraph.split(" ")
        line = ""
        for word in words:
            test = f"{line} {word}".strip()
            bbox = font.getbbox(test)
            w = bbox[2] - bbox[0] if bbox else 0
            if w > max_w and line:
                lines.append(line)
                line = word
            else:
                line = test
        if line:
            lines.append(line)
    return lines or [""]


def _render_text_layer(layer: Dict) -> np.ndarray:
    W, H = 800, 200
    img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    draw = ImageDraw.Draw(img)

    bg_hex = layer["bgColor"]
    bg_r, bg_g, bg_b = int(bg_hex[1:3], 16), int(bg_hex[3:5], 16), int(bg_hex[5:7], 16)
    bg_a = int(layer["bgOpacity"] * 255)
    draw.rectangle([(0, 0), (W, H)], fill=(bg_r, bg_g, bg_b, bg_a))

    font_size = int(layer["fontSize"] * 2)
    weight = layer.get("fontWeight", 400)
    font_path = _download_font(layer["font"], weight)

    padding = 16
    avail_w = W - padding * 2
    avail_h = H - padding * 2

    fit_size = font_size
    fit_lines = [layer["text"]]
    for size in range(font_size, 7, -1):
        try:
            f = ImageFont.truetype(font_path, size) if font_path else ImageFont.load_default()
        except Exception:
            f = ImageFont.load_default()
        lines = _wrap_text(layer["text"], f, avail_w)
        if len(lines) * size * 1.2 <= avail_h:
            fit_size = size
            fit_lines = lines
            break

    try:
        font = ImageFont.truetype(font_path, fit_size) if font_path else ImageFont.load_default()
    except Exception:
        font = ImageFont.load_default()

    fg_hex = layer["color"]
    fg_r, fg_g, fg_b = int(fg_hex[1:3], 16), int(fg_hex[3:5], 16), int(fg_hex[5:7], 16)
    line_h = fit_size * 1.2
    total_h = len(fit_lines) * line_h
    y = (H - total_h) / 2 + line_h / 2

    for i, line in enumerate(fit_lines):
        draw.text((W // 2, y + i * line_h), line, font=font,
                  fill=(fg_r, fg_g, fg_b, 255), anchor="mm")

    return np.array(img)  # (H, W, 4) RGBA uint8


def _render_image_layer(layer: Dict) -> Optional[np.ndarray]:
    src = layer.get("src", "")
    if not src:
        return None
    try:
        req = urllib.request.Request(src, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=15) as r:
            data = r.read()
        img = Image.open(BytesIO(data)).convert("RGBA").resize((800, 800), Image.LANCZOS)
        opacity = layer.get("opacity", 1.0)
        if opacity < 1.0:
            r2, g2, b2, a2 = img.split()
            a2 = a2.point(lambda x: int(x * opacity))
            img = Image.merge("RGBA", (r2, g2, b2, a2))
        print(f"[layer] Loaded image {src}", flush=True)
        return np.array(img)
    except Exception as e:
        print(f"[layer] Failed to load image {src}: {e}", flush=True)
        return None


def _render_layer(layer: Dict) -> Tuple[str, Optional[np.ndarray]]:
    arr = _render_image_layer(layer) if layer.get("type") == "image" else _render_text_layer(layer)
    return layer["id"], arr


def _rgba_to_gpu_tensor(arr: np.ndarray) -> torch.Tensor:
    """(H, W, 4) RGBA uint8 → (1, 4, H, W) float32 [0,1] on DEVICE."""
    t = torch.from_numpy(arr).float() / 255.0
    return t.permute(2, 0, 1).unsqueeze(0).to(DEVICE)


# ── Keyframe interpolation ────────────────────────────────────────────────────

def _get_corners_at_frame(layer: Dict, f: int) -> Optional[List[Tuple[float, float]]]:
    kfs = sorted(layer["keyframes"], key=lambda k: k["frame"])
    if not kfs:
        return None
    if f <= kfs[0]["frame"]:
        return [(c["x"], c["y"]) for c in kfs[0]["corners"]]
    if f >= kfs[-1]["frame"]:
        return [(c["x"], c["y"]) for c in kfs[-1]["corners"]]
    for i in range(len(kfs) - 1):
        k0, k1 = kfs[i], kfs[i + 1]
        if k0["frame"] <= f <= k1["frame"]:
            ratio = (f - k0["frame"]) / (k1["frame"] - k0["frame"])
            return [
                (k0["corners"][j]["x"] + (k1["corners"][j]["x"] - k0["corners"][j]["x"]) * ratio,
                 k0["corners"][j]["y"] + (k1["corners"][j]["y"] - k0["corners"][j]["y"]) * ratio)
                for j in range(4)
            ]
    return [(c["x"], c["y"]) for c in kfs[0]["corners"]]


def _is_static_layer(layer: Dict) -> bool:
    kfs = layer.get("keyframes", [])
    if len(kfs) <= 1:
        return True
    ref = kfs[0]["corners"]
    for kf in kfs[1:]:
        for j in range(4):
            if abs(kf["corners"][j]["x"] - ref[j]["x"]) > 0.01:
                return False
            if abs(kf["corners"][j]["y"] - ref[j]["y"]) > 0.01:
                return False
    return True


# ── GPU bilinear quad warp ────────────────────────────────────────────────────

def _build_uv_map(
    corners_t: torch.Tensor,
    gx: torch.Tensor,
    gy: torch.Tensor,
    newton_iters: int = 8,
    u_init: Optional[torch.Tensor] = None,
    v_init: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    B = corners_t.shape[0]
    H, W = gx.shape

    tl_x = corners_t[:, 0, 0].view(B, 1, 1)
    tl_y = corners_t[:, 0, 1].view(B, 1, 1)
    tr_x = corners_t[:, 1, 0].view(B, 1, 1)
    tr_y = corners_t[:, 1, 1].view(B, 1, 1)
    bl_x = corners_t[:, 3, 0].view(B, 1, 1)
    bl_y = corners_t[:, 3, 1].view(B, 1, 1)
    br_x = corners_t[:, 2, 0].view(B, 1, 1)
    br_y = corners_t[:, 2, 1].view(B, 1, 1)

    px = gx.unsqueeze(0).expand(B, -1, -1)
    py = gy.unsqueeze(0).expand(B, -1, -1)

    u = u_init.clone() if u_init is not None else torch.full((B, H, W), 0.5, dtype=torch.float32, device=DEVICE)
    v = v_init.clone() if v_init is not None else torch.full((B, H, W), 0.5, dtype=torch.float32, device=DEVICE)

    for _ in range(newton_iters):
        top_x = tl_x + (tr_x - tl_x) * u
        top_y = tl_y + (tr_y - tl_y) * u
        bot_x = bl_x + (br_x - bl_x) * u
        bot_y = bl_y + (br_y - bl_y) * u
        fx = top_x + (bot_x - top_x) * v
        fy = top_y + (bot_y - top_y) * v

        dfx_du = (tr_x - tl_x) + ((br_x - bl_x) - (tr_x - tl_x)) * v
        dfy_du = (tr_y - tl_y) + ((br_y - bl_y) - (tr_y - tl_y)) * v
        dfx_dv = bot_x - top_x
        dfy_dv = bot_y - top_y

        det = dfx_du * dfy_dv - dfx_dv * dfy_du
        det = torch.where(det.abs() < 1e-8, torch.full_like(det, 1e-8), det)

        rx, ry = px - fx, py - fy
        u = u + (dfy_dv * rx - dfx_dv * ry) / det
        v = v + (dfx_du * ry - dfy_du * rx) / det

    u = u.clamp(0.0, 1.0)
    v = v.clamp(0.0, 1.0)

    top_x = tl_x + (tr_x - tl_x) * u
    top_y = tl_y + (tr_y - tl_y) * u
    bot_x = bl_x + (br_x - bl_x) * u
    bot_y = bl_y + (br_y - bl_y) * u
    fx = top_x + (bot_x - top_x) * v
    fy = top_y + (bot_y - top_y) * v
    inside = (((px - fx) ** 2 + (py - fy) ** 2).sqrt() < 1.5).float().unsqueeze(1)

    return u, v, inside


def _sample_layer(
    layer_t: torch.Tensor,
    u: torch.Tensor,
    v: torch.Tensor,
    inside: torch.Tensor,
    B: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    grid = torch.stack([u * 2 - 1, v * 2 - 1], dim=-1)
    sampled = F.grid_sample(
        layer_t.expand(B, -1, -1, -1),
        grid, mode="bilinear", padding_mode="zeros", align_corners=True,
    )
    return sampled[:, :3], sampled[:, 3:4] * inside


# ── Main render ───────────────────────────────────────────────────────────────

@spaces.GPU(duration=10)
def render(source_video_url: str, layers_json: str, progress=gr.Progress()) -> str:
    layers = json.loads(layers_json)
    print(f"[render] {len(layers)} layer(s), source: {source_video_url}", flush=True)

    progress(0.0, desc="Downloading source video...")
    tmp_in = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    req = urllib.request.Request(source_video_url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=120) as r:
        tmp_in.write(r.read())
    tmp_in.close()

    # ── Parallel layer pre-render ──────────────────────────────────────────
    progress(0.05, desc="Pre-rendering layers...")

    # Parallel layer rendering
    layer_arrays: Dict[str, np.ndarray] = {}
    with ThreadPoolExecutor(max_workers=min(len(layers), 4)) as ex:
        for lid, arr in ex.map(_render_layer, layers):
            if arr is not None:
                layer_arrays[lid] = arr
    print(f"[render] Pre-rendered {len(layer_arrays)} layer(s)", flush=True)

    # ── GPU setup ─────────────────────────────────────────────────────────
    progress(0.1, desc="Decoding video...")
    layer_tensors: Dict[str, torch.Tensor] = {
        lid: _rgba_to_gpu_tensor(arr) for lid, arr in layer_arrays.items()
    }

    decoder      = VideoDecoder(tmp_in.name, device=str(DEVICE))
    metadata     = decoder.metadata
    total_frames = metadata.num_frames
    actual_fps   = float(metadata.average_fps) if metadata.average_fps else FPS
    print(f"[render] {total_frames} frames @ {actual_fps:.2f} fps", flush=True)

    # ── Precompute pixel grid ──────────────────────────────────────────────
    gy_grid, gx_grid = torch.meshgrid(
        torch.arange(VIDEO_HEIGHT, dtype=torch.float32, device=DEVICE),
        torch.arange(VIDEO_WIDTH,  dtype=torch.float32, device=DEVICE),
        indexing="ij",
    )

    # ── Precompute per-layer frame ranges ─────────────────────────────────
    layer_frame_ranges: Dict[str, Tuple[int, int]] = {}
    for layer in layers:
        start_f = layer["startFrame"]
        end_f   = layer["endFrame"] + 1
        layer_frame_ranges[layer["id"]] = (start_f, end_f)

    # ── Static layer: precompute warped composite (rgb*alpha, alpha) once ──
    # For static layers the result is the same for every frame — skip grid_sample entirely.
    static_composite: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
    static_uv:        Dict[str, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}

    for layer in layers:
        lid = layer["id"]
        lt  = layer_tensors.get(lid)
        if lt is None:
            continue
        if not _is_static_layer(layer):
            continue
        corners   = _get_corners_at_frame(layer, layer["keyframes"][0]["frame"])
        corners_t = torch.tensor([corners], dtype=torch.float32, device=DEVICE)
        u, v, inside = _build_uv_map(corners_t, gx_grid, gy_grid)
        rgb, alpha   = _sample_layer(lt, u, v, inside, 1)
        # store pre-multiplied: rgb_pre = rgb*alpha, alpha — both (1, C, H, W)
        static_composite[lid] = (rgb * alpha, alpha)
        print(f"[render] Static composite cached for layer {lid}", flush=True)

    # ── Animated layer warm-start ──────────────────────────────────────────
    prev_uv: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}

    # ── Pinned output buffer ───────────────────────────────────────────────
    pinned_buf = torch.empty(
        (BATCH, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.uint8, pin_memory=True
    )

    # ── Encode thread ──────────────────────────────────────────────────────
    progress(0.15, desc="Rendering frames...")
    tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp_out.close()

    encode_q: queue.Queue = queue.Queue(maxsize=4)
    encode_error: list = []

    def _encode_worker():
        try:
            with av.open(tmp_out.name, "w") as dst:
                v_stream = dst.add_stream("libx264", rate=int(round(actual_fps)))
                v_stream.width   = VIDEO_WIDTH
                v_stream.height  = VIDEO_HEIGHT
                v_stream.pix_fmt = "yuv420p"
                v_stream.options = {"preset": "ultrafast", "crf": "23"}
                while True:
                    item = encode_q.get()
                    if item is None:
                        break
                    for frame_rgb in item:
                        av_frame = av.VideoFrame.from_ndarray(frame_rgb, format="rgb24")
                        for pkt in v_stream.encode(av_frame):
                            dst.mux(pkt)
                for pkt in v_stream.encode(None):
                    dst.mux(pkt)
        except Exception as e:
            encode_error.append(e)

    enc_thread = threading.Thread(target=_encode_worker, daemon=True)
    enc_thread.start()

    # ── GPU composite loop ─────────────────────────────────────────────────
    for batch_start in range(0, total_frames, BATCH):
        batch_end = min(batch_start + BATCH, total_frames)
        B = batch_end - batch_start

        clip     = decoder[batch_start:batch_end]
        frames_f = clip.data.to(DEVICE).float() / 255.0
        if frames_f.shape[2] != VIDEO_HEIGHT or frames_f.shape[3] != VIDEO_WIDTH:
            frames_f = F.interpolate(frames_f, size=(VIDEO_HEIGHT, VIDEO_WIDTH),
                                     mode="bilinear", align_corners=False)

        for layer in layers:
            lid = layer["id"]
            lt  = layer_tensors.get(lid)
            if lt is None:
                continue

            start_f, end_f = layer_frame_ranges[lid]
            # which indices in this batch are active
            active_idx = [i for i in range(B) if start_f <= batch_start + i < end_f]
            if not active_idx:
                continue

            if lid in static_composite:
                rgb_pre, alpha = static_composite[lid]  # (1, 3, H, W), (1, 1, H, W)
                if len(active_idx) == B:
                    frames_f = rgb_pre + frames_f * (1.0 - alpha)
                else:
                    sub = rgb_pre + frames_f[active_idx] * (1.0 - alpha)
                    frames_f = frames_f.clone()
                    for out_i, src_i in enumerate(active_idx):
                        frames_f[src_i] = sub[out_i]

            else:
                # animated — build UV per-frame with warm-start
                corners_list = []
                for i in active_idx:
                    corners = _get_corners_at_frame(layer, batch_start + i)
                    corners_list.append(corners)

                Ab        = len(active_idx)
                corners_t = torch.tensor(corners_list, dtype=torch.float32, device=DEVICE)

                u_init, v_init = prev_uv.get(lid, (None, None))
                if u_init is not None and u_init.shape[0] != Ab:
                    u_init = u_init[:Ab] if u_init.shape[0] > Ab else None
                    v_init = v_init[:Ab] if v_init is not None and v_init.shape[0] > Ab else None

                u, v, inside = _build_uv_map(corners_t, gx_grid, gy_grid, u_init=u_init, v_init=v_init)
                prev_uv[lid] = (u.detach(), v.detach())

                rgb, alpha = _sample_layer(lt, u, v, inside, Ab)

                if Ab == B:
                    frames_f = rgb * alpha + frames_f * (1.0 - alpha)
                else:
                    sub = rgb * alpha + frames_f[active_idx] * (1.0 - alpha)
                    frames_f = frames_f.clone()
                    for out_i, src_i in enumerate(active_idx):
                        frames_f[src_i] = sub[out_i]

        # (B, 3, H, W) float → pinned uint8 → CPU numpy via async DMA
        out_gpu = (frames_f.permute(0, 2, 3, 1).clamp(0, 1) * 255).byte()
        pinned_buf[:B].copy_(out_gpu, non_blocking=True)
        torch.cuda.synchronize()
        encode_q.put(pinned_buf[:B].numpy().copy())

        done = batch_end / total_frames
        progress(0.15 + done * 0.82, desc=f"Rendering {batch_end}/{total_frames} frames")
        print(f"[render] Processed {batch_end}/{total_frames} frames", flush=True)

    encode_q.put(None)
    enc_thread.join()

    if encode_error:
        raise encode_error[0]

    # ── Merge audio from source via ffmpeg stream copy ─────────────────────
    tmp_final = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp_final.close()
    result = subprocess.run([
        "ffmpeg", "-y",
        "-i", tmp_out.name,
        "-i", tmp_in.name,
        "-c:v", "copy",
        "-c:a", "copy",
        "-map", "0:v:0",
        "-map", "1:a:0?",
        "-shortest",
        tmp_final.name
    ], capture_output=True)
    if result.returncode != 0:
        print(f"[render] ffmpeg audio merge failed: {result.stderr.decode()}", flush=True)
        os.rename(tmp_out.name, tmp_final.name)
    else:
        os.unlink(tmp_out.name)

    os.unlink(tmp_in.name)

    progress(1.0, desc="Done")
    print(f"[render] Output: {tmp_final.name}", flush=True)
    return tmp_final.name


# ── Gradio app ────────────────────────────────────────────────────────────────

with gr.Blocks() as demo:
    with gr.Row(visible=False):
        source_video_url = gr.Textbox(label="source_video_url")
        layers_json      = gr.Textbox(label="layers_json")
    output_video = gr.File(label="output_video")
    gr.Button("render", visible=False).click(
        fn=render,
        inputs=[source_video_url, layers_json],
        outputs=[output_video],
        api_name="render",
    )

if __name__ == "__main__":
    demo.launch()