import spaces
import gradio as gr
import os
import sys
import time
import tempfile
import shutil
import torch

_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _root)
sys.path.insert(0, os.path.join(_root, "common"))

from PIL import Image

# Weights live in a separate HF model repo (the free Space repo caps at 1GB);
# from_pretrained fetches them over HF's CDN at startup, like the original demo.
WEIGHTS_REPO = "ifire/seethrough-weights"
LAYERDIFF_SUB = "layerdiff3d"
MARIGOLD_SUB = "marigold"


def _log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# --------------- Preload models to CPU at startup ---------------
_log("Preloading LayerDiff pipeline to CPU...")
from modules.layerdiffuse.diffusers_kdiffusion_sdxl import KDiffusionStableDiffusionXLPipeline
from modules.layerdiffuse.layerdiff3d import UNetFrameConditionModel
from modules.layerdiffuse.vae import TransparentVAE, TransparentVAEDecoder, TransparentVAEEncoder

_trans_vae = TransparentVAE.from_pretrained(WEIGHTS_REPO, subfolder=f"{LAYERDIFF_SUB}/trans_vae")
_unet_ld = UNetFrameConditionModel.from_pretrained(WEIGHTS_REPO, subfolder=f"{LAYERDIFF_SUB}/unet")
_layerdiff_pipe = KDiffusionStableDiffusionXLPipeline.from_pretrained(
    WEIGHTS_REPO, subfolder=LAYERDIFF_SUB, trans_vae=_trans_vae, unet=_unet_ld, scheduler=None
)
_log("LayerDiff pipeline loaded to CPU.")

_log("Preloading Marigold pipeline to CPU...")
from modules.marigold import MarigoldDepthPipeline

_unet_mg = UNetFrameConditionModel.from_pretrained(WEIGHTS_REPO, subfolder=f"{MARIGOLD_SUB}/unet")
_marigold_pipe = MarigoldDepthPipeline.from_pretrained(WEIGHTS_REPO, subfolder=MARIGOLD_SUB, unet=_unet_mg)
_log("Marigold pipeline loaded to CPU.")

_models_on_gpu = False

from utils.inference_utils import apply_layerdiff, apply_marigold, further_extr
from utils.torch_utils import seed_everything
import utils.inference_utils as _inf


def _move_to_gpu():
    global _models_on_gpu
    if _models_on_gpu:
        _log("Models already on GPU, skipping transfer.")
        return

    t0 = time.time()
    _log("Moving LayerDiff to CUDA bf16...")
    _layerdiff_pipe.vae.to(dtype=torch.bfloat16, device="cuda")
    _layerdiff_pipe.trans_vae.to(dtype=torch.bfloat16, device="cuda")
    _layerdiff_pipe.unet.to(dtype=torch.bfloat16, device="cuda")
    _layerdiff_pipe.text_encoder.to(dtype=torch.bfloat16, device="cuda")
    _layerdiff_pipe.text_encoder_2.to(dtype=torch.bfloat16, device="cuda")
    _log(f"LayerDiff on GPU ({time.time() - t0:.1f}s)")

    t0 = time.time()
    _log("Moving Marigold to CUDA bf16...")
    _marigold_pipe.to(device="cuda", dtype=torch.bfloat16)
    _log(f"Marigold on GPU ({time.time() - t0:.1f}s)")

    # Inject into inference_utils globals so apply_* functions skip their own loading
    _inf.layerdiff_pipeline = _layerdiff_pipe
    _inf.marigold_pipeline = _marigold_pipe

    _models_on_gpu = True


_SKIP_TAGS = {"src_img", "src_head", "reconstruction"}

def _collect_layer_gallery(saved_dir):
    """Collect layer PNGs as (image, label) tuples for the gallery."""
    gallery = []
    for f in sorted(os.listdir(saved_dir)):
        if not f.endswith(".png"):
            continue
        tag = f[:-4]
        if tag.endswith("_depth") or tag in _SKIP_TAGS:
            continue
        img = Image.open(os.path.join(saved_dir, f))
        gallery.append((img, tag))
    return gallery


