import hashlib
import os
import shutil
import subprocess
import sys
import traceback
from pathlib import Path
from typing import List, Tuple

import gradio as gr
import spaces
from huggingface_hub import snapshot_download

APP_BUILD = "relit-live-readme-lite-2026-06-29"

ROOT = Path(__file__).resolve().parent
DATASETS_DEMOS = ROOT / "datasets" / "demos"
DATASETS_ENVS = ROOT / "datasets" / "envs"
RUNTIME_DIR = ROOT / "runtime_zerogpu_lite"
CACHE_DIR = RUNTIME_DIR / "cache"
TMP_DATASETS_DIR = RUNTIME_DIR / "datasets"
OUTPUT_DIR = RUNTIME_DIR / "outputs"
PREVIEW_DIR = RUNTIME_DIR / "previews"

RELIT_REPO_ID = os.environ.get("RELIT_HF_REPO", "weiqingXiao/Relit-LiVE")
WAN_REPO_ID = os.environ.get("RELIT_WAN_REPO", "Wan-AI/Wan2.1-T2V-1.3B")
DOWNLOAD_WEIGHTS = os.environ.get("RELIT_DOWNLOAD_WEIGHTS", "1") != "0"
ZEROGPU_SIZE = os.environ.get("RELIT_ZEROGPU_SIZE", "large")

CKPT_25 = ROOT / "checkpoints" / "model_frame25_480_832.ckpt"
WAN_DIR = ROOT / "models" / "Wan-AI" / "Wan2.1-T2V-1.3B"
WAN_REQUIRED = [
    WAN_DIR / "diffusion_pytorch_model.safetensors",
    WAN_DIR / "models_t5_umt5-xxl-enc-bf16.pth",
    WAN_DIR / "Wan2.1_VAE.pth",
]

VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi"}
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}


def _tail(text: str, max_chars: int = 5000) -> str:
    if not text:
        return ""
    text = text.strip()
    return text[-max_chars:]


def run_command(cmd: List[str], timeout: int | None = None) -> subprocess.CompletedProcess:
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["PYTHONPATH"] = f"{ROOT}:{env.get('PYTHONPATH', '')}"
    env.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
    env.setdefault("TOKENIZERS_PARALLELISM", "false")
    return subprocess.run(
        cmd,
        cwd=str(ROOT),
        env=env,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=timeout,
    )


def ensure_dirs() -> None:
    for path in [RUNTIME_DIR, CACHE_DIR, TMP_DATASETS_DIR, OUTPUT_DIR, PREVIEW_DIR, ROOT / "checkpoints", WAN_DIR]:
        path.mkdir(parents=True, exist_ok=True)


def ensure_weights() -> str:
    ensure_dirs()
    logs: List[str] = [f"App build: {APP_BUILD}"]

    if not DOWNLOAD_WEIGHTS:
        logs.append("Weight download disabled with RELIT_DOWNLOAD_WEIGHTS=0.")
        return "\n".join(logs)

    if not CKPT_25.exists():
        logs.append(f"Downloading Relit-LiVE 25-frame checkpoint from {RELIT_REPO_ID}...")
        snapshot_download(
            repo_id=RELIT_REPO_ID,
            local_dir=str(ROOT),
            allow_patterns=["checkpoints/model_frame25_480_832.ckpt"],
            token=os.environ.get("HF_TOKEN"),
        )
        logs.append("Relit-LiVE 25-frame checkpoint ready.")
    else:
        logs.append("Relit-LiVE 25-frame checkpoint already present.")

    if not all(path.exists() for path in WAN_REQUIRED):
        logs.append(f"Downloading Wan2.1 base model from {WAN_REPO_ID}...")
        snapshot_download(
            repo_id=WAN_REPO_ID,
            local_dir=str(WAN_DIR),
            token=os.environ.get("HF_TOKEN"),
        )
        logs.append("Wan2.1 base model ready.")
    else:
        logs.append("Wan2.1 base model already present.")

    return "\n".join(logs)


