import gradio as gr
import spaces
import subprocess
import asyncio
import json
import base64
import os
import shutil
import socket
import time
import io

import aiohttp
import websockets


CHROMIUM_PATHS = [
    "/usr/bin/chromium",
    "/usr/bin/chromium-browser",
    "/usr/bin/google-chrome-stable",
    "/usr/bin/google-chrome",
]

AUDIO_NORMALIZE = "aformat=sample_rates=44100:channel_layouts=stereo"
TARGET_FPS = 25


def get_chromium_path():
    for path in CHROMIUM_PATHS:
        if os.path.exists(path):
            return path
    raise RuntimeError(f"Chromium not found. Tried: {CHROMIUM_PATHS}")


def get_free_port():
    with socket.socket() as s:
        s.bind(("", 0))
        return s.getsockname()[1]


def get_media_duration(path: str) -> float | None:
    result = subprocess.run(
        ["ffprobe", "-v", "quiet",
         "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1",
         path],
        capture_output=True, text=True,
    )
    try:
        return float(result.stdout.strip())
    except ValueError:
        return None


def has_audio_stream(path: str) -> bool:
    result = subprocess.run(
        ["ffprobe", "-v", "quiet",
         "-show_streams", "-select_streams", "a:0",
         "-of", "default=noprint_wrappers=1",
         path],
        capture_output=True, text=True,
    )
    return bool(result.stdout.strip())


@spaces.GPU(duration=1)
def _gpu_warmup():
    pass


def _encode_frames(
    frame_bytes: list[bytes],
    timestamps: list[float],
    duration: float,
    w: int,
    h: int,
) -> str:
    """
    Decode PNG frames in memory, encode to H.264 with PyAV (libx264).
    No resize needed — Chrome always delivers frames at exactly w×h via
    Emulation.setDeviceMetricsOverride + screencast maxWidth/maxHeight.
    """
    import av
    import numpy as np
    from PIL import Image

    t0 = time.monotonic()
    print(f"[wr] encode: {len(frame_bytes)} frames → {w}x{h} @ {TARGET_FPS}fps", flush=True)

    # Select frames evenly at TARGET_FPS via binary search on timestamps
    n_out = max(1, int(duration * TARGET_FPS))
    selected = []
    for i in range(n_out):
        t = i / TARGET_FPS
        lo, hi = 0, len(timestamps) - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if timestamps[mid] < t:
                lo = mid + 1
            else:
                hi = mid
        selected.append(lo)
    print(f"[wr] encode: selected {len(selected)} frames", flush=True)

    out_path = f"/tmp/main_{os.getpid()}_{int(time.time())}.mp4"
    container = av.open(out_path, "w")
    stream = container.add_stream("libx264", rate=TARGET_FPS)
    stream.width = w
    stream.height = h
    stream.pix_fmt = "yuv420p"
    stream.options = {"crf": "15", "preset": "fast", "profile": "high", "level": "4.2"}

    for i, frame_idx in enumerate(selected):
        img = Image.open(io.BytesIO(frame_bytes[frame_idx])).convert("RGB")
        av_frame = av.VideoFrame.from_ndarray(np.array(img), format="rgb24")
        for pkt in stream.encode(av_frame):
            container.mux(pkt)
        if i % 100 == 0:
            print(f"[wr] encode: {i}/{len(selected)} frames", flush=True)

    for pkt in stream.encode():
        container.mux(pkt)
    container.close()

    print(f"[wr] encode: done in {time.monotonic()-t0:.1f}s → {out_path}", flush=True)
    return out_path


def _run_ffmpeg(cmd: list[str]):
    """Run an FFmpeg command, streaming stderr to stdout. Raises gr.Error on failure."""
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, bufsize=0)
    chunks = []
    while True:
        chunk = proc.stderr.read(512)
        if not chunk:
            break
        chunks.append(chunk)
        print(chunk.decode("utf-8", errors="replace"), end="", flush=True)
    proc.wait()
    if proc.returncode != 0:
        stderr = b"".join(chunks).decode("utf-8", errors="replace")
        raise gr.Error(f"FFmpeg failed: {stderr[-2000:]}")


