from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import gradio as gr
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download, snapshot_download
from PIL import Image, ImageDraw


MODEL_REPO = os.getenv("PIXLRELIGHT_ONNX_REPO", "Reza2kn/PIXLRelight-ONNX")
MODEL_PATH = os.getenv("PIXLRELIGHT_ONNX_PATH")
MODEL_VARIANT = os.getenv("PIXLRELIGHT_ONNX_VARIANT", "int4").lower()
MODEL_FILES = {
    "int4": "int4/pixlrelight_renderer_int4.onnx",
    "fp16": "webgpu/pixlrelight_renderer_fp16.onnx",
    "fp32": "pixlrelight_renderer.onnx",
}
ONNX_FILENAME = MODEL_FILES.get(MODEL_VARIANT, MODEL_FILES["int4"])
SIZE = 512
BASE_DIR = Path(__file__).resolve().parent
gr.set_static_paths([BASE_DIR])
SAMPLES = {
    "Warm room": str(BASE_DIR / "samples" / "room00_source.png"),
    "Studio chair": str(BASE_DIR / "samples" / "room01_source.png"),
}


DEFAULT_LIGHTS = [
    {"x": 0.28, "y": 0.22, "radius": 0.42, "intensity": 1.25, "color": "#ffd08a"},
    {"x": 0.78, "y": 0.62, "radius": 0.32, "intensity": 0.85, "color": "#86b7ff"},
]


CSS = """
.gradio-container {
  max-width: none !important;
  padding: 0 !important;
}
footer {display: none !important}
#webgpu-shell {
  min-height: 100vh;
  background: #f8fafc;
}
#webgpu-shell iframe {
  display: block;
  width: 100%;
  height: 100vh;
  min-height: 920px;
  border: 0;
  background: #f8fafc;
}
"""


def _hex_to_rgb01(value: str) -> np.ndarray:
    value = (value or "#ffffff").strip().lstrip("#")
    if len(value) != 6:
        value = "ffffff"
    return np.array([int(value[i : i + 2], 16) for i in (0, 2, 4)], dtype=np.float32) / 255.0


def _load_image(image: Any, sample_name: str | None = None) -> Image.Image:
    if image is None:
        image = SAMPLES.get(sample_name or "Warm room", next(iter(SAMPLES.values())))
    if isinstance(image, Image.Image):
        pil = image.convert("RGB")
    else:
        pil = Image.open(image).convert("RGB")
    return pil


def _resize_square(pil: Image.Image) -> Image.Image:
    pil = pil.convert("RGB")
    w, h = pil.size
    side = min(w, h)
    left = (w - side) // 2
    top = (h - side) // 2
    pil = pil.crop((left, top, left + side, top + side))
    return pil.resize((SIZE, SIZE), Image.Resampling.LANCZOS)


def _lights_from_state(state: str | list[dict[str, Any]] | None) -> list[dict[str, Any]]:
    if state is None:
        return [dict(x) for x in DEFAULT_LIGHTS]
    if isinstance(state, str):
        try:
            state = json.loads(state)
        except Exception:
            return [dict(x) for x in DEFAULT_LIGHTS]
    if not isinstance(state, list):
        return [dict(x) for x in DEFAULT_LIGHTS]
    clean = []
    for light in state[:5]:
        clean.append(
            {
                "x": float(np.clip(light.get("x", 0.5), 0.0, 1.0)),
                "y": float(np.clip(light.get("y", 0.5), 0.0, 1.0)),
                "radius": float(np.clip(light.get("radius", 0.35), 0.05, 1.0)),
                "intensity": float(np.clip(light.get("intensity", 1.0), 0.0, 2.5)),
                "color": str(light.get("color", "#ffffff")),
            }
        )
    return clean or [dict(x) for x in DEFAULT_LIGHTS]


def render_light_map(state: str | list[dict[str, Any]] | None, selected: int = 0):
    lights = _lights_from_state(state)
    yy, xx = np.mgrid[0:SIZE, 0:SIZE].astype(np.float32)
    light_map = np.zeros((SIZE, SIZE, 3), dtype=np.float32)
    ambient = np.full_like(light_map, 0.055)
    light_map += ambient
    for light in lights:
        cx = light["x"] * (SIZE - 1)
        cy = light["y"] * (SIZE - 1)
        sigma = max(1.0, light["radius"] * SIZE * 0.42)
        falloff = np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2.0 * sigma * sigma))
        light_map += falloff[..., None] * _hex_to_rgb01(light["color"]) * light["intensity"]
    light_map = light_map / max(1.0, float(light_map.max()))

    preview = Image.fromarray((np.clip(light_map, 0, 1) * 255).astype(np.uint8), "RGB")
    draw = ImageDraw.Draw(preview)
    for idx, light in enumerate(lights):
        cx = int(light["x"] * (SIZE - 1))
        cy = int(light["y"] * (SIZE - 1))
        r = 11 if idx == selected else 8
        draw.ellipse((cx - r, cy - r, cx + r, cy + r), outline="white", width=3)
        draw.ellipse((cx - r + 3, cy - r + 3, cx + r - 3, cy + r - 3), fill=tuple((_hex_to_rgb01(light["color"]) * 255).astype(int)))
    return preview


