from __future__ import annotations

import os
import subprocess
import sys
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import trimesh

from sphereparam import (
    REPO_ROOT,
    check_tools,
    compute_spherical_parameterization,
    load_mesh,
    normalize_for_display,
    required_tools,
    sphere_to_uv,
    uv_xy_colors,
)


EXAMPLE_MESH = REPO_ROOT / "demos" / "data" / "bunny.orig.m"
MAX_VERTICES = int(os.environ.get("SPHEREPARAM_MAX_VERTICES", "50000"))


def ensure_executable_permissions() -> None:
    for path in [
        REPO_ROOT / "demos" / "bin" / "meshtopm.sh",
        REPO_ROOT / "bin" / "unix" / "MeshSimplify",
        REPO_ROOT / "bin" / "unix" / "reverselines",
        REPO_ROOT / "bin" / "unix" / "Filterprog",
        REPO_ROOT / "bin" / "unix" / "SphereParam",
    ]:
        if path.exists():
            path.chmod(path.stat().st_mode | 0o755)


def ensure_native_tools() -> None:
    ensure_executable_permissions()
    try:
        check_tools()
        return
    except FileNotFoundError:
        pass

    env = os.environ.copy()
    env["PYTHON_BIN"] = sys.executable
    subprocess.run(["bash", str(REPO_ROOT / "compile.sh")], cwd=REPO_ROOT, check=True, env=env)
    ensure_executable_permissions()
    check_tools()


def rgba_colors(rgb: np.ndarray) -> np.ndarray:
    rgb255 = np.clip(np.asarray(rgb) * 255.0, 0, 255).astype(np.uint8)
    alpha = np.full((len(rgb255), 1), 255, dtype=np.uint8)
    return np.column_stack([rgb255, alpha])


def export_glb(path: Path, vertices: np.ndarray, faces: np.ndarray, colors: np.ndarray) -> Path:
    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
    mesh.visual.vertex_colors = rgba_colors(colors)
    mesh.export(path)
    return path


def export_scene_glb(
    path: Path,
    mesh_vertices: np.ndarray,
    sphere_vertices: np.ndarray,
    faces: np.ndarray,
    colors: np.ndarray,
) -> Path:
    scene = trimesh.Scene()
    original = trimesh.Trimesh(vertices=mesh_vertices + np.array([-1.35, 0.0, 0.0]), faces=faces, process=False)
    original.visual.vertex_colors = rgba_colors(colors)
    sphere = trimesh.Trimesh(vertices=sphere_vertices + np.array([1.35, 0.0, 0.0]), faces=faces, process=False)
    sphere.visual.vertex_colors = rgba_colors(colors)
    scene.add_geometry(original, node_name="input mesh")
    scene.add_geometry(sphere, node_name="sphere")
    scene.export(path)
    return path


@spaces.GPU(duration=120)
def parameterize(mesh_file: str, effort: int, base_faces: int, threads: int) -> tuple[str, str, str, str, str]:
    if not mesh_file:
        raise gr.Error("Upload a closed triangular genus-0 mesh first.")

    mesh_path = Path(mesh_file)
    ensure_native_tools()

    vertices, faces = load_mesh(mesh_path)
    if len(vertices) > MAX_VERTICES:
        raise gr.Error(
            f"This demo is capped at {MAX_VERTICES:,} vertices; uploaded mesh has {len(vertices):,}."
        )

    base_faces_arg = None if int(base_faces) <= 0 else int(base_faces)
    threads_arg = None if int(threads) <= 0 else int(threads)
    sphere_vertices = compute_spherical_parameterization(
        vertices,
        faces,
        effort=int(effort),
        base_faces=base_faces_arg,
        threads=threads_arg,
        verbose=False,
    )

    uv = sphere_to_uv(sphere_vertices)
    colors = uv_xy_colors(uv)
    display_vertices = normalize_for_display(vertices)

    out_dir = Path(tempfile.mkdtemp(prefix="sphereparam_space_"))
    npz_path = out_dir / f"{mesh_path.stem}_sphereparam.npz"
    np.savez_compressed(
        npz_path,
        vertices=vertices,
        faces=faces,
        sphere_vertices=sphere_vertices,
        uv=uv,
        source_mesh=str(mesh_path),
        effort=int(effort),
        base_faces=-1 if base_faces_arg is None else base_faces_arg,
    )

    original_glb = export_glb(out_dir / "input_uv_colored.glb", display_vertices, faces, colors)
    sphere_glb = export_glb(out_dir / "sphere_uv_colored.glb", sphere_vertices, faces, colors)
    comparison_glb = export_scene_glb(out_dir / "input_and_sphere.glb", display_vertices, sphere_vertices, faces, colors)

    status = (
        f"Done: {len(vertices):,} vertices, {len(faces):,} faces. "
        "Colors encode spherical UV coordinates and are shared between the input mesh and sphere."
    )
    return str(comparison_glb), str(original_glb), str(sphere_glb), str(npz_path), status


def available_tools_status() -> str:
    tools = required_tools()
    present = [name for name, path in tools.items() if path.exists()]
    missing = [name for name, path in tools.items() if not path.exists()]
    if not missing:
        return "Native tools are compiled and ready."
    return "Native tools will compile on first run. Missing now: " + ", ".join(missing)


with gr.Blocks(title="Spherical Parameterization") as demo:
    gr.Markdown(
        "# Spherical Parameterization\n"
        "Upload a closed genus-0 triangle mesh, run Hoppe/Praun spherical parameterization, "
        "and inspect the UV-colored mesh and sphere as GLB previews."
    )

    with gr.Row():
        with gr.Column(scale=1):
            mesh_input = gr.File(
                label="Mesh",
                file_types=[".m", ".ply", ".obj", ".off", ".stl", ".glb", ".gltf"],
                type="filepath",
            )
            effort = gr.Slider(0, 5, value=2, step=1, label="Effort")
            base_faces = gr.Number(value=0, precision=0, label="Base faces, 0 for default")
            threads = gr.Number(value=0, precision=0, label="Threads, 0 for default")
            run = gr.Button("Parameterize", variant="primary")
            status = gr.Textbox(label="Status", value=available_tools_status(), interactive=False)
            npz_output = gr.File(label="Download .npz")
        with gr.Column(scale=2):
            comparison = gr.Model3D(label="Input and sphere")

    with gr.Row():
        original = gr.Model3D(label="Input mesh")
        sphere = gr.Model3D(label="Sphere")

    gr.Examples(
        examples=[[str(EXAMPLE_MESH), 2, 0, 0]],
        inputs=[mesh_input, effort, base_faces, threads],
    )

    run.click(
        parameterize,
        inputs=[mesh_input, effort, base_faces, threads],
        outputs=[comparison, original, sphere, npz_output, status],
    )


if __name__ == "__main__":
    demo.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"))