def list_samples() -> List[str]:
    if not DATASETS_DEMOS.exists():
        return []
    samples = []
    for item in sorted(DATASETS_DEMOS.iterdir()):
        if not item.is_dir():
            continue
        # relit_inference.py expects at least images_4 plus intrinsic buffers.
        if (item / "images_4").exists():
            samples.append(item.name)
    return samples


def list_envs() -> List[str]:
    if not DATASETS_ENVS.exists():
        return []
    envs = []
    for item in sorted(DATASETS_ENVS.iterdir()):
        if not item.is_dir():
            continue
        if (item / "ldr_video_fix_first_frame.mp4").exists() or (item / "hdr_log_video_fix_first_frame.mp4").exists():
            envs.append(item.name)
    return envs


def sample_frame_paths(sample_name: str, limit: int = 8) -> List[str]:
    images_dir = DATASETS_DEMOS / sample_name / "images_4"
    if not images_dir.exists():
        return []
    files = []
    for ext in sorted(IMAGE_EXTS):
        files.extend(images_dir.glob(f"*{ext}"))
    files = sorted(files)
    if len(files) <= limit:
        return [str(f) for f in files]
    # Spread frames across the sequence.
    idxs = sorted(set(round(i * (len(files) - 1) / (limit - 1)) for i in range(limit)))
    return [str(files[i]) for i in idxs]


def sample_info(sample_name: str) -> str:
    sample_dir = DATASETS_DEMOS / sample_name
    if not sample_dir.exists():
        return "Sample not found."
    parts = []
    for sub in ["images_4", "Base Color", "normal", "Roughness", "depth", "Metallic"]:
        d = sample_dir / sub
        if d.exists():
            count = sum(1 for p in d.iterdir() if p.suffix.lower() in IMAGE_EXTS | {".exr"})
            parts.append(f"{sub}: {count} files")
    return "\n".join(parts) or "No recognized Relit-LiVE sample folders found."


def update_preview(sample_name: str) -> Tuple[List[str], str]:
    if not sample_name:
        return [], "No sample selected."
    return sample_frame_paths(sample_name), sample_info(sample_name)


def validate_runtime(sample_name: str, env_name: str) -> None:
    missing = []
    if not (ROOT / "relit_inference.py").exists():
        missing.append("relit_inference.py")
    if not CKPT_25.exists():
        missing.append(str(CKPT_25.relative_to(ROOT)))
    for path in WAN_REQUIRED:
        if not path.exists():
            missing.append(str(path.relative_to(ROOT)))
    sample_dir = DATASETS_DEMOS / sample_name
    env_dir = DATASETS_ENVS / env_name
    if not sample_dir.exists():
        missing.append(f"datasets/demos/{sample_name}")
    if not (sample_dir / "images_4").exists():
        missing.append(f"datasets/demos/{sample_name}/images_4")
    if not env_dir.exists():
        missing.append(f"datasets/envs/{env_name}")
    if missing:
        raise RuntimeError("Missing required files:\n- " + "\n- ".join(missing))


def make_single_sample_dataset(sample_name: str, job_dir: Path) -> Path:
    dataset_dir = job_dir / "dataset"
    dataset_dir.mkdir(parents=True, exist_ok=True)
    src = DATASETS_DEMOS / sample_name
    dst = dataset_dir / sample_name
    if dst.exists() or dst.is_symlink():
        if dst.is_symlink() or dst.is_file():
            dst.unlink()
        else:
            shutil.rmtree(dst)
    try:
        dst.symlink_to(src.resolve(), target_is_directory=True)
    except Exception:
        shutil.copytree(src, dst)
    return dataset_dir


def cache_key(sample_name: str, env_name: str, mode: str, num_frames: int, steps: int, cfg_scale: float, quality: int) -> str:
    raw = f"{APP_BUILD}|{sample_name}|{env_name}|{mode}|{num_frames}|{steps}|{cfg_scale}|{quality}"
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


def result_path_for(key: str, num_frames: int) -> Path:
    suffix = ".png" if int(num_frames) == 1 else ".mp4"
    return OUTPUT_DIR / key / f"relit_live_result{suffix}"