def _assemble_final(
    main_path: str,
    output_path: str,
    w: int,
    h: int,
    duration: float,
    audio_path: str | None,
    pre_rec_path: str | None,
    post_rec_path: str | None,
    postroll_path: str | None,
    watermark_path: str | None,
):
    """
    Final assembly via FFmpeg:
    - main_path is already-encoded H.264 video (no audio)
    - stream-copy main video when there are no segments to concat
    - re-encode only when mixing multiple video segments or adding watermark
    """
    has_segments = any([pre_rec_path, post_rec_path, postroll_path])
    has_watermark = watermark_path is not None

    # ── Simple: no segments ───────────────────────────────────────────────────
    if not has_segments and not audio_path and not has_watermark:
        shutil.copy2(main_path, output_path)
        return

    if not has_segments and not has_watermark:
        # Stream-copy video, just add audio
        _run_ffmpeg([
            "ffmpeg", "-y",
            "-i", main_path,
            "-i", audio_path,
            "-c:v", "copy",
            "-c:a", "aac", "-b:a", "192k",
            "-shortest",
            "-movflags", "+faststart",
            output_path,
        ])
        return

    # ── Simple with watermark only ─────────────────────────────────────────────
    if not has_segments and has_watermark:
        # Watermark overlay requires re-encoding
        inputs = ["-i", main_path, "-i", watermark_path]
        cmd = ["ffmpeg", "-y"] + inputs

        # Loop watermark and overlay with proper alpha handling (eval=frame for per-frame alpha)
        filter_parts = [f"[1:v]loop=loop=-1:size=1:start=0[wm];[0:v][wm]overlay=0:0:eval=frame:alpha=straight:shortest=1[outv]"]

        cmd += ["-filter_complex", ";".join(filter_parts), "-map", "[outv]"]

        if audio_path:
            cmd += ["-i", audio_path, "-map", "2:a", "-c:a", "aac", "-b:a", "192k", "-shortest"]

        cmd += [
            "-c:v", "libx264", "-crf", "15", "-preset", "fast",
            "-pix_fmt", "yuv420p", "-profile:v", "high", "-level:v", "4.2",
            "-r", "25",
            "-movflags", "+faststart",
            output_path,
        ]
        _run_ffmpeg(cmd)
        return

    # ── Complex: multiple segments ────────────────────────────────────────────
    scale_video = (
        f"scale={w}:{h}:force_original_aspect_ratio=decrease,"
        f"pad={w}:{h}:(ow-iw)/2:(oh-ih)/2,setsar=1"
    )

    inputs = ["-i", main_path]
    next_idx = 1

    audio_idx = None
    if audio_path:
        inputs += ["-i", audio_path]
        audio_idx = next_idx
        next_idx += 1

    watermark_idx = None
    if watermark_path:
        inputs += ["-i", watermark_path]
        watermark_idx = next_idx
        next_idx += 1

    video_files = []
    for label, path in [("pre", pre_rec_path), ("postrec", post_rec_path), ("postroll", postroll_path)]:
        if path:
            inputs += ["-i", path]
            video_files.append((label, path, next_idx))
            next_idx += 1

    pre_info      = next((v for v in video_files if v[0] == "pre"),      None)
    postrec_info  = next((v for v in video_files if v[0] == "postrec"),  None)
    postroll_info = next((v for v in video_files if v[0] == "postroll"), None)

    def probe(info):
        if not info:
            return None
        _, path, _ = info
        return {
            "has_audio": has_audio_stream(path),
            "duration":  get_media_duration(path) or 0.0,
        }

    pre_meta      = probe(pre_info)
    postrec_meta  = probe(postrec_info)
    postroll_meta = probe(postroll_info)

    segments = []

    if pre_info:
        _, _, idx = pre_info
        segments.append({
            "vf":        f"[{idx}:v]{scale_video}",
            "has_audio": pre_meta["has_audio"],
            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
            "duration":  pre_meta["duration"],
        })

    segments.append({
        "vf":        "[0:v]fps=25,setsar=1,setpts=PTS-STARTPTS",
        "has_audio": bool(audio_path),
        "af_real":   (
            f"[{audio_idx}:a]{AUDIO_NORMALIZE},atrim=duration={duration:.4f},asetpts=PTS-STARTPTS"
            if audio_path else None
        ),
        "duration":  duration,
    })

    if postrec_info:
        _, _, idx = postrec_info
        segments.append({
            "vf":        f"[{idx}:v]{scale_video}",
            "has_audio": postrec_meta["has_audio"],
            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
            "duration":  postrec_meta["duration"],
        })

    if postroll_info:
        _, _, idx = postroll_info
        segments.append({
            "vf":        f"[{idx}:v]{scale_video}",
            "has_audio": postroll_meta["has_audio"],
            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
            "duration":  postroll_meta["duration"],
        })

    need_audio_out = any(s["has_audio"] for s in segments)

    fp = []
    v_labels = []
    a_labels = []

    for i, seg in enumerate(segments):
        vl = f"v{i}"
        al = f"a{i}"
        fp.append(f"{seg['vf']}[{vl}]")
        v_labels.append(f"[{vl}]")

        if need_audio_out:
            if seg["has_audio"]:
                fp.append(f"{seg['af_real']}[{al}]")
            else:
                d = f"{seg['duration']:.4f}"
                fp.append(f"anullsrc=r=44100:cl=stereo,atrim=duration={d},asetpts=PTS-STARTPTS[{al}]")
            a_labels.append(f"[{al}]")

    n = len(segments)
    fp.append(f"{''.join(v_labels)}concat=n={n}:v=1:a=0[concatv]")
    if need_audio_out:
        fp.append(f"{''.join(a_labels)}concat=n={n}:v=0:a=1[outa]")

    # Apply watermark overlay if provided
    if watermark_idx is not None:
        # Loop watermark and overlay with proper alpha handling (eval=frame for per-frame alpha)
        fp.append(f"[{watermark_idx}:v]loop=loop=-1:size=1:start=0[wm]")
        fp.append(f"[concatv][wm]overlay=0:0:eval=frame:alpha=straight:shortest=1[outv]")
    else:
        fp.append(f"[concatv]null[outv]")

    cmd = ["ffmpeg"] + inputs + ["-filter_complex", ";".join(fp), "-map", "[outv]"]
    if need_audio_out:
        cmd += ["-map", "[outa]", "-c:a", "aac", "-b:a", "192k"]
    cmd += [
        "-c:v", "libx264", "-crf", "15", "-preset", "fast",
        "-pix_fmt", "yuv420p", "-profile:v", "high", "-level:v", "4.2",
        "-r", "25",
        "-movflags", "+faststart", "-y", output_path,
    ]
    _run_ffmpeg(cmd)


