#!/usr/bin/env python3
"""Lotus 1 Space — depth + normal endpoints."""

# ── Compatibility shims for huggingface_hub >= 0.30 ──────────────────────────
# HfFolder and cached_download were removed in hub 0.30, but diffusers==0.28.0
# and gradio 5.x oauth module still import them. Restore stubs/aliases BEFORE
# any of those packages are imported.
import huggingface_hub as _hfhub

try:
    _hfhub.HfFolder
except AttributeError:
    class _HfFolderStub:
        @staticmethod
        def get_token(): return None
        @staticmethod
        def save_token(token): pass
        @staticmethod
        def delete_token(): pass
    _hfhub.HfFolder = _HfFolderStub

try:
    _hfhub.cached_download
except AttributeError:
    from huggingface_hub import hf_hub_download
    def _cached_download_compat(url_or_filename=None, *args, **kwargs):
        # diffusers uses cached_download(url, ...) or cached_download(repo_id, ...)
        # Best-effort redirect to hf_hub_download for repo-based calls
        if url_or_filename and not url_or_filename.startswith("http"):
            return hf_hub_download(url_or_filename, *args, **kwargs)
        # For URL-based calls, just return the path as-is (shouldn't reach here at inference)
        raise NotImplementedError(f"cached_download URL mode not supported: {url_or_filename}")
    _hfhub.cached_download = _cached_download_compat
    # Also patch the diffusers import path
    import sys
    if "huggingface_hub" in sys.modules:
        sys.modules["huggingface_hub"].cached_download = _cached_download_compat
# ─────────────────────────────────────────────────────────────────────────────

import os

# ── Patch gradio_client bool-schema bug ──────────────────────────────────────
# gradio_client <= ~1.7.x has a bug where `"const" in schema` raises
# TypeError when schema is a bool (e.g. additionalProperties: true).
# This crashes GET /gradio_api/info, making the space unreachable via API.
# Monkey-patch _json_schema_to_python_type to guard against non-dict schemas.
try:
    import gradio_client.utils as _gcu
    _orig_j2p = _gcu._json_schema_to_python_type
    def _safe_j2p(schema, defs=None):
        if not isinstance(schema, dict):
            return "Any"
        return _orig_j2p(schema, defs)
    _gcu._json_schema_to_python_type = _safe_j2p
except Exception as _e:
    print(f"[warn] gradio_client patch failed: {_e}")
# ─────────────────────────────────────────────────────────────────────────────

import tempfile

import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image
from infer import load_pipe, infer_pipe

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
SEED = 3

_pipes = {}

def _get_pipe(task: str):
    if task not in _pipes:
        pipe_g, _pipe_d = load_pipe(task, device)
        _pipes[task] = pipe_g
    return _pipes[task]


def _save_16bit_png(img: Image.Image) -> str:
    """Save the depth map as a true 16-bit PNG and return the path.

    colorize_depth_map (modified previously) already returns a uint16 PIL
    Image. The default gr.Image output would re-encode it as 8-bit webp,
    which destroys the upper byte and yields visible stepping in the
    fabricated relief (~20 µm per step over a 5 mm relief). Returning
    a file path via gr.File bypasses Gradio's image processing entirely.
    """
    arr = np.array(img)
    fd, tmp_path = tempfile.mkstemp(suffix='_depth16.png')
    os.close(fd)
    if arr.dtype == np.uint16:
        # Native 16-bit path — preserve all 65,536 levels
        Image.fromarray(arr, mode="I;16").save(tmp_path, format="PNG")
    elif arr.dtype == np.uint8:
        # Up-promote to 16-bit (no extra precision, but consistent output)
        arr16 = arr.astype(np.uint16) * 257  # 0->0, 255->65535
        Image.fromarray(arr16, mode="I;16").save(tmp_path, format="PNG")
    else:
        # Float / other — normalise to 16-bit
        a = arr.astype(np.float32)
        rng = max(a.max() - a.min(), 1e-9)
        a = (a - a.min()) / rng
        Image.fromarray((a * 65535).astype(np.uint16),
                         mode="I;16").save(tmp_path, format="PNG")
    return tmp_path


@spaces.GPU
def infer_depth(image_path):
    pipe = _get_pipe("depth")
    img = infer_pipe(pipe, image_path, "depth", SEED, device)
    return _save_16bit_png(img)


def _save_png(img: Image.Image) -> str:
    """Save an RGB / L image as lossless PNG and return the path.

    Used for the normal map: Gradio's gr.Image output would re-encode via
    webp (lossy), and even tiny per-pixel noise on flat regions becomes
    visible orange-peel artifacts after Poisson integration of the
    normal-derived gradient field.
    """
    fd, tmp_path = tempfile.mkstemp(suffix='_normal.png')
    os.close(fd)
    img.save(tmp_path, format="PNG")
    return tmp_path


@spaces.GPU
def infer_normal(image_path):
    pipe = _get_pipe("normal")
    img = infer_pipe(pipe, image_path, "normal", SEED, device)
    return _save_png(img)


def _save_normal_npy(arr) -> str:
    """Save the raw float normal prediction as a .npy (float16, exact).

    The model computes in float16, so float16 storage is lossless and only
    ~6 MB for 1024x1024x3. Values are in [0,1] — the same convention as the
    8-bit PNG, where the surface normal = value * 2 - 1. Returning the float
    array bypasses the 8-bit quantisation that stair-steps the integrated
    bas-relief into visible 'orange-peel' grain.
    """
    fd, tmp_path = tempfile.mkstemp(suffix='_normal_f16.npy')
    os.close(fd)
    np.save(tmp_path, np.asarray(arr, dtype=np.float16))
    return tmp_path


@spaces.GPU
def infer_normal16(image_path):
    pipe = _get_pipe("normal")
    _img, npy = infer_pipe(pipe, image_path, "normal", SEED, device,
                           return_float=True)
    return _save_normal_npy(npy)


with gr.Blocks(title="Lotus 1 - Depth + Normal") as demo:
    gr.Markdown("# Lotus 1 - Depth & Normal")
    gr.Markdown("API: `/depth` returns a **16-bit grayscale PNG** depth map, "
                "`/normal` returns a **lossless RGB PNG** normal map, and "
                "`/normal16` returns the **raw float16 normal** as a `.npy` "
                "(values in [0,1]; normal = value*2-1). All via gr.File so "
                "Gradio doesn't re-encode them.")

    with gr.Tab("Depth"):
        d_in = gr.Image(label="Input", type="filepath")
        d_out = gr.File(label="Depth (16-bit PNG)")
        d_btn = gr.Button("Run depth")
        d_btn.click(infer_depth, inputs=d_in, outputs=d_out, api_name="depth")

    with gr.Tab("Normal"):
        n_in = gr.Image(label="Input", type="filepath")
        n_out = gr.File(label="Normal (lossless PNG)")
        n_btn = gr.Button("Run normal")
        n_btn.click(infer_normal, inputs=n_in, outputs=n_out, api_name="normal")

    with gr.Tab("Normal (float16)"):
        n16_in = gr.Image(label="Input", type="filepath")
        n16_out = gr.File(label="Normal (float16 .npy)")
        n16_btn = gr.Button("Run normal (float16)")
        n16_btn.click(infer_normal16, inputs=n16_in, outputs=n16_out,
                      api_name="normal16")


if __name__ == "__main__":
    demo.queue(max_size=10).launch(server_name="0.0.0.0", server_port=7860,
                                    show_error=True)