#!/usr/bin/env python3
"""Gradio wrapper for Molecular Gigafactory GPU rendering on HF Space.
Auto-starts render on boot; persists progress via checkpoint file.
Uploads final MP4 to HF Dataset repo so it survives Space restarts.
"""

import gradio as gr
import subprocess
import threading
import json
import os
import sys
import time
import shutil
from pathlib import Path
from huggingface_hub import HfApi, upload_file, create_repo

OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/output")
PROGRESS_FILE = os.path.join(OUTPUT_DIR, "progress.json")
STATUS_FILE = os.path.join(OUTPUT_DIR, "status.json")
HF_DATASET = "MolecularReality/mol-giga-output"
HF_TOKEN = os.environ.get("HF_TOKEN")

PHASES = ["A", "B", "C", "D", "E"]
RENDER_RES = "480"
RENDER_SAMPLES = "16"

# ---------------------------------------------------------------------------
# Render engine
# ---------------------------------------------------------------------------
def run_blender_phase(phase):
    """Run Blender for a single phase. Returns True on success.
    Streams output in real-time so logs are visible immediately."""
    phase_dir = os.path.join(OUTPUT_DIR, f"phase_{phase}")
    os.makedirs(phase_dir, exist_ok=True)

    cmd = [
        "blender", "--background",
        "--python", "/app/scene.py", "--",
        "--render-all", f"--phase={phase}",
        f"--output={OUTPUT_DIR}",
        f"--res={RENDER_RES}",
        f"--samples={RENDER_SAMPLES}",
        "--cpu",
    ]

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        universal_newlines=True,
    )

    stdout_lines = []
    noisy_markers = ('Synchronizing object', 'Building BVH', 'Building OptiX',
                     'Updating Mesh', 'Updating Objects', 'Updating Primitive',
                     'Updating Images', 'Updating Camera', 'Updating Lookup',
                     'Updating Lights', 'Updating Integrator', 'Updating Film',
                     'Updating Baking', 'Updating Device', 'Copying BVH',
                     'Copying Mesh', 'Computing normals', 'Computing distribution',
                     'Computing tree', 'Writing constant memory')
    log_path = os.path.join(OUTPUT_DIR, f"blender_{phase}.log")
    def log_output():
        with open(log_path, "w") as logf:
            for line in proc.stdout:
                stripped = line.rstrip()
                stdout_lines.append(line)
                logf.write(line)
                logf.flush()
                if any(m in stripped for m in noisy_markers):
                    continue
                print(f"[Blender-{phase}] {stripped}")

    reader = threading.Thread(target=log_output, daemon=True)
    reader.start()

    try:
        returncode = proc.wait(timeout=7200)
    except subprocess.TimeoutExpired:
        proc.kill()
        reader.join(timeout=5)
        print(f"Phase {phase} TIMEOUT after 7200s")
        return False

    reader.join(timeout=5)
    full_output = ''.join(stdout_lines)

    if returncode != 0:
        err_msg = full_output[-2000:] if full_output else "(no output)"
        print(f"Phase {phase} FAILED (exit {returncode}):\n{err_msg}")
        save_status({"phase": phase, "progress": 0,
                     "message": f"Phase {phase} FAILED: {err_msg[-300:]}"})
        return False

    # Encode phase frames to MP4
    phase_mp4 = os.path.join(OUTPUT_DIR, f"phase_{phase}.mp4")
    # Find any frame PNG — phases start at different frame numbers
    png_files = sorted([f for f in os.listdir(phase_dir) if f.endswith('.png')])
    if png_files:
        print(f"Encoding {len(png_files)} frames to {phase_mp4}...")
        # Build ffmpeg input from actual file list (handles non-0001 start frames)
        enc = subprocess.run([
            "ffmpeg", "-y", "-framerate", "24",
            "-pattern_type", "glob", "-i", os.path.join(phase_dir, "frame_*.png"),
            "-c:v", "libx264", "-preset", "fast", "-crf", "18",
            "-pix_fmt", "yuv420p",
            phase_mp4,
        ], capture_output=True, text=True, timeout=600)
        if enc.returncode != 0:
            print(f"ffmpeg phase {phase} FAILED:\n{enc.stderr[-1000:]}")
            return False

    return os.path.exists(phase_mp4)