@spaces.GPU(duration=150)
def inference(image: Image.Image, resolution: int = 720, seed: int = 42, split_left_right: bool = False):
    t_start = time.time()
    if image is None:
        raise gr.Error("Please upload an image.")

    # Snap to nearest multiple of 16 for clean latent dimensions
    resolution = max(64, min(resolution, 1280))
    resolution = round(resolution / 16) * 16
    _log(f"Resolution: {resolution}, Seed: {seed}, split_left_right: {split_left_right}, Image: {image.size}")

    _move_to_gpu()
    seed_everything(seed)

    tmpdir = tempfile.mkdtemp(prefix="seethrough_")
    try:
        input_path = os.path.join(tmpdir, "input.png")
        image.save(input_path)

        t0 = time.time()
        _log("Running LayerDiff...")
        apply_layerdiff(
            input_path, WEIGHTS_REPO,
            save_dir=tmpdir, seed=seed, resolution=resolution,
        )
        _log(f"LayerDiff done ({time.time() - t0:.1f}s)")

        t0 = time.time()
        _log("Running Marigold depth...")
        apply_marigold(
            input_path, WEIGHTS_REPO,
            save_dir=tmpdir, seed=seed, resolution=resolution,
        )
        _log(f"Marigold done ({time.time() - t0:.1f}s)")

        saved = os.path.join(tmpdir, "input")

        # Collect gallery before PSD assembly (further_extr may modify files)
        gallery = _collect_layer_gallery(saved)

        t0 = time.time()
        _log("Running PSD assembly...")
        further_extr(saved, rotate=False, save_to_psd=True, tblr_split=split_left_right)
        _log(f"PSD assembly done ({time.time() - t0:.1f}s)")

        psd_path = saved + ".psd"
        if os.path.exists(psd_path):
            output_path = os.path.join(
                tempfile.gettempdir(), "seethrough_output.psd"
            )
            shutil.copy2(psd_path, output_path)
            _log(f"Total inference time: {time.time() - t_start:.1f}s")
            return output_path, gallery

        raise gr.Error("PSD generation failed — no output file produced.")
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


with gr.Blocks(title="See-through: Layer Decomposition") as demo:
    gr.Markdown(
        "# See-through: Single-image Layer Decomposition for Anime Characters\n\n"
        '<a href="https://github.com/shitagaki-lab/see-through" target="_blank">GitHub</a> | '
        '<a href="https://arxiv.org/abs/2602.03749" target="_blank">Paper (arXiv:2602.03749)</a>\n\n'
        "Upload an anime character illustration to decompose it into "
        "fully-inpainted semantic layers with depth ordering, "
        "exported as a layered PSD file.\n\n"
        "**Note:** 720 resolution is the default; higher resolutions take longer. "
        "Weights stream from the ifire/seethrough-weights model repo (git-lfs) "
        "on first boot — no runtime download per request."
    )
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Upload image (non-square images will be padded)")
            resolution = gr.Slider(
                minimum=512, maximum=1280, value=720, step=16,
                label="Resolution",
                info="720 default; higher takes longer.",
            )
            seed = gr.Slider(minimum=0, maximum=9999, value=42, step=1, label="Seed")
            split_left_right = gr.Checkbox(
                value=False,
                label="Split left/right arms & legs",
                info="Separate left and right limbs into individual layers.",
            )
            run_btn = gr.Button("Run", variant="primary")
        with gr.Column(scale=2):
            psd_output = gr.File(label="Download layered PSD")
            gallery_output = gr.Gallery(label="Separated layers", columns=4, height="auto")

    run_btn.click(
        fn=inference,
        inputs=[input_image, resolution, seed, split_left_right],
        outputs=[psd_output, gallery_output],
    )

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