def _make_intrinsics(source_arr: np.ndarray, state: str | list[dict[str, Any]] | None) -> np.ndarray:
    light_img = np.asarray(render_light_map(state)).astype(np.float32) / 255.0
    albedo = np.clip(source_arr * 0.85 + 0.15, 0, 1)
    irradiance = light_img
    residual = np.clip((light_img - 0.5) * 0.45 + source_arr * 0.12 + 0.5, 0, 1)
    x = np.concatenate([albedo, irradiance, residual], axis=2)
    return np.transpose(x, (2, 0, 1))[None].astype(np.float32)


_SESSION: ort.InferenceSession | None = None


def get_session() -> ort.InferenceSession:
    global _SESSION
    if _SESSION is None:
        if MODEL_PATH:
            model_path = MODEL_PATH
        else:
            repo_dir = snapshot_download(
                MODEL_REPO,
                allow_patterns=[ONNX_FILENAME, ONNX_FILENAME + ".data"],
            )
            model_path = str(Path(repo_dir) / ONNX_FILENAME)
        providers = ["CPUExecutionProvider"]
        _SESSION = ort.InferenceSession(model_path, providers=providers)
    return _SESSION


def load_sample(sample_name: str):
    pil = _load_image(None, sample_name)
    state = json.dumps([dict(x) for x in DEFAULT_LIGHTS])
    return pil, state, render_light_map(state, 0), 0, 0, 0.28, 0.22, 0.42, 1.25, "#ffd08a"


def select_light(selected: int, state: str):
    lights = _lights_from_state(state)
    selected = int(np.clip(selected, 0, len(lights) - 1))
    light = lights[selected]
    return selected, light["x"], light["y"], light["radius"], light["intensity"], light["color"], render_light_map(lights, selected)


def update_light(selected: int, x: float, y: float, radius: float, intensity: float, color: str, state: str):
    lights = _lights_from_state(state)
    selected = int(np.clip(selected, 0, len(lights) - 1))
    lights[selected].update(x=float(x), y=float(y), radius=float(radius), intensity=float(intensity), color=color)
    return json.dumps(lights), render_light_map(lights, selected)


def add_light(state: str):
    lights = _lights_from_state(state)
    if len(lights) < 5:
        lights.append({"x": 0.5, "y": 0.35, "radius": 0.28, "intensity": 0.9, "color": "#ffffff"})
    selected = len(lights) - 1
    light = lights[selected]
    return json.dumps(lights), selected, selected, light["x"], light["y"], light["radius"], light["intensity"], light["color"], render_light_map(lights, selected)


def remove_light(selected: int, state: str):
    lights = _lights_from_state(state)
    if len(lights) > 1:
        lights.pop(int(np.clip(selected, 0, len(lights) - 1)))
    selected = int(np.clip(selected, 0, len(lights) - 1))
    light = lights[selected]
    return json.dumps(lights), selected, selected, light["x"], light["y"], light["radius"], light["intensity"], light["color"], render_light_map(lights, selected)


def move_light(evt: gr.SelectData, selected: int, state: str):
    lights = _lights_from_state(state)
    selected = int(np.clip(selected, 0, len(lights) - 1))
    x, y = evt.index
    lights[selected]["x"] = float(np.clip(x / max(1, SIZE - 1), 0, 1))
    lights[selected]["y"] = float(np.clip(y / max(1, SIZE - 1), 0, 1))
    light = lights[selected]
    return json.dumps(lights), light["x"], light["y"], render_light_map(lights, selected)


def relight(image, sample_name: str, state: str):
    source = _resize_square(_load_image(image, sample_name))
    source_arr = np.asarray(source).astype(np.float32) / 255.0
    source_tensor = np.transpose(source_arr, (2, 0, 1))[None].astype(np.float32)
    intrinsics = _make_intrinsics(source_arr, state)
    outputs = get_session().run(
        ["rgb", "rgb_gain", "rgb_bias"],
        {"source_images": source_tensor, "target_intrinsics": intrinsics},
    )
    pred = np.clip(outputs[0][0].transpose(1, 2, 0), 0, 1)
    gain = outputs[1][0].transpose(1, 2, 0)
    bias = outputs[2][0].transpose(1, 2, 0)
    gain_vis = np.clip((gain - gain.min()) / max(1e-6, gain.max() - gain.min()), 0, 1)
    bias_vis = np.clip((bias - bias.min()) / max(1e-6, bias.max() - bias.min()), 0, 1)
    return (
        Image.fromarray((pred * 255).astype(np.uint8), "RGB"),
        Image.fromarray((gain_vis * 255).astype(np.uint8), "RGB"),
        Image.fromarray((bias_vis * 255).astype(np.uint8), "RGB"),
    )


with gr.Blocks(css=CSS, theme=gr.themes.Base()) as demo:
    gr.HTML(
        f"""
        <div id="webgpu-shell">
          <iframe
            title="PIXLRelight WebGPU INT4 interactive light studio"
            src="/gradio_api/file={BASE_DIR / 'webgpu.html'}"
            allow="webgpu; fullscreen"
          ></iframe>
        </div>
        """
    )


if __name__ == "__main__":
    demo.launch(show_error=True)