def encode_final():
    """Concat phase MP4s into final movie."""
    concat_list = os.path.join(OUTPUT_DIR, "concat.txt")
    mp4_paths = []
    for phase in PHASES:
        mp4 = os.path.join(OUTPUT_DIR, f"phase_{phase}.mp4")
        if os.path.exists(mp4):
            mp4_paths.append(mp4)

    if not mp4_paths:
        return None

    with open(concat_list, "w") as f:
        for mp4 in mp4_paths:
            f.write(f"file '{mp4}'\n")

    final_mp4 = os.path.join(OUTPUT_DIR, "Molecular-Gigafactory.mp4")
    subprocess.run([
        "ffmpeg", "-y", "-f", "concat", "-safe", "0",
        "-i", concat_list, "-c", "copy", final_mp4,
    ], capture_output=True, text=True, timeout=300)

    return final_mp4 if os.path.exists(final_mp4) else None


def upload_result(mp4_path):
    """Upload final MP4 to HF Dataset."""
    if not HF_TOKEN:
        print("WARNING: HF_TOKEN not set, cannot upload")
        return False
    try:
        api = HfApi(token=HF_TOKEN)
        create_repo(HF_DATASET, repo_type="dataset", exist_ok=True, token=HF_TOKEN)
        upload_file(
            path_or_fileobj=mp4_path,
            path_in_repo="Molecular-Gigafactory.mp4",
            repo_id=HF_DATASET,
            repo_type="dataset",
            token=HF_TOKEN,
        )
        # Also upload poster frame
        poster = os.path.join(OUTPUT_DIR, "poster.jpg")
        subprocess.run([
            "ffmpeg", "-y", "-i", mp4_path, "-ss", "30", "-vframes", "1",
            "-q:v", "2", poster,
        ], capture_output=True, timeout=60)
        if os.path.exists(poster):
            upload_file(
                path_or_fileobj=poster,
                path_in_repo="Molecular-Gigafactory-poster.jpg",
                repo_id=HF_DATASET,
                repo_type="dataset",
                token=HF_TOKEN,
            )
        return True
    except Exception as e:
        print(f"Upload failed: {e}")
        return False


def save_status(data):
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    with open(STATUS_FILE, "w") as f:
        json.dump(data, f)


def load_status():
    try:
        with open(STATUS_FILE) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {"phase": "idle", "progress": 0, "message": "Ready"}


# ---------------------------------------------------------------------------
# Background render thread
# ---------------------------------------------------------------------------
def render_test_poster():
    """Minimal smoke test: 64x64, 1 sample, CPU Cycles. Returns True on success."""
    print("SMOKE TEST: 64x64, 1 sample, CPU...")
    save_status({"phase": "smoke", "progress": 0, "message": "Smoke test running..."})
    smoke_cmd = [
        "blender", "--background",
        "--python", "/app/scene.py", "--",
        "--smoke",
        f"--output={OUTPUT_DIR}",
    ]
    proc = subprocess.run(smoke_cmd, capture_output=True, text=True, timeout=900)
    smoke_path = os.path.join(OUTPUT_DIR, "smoke_test.png")
    if os.path.exists(smoke_path):
        print(f"SMOKE TEST PASSED: {os.path.getsize(smoke_path)} bytes")
        save_status({"phase": "smoke_ok", "progress": 0,
                     "message": f"Smoke passed: {os.path.getsize(smoke_path)} bytes"})
        if HF_TOKEN:
            try:
                api = HfApi(token=HF_TOKEN)
                create_repo(HF_DATASET, repo_type="dataset", exist_ok=True, token=HF_TOKEN)
                upload_file(path_or_fileobj=smoke_path, path_in_repo="smoke_test.png",
                            repo_id=HF_DATASET, repo_type="dataset", token=HF_TOKEN)
                upload_file(path_or_fileobj=smoke_path, path_in_repo="Molecular-Gigafactory-poster.jpg",
                            repo_id=HF_DATASET, repo_type="dataset", token=HF_TOKEN)
            except Exception as e:
                print(f"Smoke upload failed: {e}")
        return True
    else:
        # STAY in failed state so we can read the error
        err_detail = proc.stdout[-800:] if proc.stdout else "(no stdout)"
        err_detail += "\n--- STDERR ---\n"
        err_detail += proc.stderr[-400:] if proc.stderr else "(no stderr)"
        print(f"SMOKE TEST FAILED:\n{err_detail}")
        save_status({"phase": "smoke_failed", "progress": 0,
                     "message": f"Smoke FAILED (exit {proc.returncode})\n{err_detail}"})
        return False