def _cp(path: str | None) -> str | None:
    """Copy an uploaded file to /tmp before asyncio yields give Gradio a cleanup opportunity."""
    if not path:
        return None
    dst = f"/tmp/wr_{os.getpid()}_{int(time.time() * 1000)}_{os.path.basename(path)}"
    shutil.copy2(path, dst)
    return dst


async def _record_async(
    url: str, w: int, h: int, duration: int,
    output_path: str,
    audio_path: str | None,
    pre_rec_path: str | None,
    post_rec_path: str | None,
    postroll_path: str | None,
    watermark_path: str | None,
):
    t0 = time.monotonic()
    def log(msg):
        print(f"[wr {time.monotonic()-t0:.1f}s] {msg}", flush=True)

    port = get_free_port()
    chromium = get_chromium_path()
    log(f"launching chromium on port {port}, url={url}, size={w}x{h}, duration={duration}")

    proc = subprocess.Popen(
        [
            chromium,
            f"--remote-debugging-port={port}",
            "--headless",
            "--no-sandbox",
            "--disable-dev-shm-usage",
            "--disable-setuid-sandbox",
            "--use-gl=swiftshader",
            "--enable-webgl",
            "--hide-scrollbars",
            "--autoplay-policy=no-user-gesture-required",
            f"--window-size={w},{h}",
        ],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

    try:
        ws_url = None
        async with aiohttp.ClientSession() as session:
            for attempt in range(20):
                await asyncio.sleep(0.5)
                try:
                    async with session.get(f"http://localhost:{port}/json/list") as resp:
                        pages = await resp.json()
                        if pages:
                            ws_url = pages[0]["webSocketDebuggerUrl"]
                            log(f"chrome devtools ready after {attempt+1} attempts")
                            break
                except Exception as e:
                    log(f"devtools attempt {attempt+1} failed: {e}")
                    continue

        if not ws_url:
            raise gr.Error("Could not connect to Chrome DevTools.")

        # frames: list of (timestamp_seconds, raw_png_bytes)
        frames: list[tuple[float, bytes]] = []
        msg_id = 1

        async with websockets.connect(ws_url, max_size=50_000_000) as ws:

            async def send(method, params=None):
                nonlocal msg_id
                current_id = msg_id
                await ws.send(json.dumps({"id": current_id, "method": method, "params": params or {}}))
                msg_id += 1
                return current_id

            await send("Page.enable")
            await send("Emulation.setDeviceMetricsOverride", {
                "width": w, "height": h, "deviceScaleFactor": 1, "mobile": False,
            })

            log(f"navigating to {url}")
            nav_id = await send("Page.navigate", {"url": url})
            deadline = time.monotonic() + 10
            while time.monotonic() < deadline:
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
                    if json.loads(raw).get("id") == nav_id:
                        log("Page.navigate acknowledged")
                        break
                except asyncio.TimeoutError:
                    continue

            log("waiting for Page.loadEventFired (max 15s)")
            deadline = time.monotonic() + 15
            while time.monotonic() < deadline:
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
                    if json.loads(raw).get("method") == "Page.loadEventFired":
                        log("Page.loadEventFired received")
                        break
                except asyncio.TimeoutError:
                    continue
            else:
                log("load timeout — starting screencast anyway")

            await send("Runtime.enable")
            await send("Runtime.evaluate", {
                "expression": (
                    "document.querySelectorAll('audio,video')"
                    ".forEach(m => { m.muted = false; m.play().catch(() => {}); })"
                ),
                "awaitPromise": False,
            })

            log(f"starting screencast for {duration}s")
            await send("Page.startScreencast", {
                "format": "png",
                "maxWidth": w,
                "maxHeight": h,
                "everyNthFrame": 1,
            })

            start = time.monotonic()
            last_log = start
            while time.monotonic() - start < duration:
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
                    obj = json.loads(raw)
                    if obj.get("method") == "Page.screencastFrame":
                        ts = time.monotonic() - start
                        # Decode base64 to raw bytes immediately — avoids storing base64 overhead
                        frames.append((ts, base64.b64decode(obj["params"]["data"])))
                        await send("Page.screencastFrameAck", {
                            "sessionId": obj["params"]["sessionId"]
                        })
                        now = time.monotonic()
                        if now - last_log >= 5:
                            log(f"recording... {ts:.1f}s elapsed, {len(frames)} frames")
                            last_log = now
                except asyncio.TimeoutError:
                    continue

            log(f"screencast done: {len(frames)} frames captured")
            await send("Page.stopScreencast")

    finally:
        proc.terminate()
        proc.wait()
        log("chromium terminated")

    if not frames:
        raise gr.Error("No frames captured. Check the URL.")

    timestamps  = [f[0] for f in frames]
    frame_bytes = [f[1] for f in frames]

    log(f"encoding {len(frames)} frames (PyAV/libx264, no disk I/O)...")
    main_path = _encode_frames(frame_bytes, timestamps, float(duration), w, h)
    log(f"main recording encoded: {main_path}")

    log("assembling final video (audio + segments + watermark)...")
    _assemble_final(main_path, output_path, w, h, float(duration),
                    audio_path, pre_rec_path, post_rec_path, postroll_path, watermark_path)

    try:
        os.unlink(main_path)
    except OSError:
        pass

    log("complete")
    return output_path


async def record(
    url: str,
    width: float,
    height: float,
    duration: int,
    audio_file: str | None,
    pre_rec_file: str | None,
    post_rec_file: str | None,
    postroll_file: str | None,
    watermark_file: str | None,
) -> str:
    if not url.strip():
        raise gr.Error("Please enter a URL.")
    if not url.startswith(("http://", "https://")):
        url = "https://" + url

    # Copy uploaded files synchronously before the first await, so Gradio's
    # event-loop cleanup (triggered at await yield points) can't delete them.
    audio_file    = _cp(audio_file)
    pre_rec_file  = _cp(pre_rec_file)
    post_rec_file = _cp(post_rec_file)
    postroll_file = _cp(postroll_file)
    watermark_file = _cp(watermark_file)
    print(f"[wr] files copied: audio={audio_file} pre={pre_rec_file} post={post_rec_file} postroll={postroll_file} watermark={watermark_file}", flush=True)

    w = max(2, (int(width)  // 2) * 2)
    h = max(2, (int(height) // 2) * 2)

    if audio_file:
        audio_dur = get_media_duration(audio_file)
        if audio_dur:
            duration = int(min(audio_dur, 120))

    output_path = f"/tmp/recording_{os.getpid()}_{int(time.time())}.mp4"
    return await _record_async(
        url, w, h, duration, output_path,
        audio_file, pre_rec_file, post_rec_file, postroll_file, watermark_file,
    )


with gr.Blocks(title="Website Renderer") as demo:
    gr.Markdown(
        "# Website Renderer\n"
        "Recording starts after full page load. "
        "Uploading audio overrides the duration slider (max 2 min)."
    )

    url_input = gr.Textbox(label="URL", placeholder="https://example.com")

    with gr.Row():
        width_input    = gr.Number(label="Width",  value=540, precision=0, minimum=100, maximum=3840)
        height_input   = gr.Number(label="Height", value=960, precision=0, minimum=100, maximum=3840)
        duration_input = gr.Slider(minimum=5, maximum=120, value=30, step=5, label="Duration (seconds)")

    with gr.Row():
        audio_input = gr.Audio(
            label="Audio track (optional — overrides duration)",
            type="filepath",
            sources=["upload"],
        )

    with gr.Row():
        pre_rec_input  = gr.File(
            label="Pre-recording video (optional)",
            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
        )
        post_rec_input = gr.File(
            label="Post-recording video (optional)",
            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
        )
        postroll_input = gr.File(
            label="Post-roll video (optional)",
            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
        )

    watermark_input = gr.File(
        label="Watermark (optional — PNG with transparency)",
        file_types=[".png"],
    )

    record_btn   = gr.Button("Record", variant="primary")
    video_output = gr.Video(label="Recorded Video")

    record_btn.click(
        fn=record,
        inputs=[
            url_input, width_input, height_input, duration_input,
            audio_input, pre_rec_input, post_rec_input, postroll_input, watermark_input,
        ],
        outputs=video_output,
    )

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