def estimate_duration(sample_name: str, env_name: str, mode: str, num_frames: int, steps: int, cfg_scale: float, quality: int) -> int:
    # Includes model load time because we intentionally follow the official README CLI path.
    frames = int(num_frames)
    steps = int(steps)
    return min(1800, max(180, int(180 + frames * steps * 1.2)))


@spaces.GPU(duration=estimate_duration, size=ZEROGPU_SIZE)
def run_relit_cli(sample_name: str, env_name: str, mode: str, num_frames: int, steps: int, cfg_scale: float, quality: int) -> Tuple[str, str]:
    validate_runtime(sample_name, env_name)
    key = cache_key(sample_name, env_name, mode, int(num_frames), int(steps), float(cfg_scale), int(quality))
    output_path = result_path_for(key, int(num_frames))
    if output_path.exists() and output_path.stat().st_size > 0:
        return str(output_path), f"Cache hit: {key}"

    job_dir = TMP_DATASETS_DIR / key
    if job_dir.exists():
        shutil.rmtree(job_dir)
    job_dir.mkdir(parents=True, exist_ok=True)
    dataset_dir = make_single_sample_dataset(sample_name, job_dir)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    cmd = [
        sys.executable,
        str(ROOT / "relit_inference.py"),
        "--dataset_path", str(dataset_dir),
        "--ckpt_path", str(CKPT_25),
        "--output_dir", str(output_path.parent),
        "--output_path", str(output_path),
        "--cfg_scale", str(float(cfg_scale)),
        "--height", "480",
        "--width", "832",
        "--num_frames", str(int(num_frames)),
        "--padding_resolution",
        "--use_ref_image",
        "--env_map_path", str(DATASETS_ENVS / env_name),
        "--frame_interval", "1",
        "--num_inference_steps", str(int(steps)),
        "--quality", str(int(quality)),
        "--dataloader_num_workers", "0",
    ]

    if mode == "Rotating light":
        cmd.append("--use_rotate_light")
    elif mode == "Fixed first frame + width light rotation":
        cmd.append("--use_fixed_frame_and_w_rotate_light")
    elif mode == "Fixed first frame + height light rotation":
        cmd.append("--use_fixed_frame_and_h_rotate_light")

    proc = run_command(cmd, timeout=3600)
    stdout_tail = _tail(proc.stdout)
    stderr_tail = _tail(proc.stderr)
    log = (
        "Command:\n" + " ".join(cmd) +
        "\n\nSTDOUT tail:\n" + stdout_tail +
        "\n\nSTDERR tail:\n" + stderr_tail
    )

    if proc.returncode != 0:
        raise RuntimeError(f"relit_inference.py failed with exit code {proc.returncode}.\n\n{log}")
    if not output_path.exists() or output_path.stat().st_size == 0:
        produced = sorted(output_path.parent.glob("*.mp4")) + sorted(output_path.parent.glob("*.png"))
        if produced:
            output_path = produced[-1]
        else:
            raise RuntimeError("relit_inference.py finished but no output file was produced.\n\n" + log)

    return str(output_path), log


def run_demo(sample_name: str, env_name: str, mode: str, num_frames: int, steps: int, cfg_scale: float, quality: int):
    try:
        ensure_weights()
        key = cache_key(sample_name, env_name, mode, int(num_frames), int(steps), float(cfg_scale), int(quality))
        cached = result_path_for(key, int(num_frames))
        if cached.exists() and cached.stat().st_size > 0:
            status = f"✅ Cached result loaded.\nKey: {key}"
            if int(num_frames) == 1:
                return gr.update(value=str(cached), visible=True), gr.update(value=None, visible=False), status
            return gr.update(value=None, visible=False), gr.update(value=str(cached), visible=True), status

        output, log = run_relit_cli(sample_name, env_name, mode, int(num_frames), int(steps), float(cfg_scale), int(quality))
        status = "✅ Relit-LiVE inference completed.\n\n" + _tail(log, 2500)
        if int(num_frames) == 1 or output.lower().endswith(".png"):
            return gr.update(value=output, visible=True), gr.update(value=None, visible=False), status
        return gr.update(value=None, visible=False), gr.update(value=output, visible=True), status
    except Exception as exc:
        status = "❌ Inference failed.\n\n" + str(exc) + "\n\n" + traceback.format_exc()[-3000:]
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), status