def live_preview_uploader(stop_event):
    """Watch for new preview frames and upload them to HF Dataset in background."""
    preview_path = os.path.join(OUTPUT_DIR, "preview", "latest_frame.png")
    uploaded_mtime = 0
    while not stop_event.is_set():
        try:
            if os.path.exists(preview_path):
                mtime = os.path.getmtime(preview_path)
                if mtime > uploaded_mtime:
                    # Small delay to let file write complete
                    time.sleep(1)
                    if HF_TOKEN:
                        api = HfApi(token=HF_TOKEN)
                        create_repo(HF_DATASET, repo_type="dataset", exist_ok=True, token=HF_TOKEN)
                        upload_file(
                            path_or_fileobj=preview_path,
                            path_in_repo="live_frame.png",
                            repo_id=HF_DATASET,
                            repo_type="dataset",
                            token=HF_TOKEN,
                        )
                    uploaded_mtime = mtime
        except Exception as e:
            print(f"Live preview upload error: {e}")
        stop_event.wait(10)


def render_all_phases():
    """Render all phases, encode, upload. Runs in background thread."""
    # Start live preview uploader
    stop_event = threading.Event()
    preview_thread = threading.Thread(target=live_preview_uploader, args=(stop_event,), daemon=True)
    preview_thread.start()

    # Smoke test must pass before we waste time on full render
    if not render_test_poster():
        stop_event.set()
        print("SMOKE TEST FAILED — aborting render")
        return
    save_status({"phase": "starting", "progress": 0, "message": "Initializing..."})

    completed_phases = []
    for phase in PHASES:
        # Check if phase already done (from checkpoint)
        progress = {}
        try:
            with open(PROGRESS_FILE) as f:
                progress = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            pass

        phase_key = f"phase_{phase}"
        if progress.get(phase_key, {}).get("complete"):
            msg = f"Phase {phase} already complete, skipping"
            print(msg)
            save_status({"phase": phase, "progress": int((len(completed_phases) + 1) / len(PHASES) * 100),
                         "message": msg})
            completed_phases.append(phase)
            continue

        msg = f"Rendering phase {phase}..."
        print(msg)
        save_status({"phase": phase, "progress": int(len(completed_phases) / len(PHASES) * 100),
                     "message": msg})

        ok = run_blender_phase(phase)
        if ok:
            completed_phases.append(phase)
            save_status({"phase": phase, "progress": int(len(completed_phases) / len(PHASES) * 100),
                         "message": f"Phase {phase} done"})
        else:
            # run_blender_phase already saved detailed error to status
            stop_event.set()
            return

    # All phases done — encode final
    save_status({"phase": "encoding", "progress": 85, "message": "Encoding final MP4..."})
    final_mp4 = encode_final()
    if not final_mp4:
        save_status({"phase": "error", "progress": 85, "message": "Final encode failed"})
        stop_event.set()
        return

    # Upload
    save_status({"phase": "uploading", "progress": 90, "message": "Uploading to HF Dataset..."})
    ok = upload_result(final_mp4)
    if ok:
        save_status({"phase": "done", "progress": 100,
                     "message": f"Complete! MP4 uploaded to {HF_DATASET}"})
    else:
        save_status({"phase": "done", "progress": 100,
                     "message": f"Render complete but upload failed. MP4 at {final_mp4}"})
    stop_event.set()


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
def get_status():
    s = load_status()
    return json.dumps(s, indent=2)


