"""Gradio demo: compress PBR texture sets to Intel TSNC format."""

from __future__ import annotations

import tempfile
from pathlib import Path

import install_dependencies

install_dependencies.install_private_package()

try:
    import spaces

    ON_SPACES = True
except ImportError:
    ON_SPACES = False

    def _gpu_decorator(**kwargs):
        def wrapper(fn):
            return fn

        return wrapper

    class _SpacesShim:
        GPU = staticmethod(_gpu_decorator)

    spaces = _SpacesShim()

import gradio as gr
import numpy as np
import torch
from numpy import iinfo

from tsnc.hypernetwork import Hypernetwork, functional_bcf1, write_tsnc

HF_MODEL_ID = "belcour/egsr2026_hypernetwork"
CFG_FILE = "karajan_ZDitTv2_2026_06_10_lod.cfg"

# ZeroGPU: CUDA is not visible in the main process until after .to("cuda").
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading hypernetwork on {device}...")
hypernetwork = Hypernetwork.from_pretrained(HF_MODEL_ID, CFG_FILE)
hypernetwork.eval().to(device)
print("Hypernetwork ready.")


def _validate_inputs(diffuse, normal, arm) -> tuple[int, int]:
    if diffuse is None or normal is None or arm is None:
        raise gr.Error("Please upload diffuse, normal, and ARM textures.")

    images = [("diffuse", diffuse), ("normal", normal), ("arm", arm)]
    shapes = {name: img.shape[:2] for name, img in images}
    ref_shape = shapes["diffuse"]

    for name, shape in shapes.items():
        if shape != ref_shape:
            raise gr.Error(
                f"All textures must have the same resolution. "
                f"Diffuse is {ref_shape[0]}x{ref_shape[1]}, {name} is {shape[0]}x{shape[1]}."
            )

    height, width = ref_shape
    if height != width:
        raise gr.Error(f"Textures must be square (got {height}x{width}).")

    return height, width


def _numpy_to_tensor(img: np.ndarray) -> torch.Tensor:
    if iinfo(img.dtype).max > 255:
        img = img.astype(np.float32) / 65535.0
    else:
        img = img.astype(np.float32) / 255.0

    if img.ndim == 2:
        img = img[..., None]
    else:
        img = img[..., :3]

    return torch.tensor(img, dtype=torch.float32)


def _stack_textures(diffuse: np.ndarray, normal: np.ndarray, arm: np.ndarray) -> torch.Tensor:
    tensors = [_numpy_to_tensor(img) for img in (diffuse, normal, arm)]
    ts_tensor = torch.cat(tensors, dim=-1).to(device)
    return ts_tensor[None].permute(0, 3, 1, 2)


def _write_tsnc_to_temp(bcf1s, wb) -> str:
    outdir = tempfile.mkdtemp(prefix="tsnc_")
    out_name = "compressed.tsnc"
    write_tsnc(bcf1s, wb, cfg=hypernetwork.cfg, outdir=outdir, out_name=out_name)
    return str(Path(outdir) / out_name)


def _gpu_duration(diffuse, normal, arm) -> int:
    """Scale GPU quota with texture resolution (square maps)."""
    for img in (diffuse, normal, arm):
        if img is not None:
            return max(120, int(img.shape[0] * img.shape[1] / 8192))
    return 120


@spaces.GPU(duration=_gpu_duration)
def compress(diffuse, normal, arm):
    _height, width = _validate_inputs(diffuse, normal, arm)

    with torch.inference_mode():
        ts_tensor = _stack_textures(diffuse, normal, arm)
        bcf1s, wb = hypernetwork(ts_tensor)
        ts_out, _ = functional_bcf1(
            bcf1s,
            wb,
            W=width,
            lod=0.0,
            subpixel_shift=False,
            activation=hypernetwork.activation,
        )

    tsnc_path = _write_tsnc_to_temp(bcf1s, wb)
    preview_diff = (ts_out[0, ..., 0:3].detach().cpu().clamp(0.0, 1.0).numpy() * 255.0).astype(
        np.uint8
    )
    preview_norm = (ts_out[0, ..., 3:6].detach().cpu().clamp(0.0, 1.0).numpy() * 255.0).astype(
        np.uint8
    )
    preview_arm = (ts_out[0, ..., 6:9].detach().cpu().clamp(0.0, 1.0).numpy() * 255.0).astype(
        np.uint8
    )

    return tsnc_path, preview_diff, preview_norm, preview_arm


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="TSNC Hypernetwork") as demo:
        gr.Markdown(
            """
            # TSNC Hypernetwork

            Compress a PBR texture set (diffuse, normal, ARM) into Intel's
            [Texture Set Neural Compression (TSNC)](https://github.com/GameTechDev/TextureSetNeuralCompressionSample) format.

            **Requirements:** square textures, same resolution for all three maps.
            """
        )

        with gr.Row():
            diffuse_in = gr.Image(label="Diffuse", type="numpy", image_mode="RGB")
            normal_in = gr.Image(label="Normal", type="numpy", image_mode="RGB")
            arm_in = gr.Image(label="ARM", type="numpy", image_mode="RGB")

        compress_btn = gr.Button("Compress", variant="primary")

        tsnc_out = gr.File(label="TSNC file")

        with gr.Row():
            diff_out = gr.Image(label="Reconstructed diffuse", type="numpy")
            norm_out = gr.Image(label="Reconstructed normal", type="numpy")
            arm_out = gr.Image(label="Reconstructed ARM", type="numpy")

        compress_btn.click(
            fn=compress,
            inputs=[diffuse_in, normal_in, arm_in],
            outputs=[tsnc_out, diff_out, norm_out, arm_out],
        )

    return demo


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