def startup_message() -> str:
    try:
        return ensure_weights()
    except Exception as exc:
        return "Runtime setup failed:\n" + str(exc)


SAMPLES = list_samples()
ENVS = list_envs()
DEFAULT_SAMPLE = SAMPLES[0] if SAMPLES else None
DEFAULT_ENV = "Pink_Sunrise" if "Pink_Sunrise" in ENVS else (ENVS[0] if ENVS else None)

css = """
#notice {border: 1px solid #ddd; border-radius: 12px; padding: 12px;}
"""

with gr.Blocks(title="Relit-LiVE ZeroGPU README Demo", css=css) as demo:
    gr.Markdown(
        "# Relit-LiVE ZeroGPU Demo\n"
        "This Space follows the repository README inference path: `python relit_inference.py`. "
        "It uses curated samples already present in `datasets/demos` and environments from `datasets/envs`. "
        "The full Gradio inverse-rendering pipeline is intentionally not used here."
    )

    with gr.Accordion("Runtime setup", open=False):
        setup_box = gr.Textbox(value=startup_message(), label="Setup log", lines=10)
        refresh_setup = gr.Button("Refresh setup check")
        refresh_setup.click(fn=startup_message, inputs=[], outputs=[setup_box], queue=False)

    if not SAMPLES or not ENVS:
        gr.Markdown(
            "## Missing demo assets\n"
            "This app expects the cloned repository to contain `datasets/demos/` and `datasets/envs/`."
        )
    else:
        with gr.Row():
            with gr.Column(scale=1):
                sample = gr.Dropdown(SAMPLES, value=DEFAULT_SAMPLE, label="Demo sample")
                env = gr.Dropdown(ENVS, value=DEFAULT_ENV, label="Lighting environment")
                mode = gr.Radio(
                    [
                        "Basic relighting",
                        "Rotating light",
                        "Fixed first frame + width light rotation",
                        "Fixed first frame + height light rotation",
                    ],
                    value="Basic relighting",
                    label="README inference mode",
                )
                num_frames = gr.Radio([1, 9, 17, 25], value=9, label="Frames")
                steps = gr.Slider(8, 50, value=16, step=1, label="Inference steps")
                cfg_scale = gr.Slider(1.0, 5.0, value=1.0, step=0.1, label="CFG scale")
                quality = gr.Slider(5, 10, value=8, step=1, label="Video quality")
                run_btn = gr.Button("Run README inference", variant="primary")

            with gr.Column(scale=1):
                preview = gr.Gallery(label="Input sample frames", columns=4, height=320, value=sample_frame_paths(DEFAULT_SAMPLE) if DEFAULT_SAMPLE else [])
                info = gr.Textbox(label="Sample contents", lines=8, value=sample_info(DEFAULT_SAMPLE) if DEFAULT_SAMPLE else "")

        with gr.Row():
            result_image = gr.Image(label="Relit result", visible=False)
            result_video = gr.Video(label="Relit result", visible=False)

        status = gr.Textbox(label="Status / logs", lines=16)

        sample.change(fn=update_preview, inputs=[sample], outputs=[preview, info], queue=False)
        run_btn.click(
            fn=run_demo,
            inputs=[sample, env, mode, num_frames, steps, cfg_scale, quality],
            outputs=[result_image, result_video, status],
        )

    gr.Markdown(
        "---\n"
        "**ZeroGPU scope:** this demo does not accept arbitrary uploads because arbitrary inputs require the full inverse-rendering pipeline and Cosmos. "
        "Here we run the README `relit_inference.py` path on prepared Relit-LiVE samples."
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
        show_error=True,
    )