def get_diagnostics():
    """Return diagnostic info: dir listing, recent Blender log."""
    lines = []
    lines.append("=== STATUS ===")
    lines.append(json.dumps(load_status(), indent=2))
    lines.append("")
    lines.append("=== OUTPUT DIR ===")
    try:
        for root, dirs, files in os.walk(OUTPUT_DIR):
            rel = os.path.relpath(root, OUTPUT_DIR)
            if rel == ".":
                rel = OUTPUT_DIR
            lines.append(f"\n{rel}/")
            file_list = sorted(files)
            if len(file_list) > 20:
                lines.append(f"  ({len(file_list)} files, showing first/last 10)")
                for f in file_list[:10]:
                    fp = os.path.join(root, f)
                    sz = os.path.getsize(fp)
                    lines.append(f"  {f}  ({sz} bytes)")
                lines.append(f"  ...")
                for f in file_list[-10:]:
                    fp = os.path.join(root, f)
                    sz = os.path.getsize(fp)
                    lines.append(f"  {f}  ({sz} bytes)")
            else:
                for f in file_list:
                    fp = os.path.join(root, f)
                    sz = os.path.getsize(fp)
                    lines.append(f"  {f}  ({sz} bytes)")
    except Exception as e:
        lines.append(f"Error: {e}")
    lines.append("")
    lines.append("=== BLENDER LOG (last 3KB) ===")
    for phase in PHASES:
        log_path = os.path.join(OUTPUT_DIR, f"blender_{phase}.log")
        if os.path.exists(log_path):
            with open(log_path) as f:
                content = f.read()
                lines.append(f"\n--- blender_{phase}.log ({len(content)} bytes) ---")
                lines.append(content[-3000:] if len(content) > 3000 else content)
    return "\n".join(lines)


def start_render():
    s = load_status()
    if s.get("phase") in ("running", "starting", "encoding", "uploading"):
        return "Render already in progress!"
    threading.Thread(target=render_all_phases, daemon=True).start()
    return "Render started! Check status below."


def get_download_file():
    """Return the final MP4 path for download, or latest phase MP4 if not ready."""
    final_mp4 = os.path.join(OUTPUT_DIR, "Molecular-Gigafactory.mp4")
    if os.path.exists(final_mp4):
        return final_mp4
    # Fall back to latest phase MP4
    for phase in reversed(PHASES):
        phase_mp4 = os.path.join(OUTPUT_DIR, f"phase_{phase}.mp4")
        if os.path.exists(phase_mp4):
            return phase_mp4
    return None


with gr.Blocks(title="Molecular Gigafactory Render") as demo:
    gr.Markdown("""
    # Molecular Gigafactory — GPU Render Service
    Auto-starts rendering on Space boot. Results uploaded to HF Dataset.
    """)

    with gr.Row():
        start_btn = gr.Button("Start Render", variant="primary")
        refresh_btn = gr.Button("Refresh Status")
        diag_btn = gr.Button("Diagnostics")

    status_display = gr.Textbox(
        value=get_status(), label="Status", lines=6
    )

    diag_display = gr.Textbox(
        value="", label="Diagnostics", lines=20, visible=False
    )

    with gr.Row():
        download_btn = gr.Button("Download MP4")
        download_output = gr.File(label="Download")

    start_btn.click(start_render, outputs=[status_display])
    refresh_btn.click(get_status, outputs=[status_display])
    diag_btn.click(get_diagnostics, outputs=[diag_display])
    download_btn.click(get_download_file, outputs=[download_output])

    # Manual refresh via button above
    demo.load(get_status, outputs=[status_display])


# ---------------------------------------------------------------------------
# Auto-start on boot
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    status = load_status()
    if status.get("phase") in ("idle", "error", None) or status.get("progress", 0) < 100:
        print(f"Auto-starting render (current status: {status.get('phase')})")
        threading.Thread(target=render_all_phases, daemon=True).start()
    else:
        print(f"Render already complete (status: {status.get('phase')})")

    demo.queue(default_concurrency_limit=1)
    demo.launch(server_name="0.0.0.0", server_port=7860, theme="soft")