import os
import sys
import glob
import importlib.util
import subprocess

APP_DIR = os.path.dirname(os.path.abspath(__file__))

# TRELLIS runtime configuration (matches the reference setup: spconv + flash-attn).
os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("SPARSE_BACKEND", "spconv")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("SPARSE_ATTN_BACKEND", "flash_attn")
os.environ["SPACEFLOW_KEEP_ON_GPU"] = "1"  # all weights stay resident on the ZeroGPU device

# TRELLIS needs the mip-splatting fork of diff-gaussian-rasterization (kernel_size / subpixel_offset API).
# A prebuilt sm_120 wheel ships with the Space; install it before torch is imported.
if importlib.util.find_spec("diff_gaussian_rasterization") is None:
    wheel = glob.glob(os.path.join(APP_DIR, "wheels", "diff_gaussian_rasterization-*.whl"))[0]
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", wheel])

import spaces  # noqa: E402  (must come before torch)

import copy  # noqa: E402
import math  # noqa: E402
import json  # noqa: E402
import random  # noqa: E402
import shutil  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import trimesh  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402

sys.path.insert(0, APP_DIR)
os.chdir(APP_DIR)
os.makedirs("debug", exist_ok=True)

import run_local_tau as rlt  # noqa: E402
from third_party.TRELLIS.trellis.pipelines import TrellisTextTo3DPipeline  # noqa: E402
from third_party.PartField.partfield.model_trainer_pvcnn_only_demo import Model as PartFieldModel  # noqa: E402
from third_party.PartField.partfield.dataloader import Demo_Dataset  # noqa: E402
from lib.opt import self_similarity  # noqa: E402
from lib.util import pointcloud  # noqa: E402
from lib.util.checkpoint_compat import partfield_checkpoint_scope  # noqa: E402

CFG = OmegaConf.load("config/default.yaml")
LOW_CONTROL_BBOX_MARGIN = 0.07  # sq_ui default (DEFAULT_LOW_CONTROL_BBOX_MARGIN)

# ---------------------------------------------------------------------------
# Models (loaded once at module scope, placed on the GPU)
# ---------------------------------------------------------------------------
# TRELLIS' text conditioner encodes the empty prompt at init; defer that forward
# pass to the GPU worker instead of running it at import time.
_orig_encode_text = TrellisTextTo3DPipeline.encode_text
TrellisTextTo3DPipeline.encode_text = lambda self, text: None
PIPE = TrellisTextTo3DPipeline.from_pretrained("config/trellis_pipeline")
TrellisTextTo3DPipeline.encode_text = _orig_encode_text
PIPE.cuda()
PIPE.text_cond_model["model"].to("cuda")

PF_CFG = OmegaConf.load("third_party/PartField/config.yaml")
PARTFIELD = PartFieldModel(PF_CFG, None)
_pf_ckpt = hf_hub_download(
    "mikaelaangel/partfield-ckpt", "model_objaverse.ckpt",
    revision="90b9b1e08b6a12fdcb6ee26b4854a26235e1765f",
)
with partfield_checkpoint_scope():
    _pf_state = torch.load(_pf_ckpt, map_location="cpu", weights_only=True)["state_dict"]
_missing, _unexpected = PARTFIELD.load_state_dict(_pf_state, strict=False)
print(f"[PartField] missing keys: {len(_missing)}, unexpected keys: {len(_unexpected)}")
del _pf_state
PARTFIELD.eval().to("cuda")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _seed_all(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _ensure_null_cond():
    if PIPE.text_cond_model.get("null_cond") is None:
        PIPE.text_cond_model["null_cond"] = PIPE.encode_text([""])


def _read_metadata(npz_path):
    with np.load(npz_path, allow_pickle=False) as d:
        if "spaceflow_metadata.json" in d.files:
            try:
                return json.loads(bytes(d["spaceflow_metadata.json"]).decode("utf-8"))
            except Exception:
                return {}
    return {}


def _load_arrays(npz_path):
    with np.load(npz_path, allow_pickle=False) as d:
        arrays = {k: np.asarray(d[k]) for k in d.files if k != "spaceflow_metadata.json"}
    n = arrays["scales"].shape[0]
    arrays.setdefault("control_levels", np.ones((n,)))
    return arrays


def _subset(arrays, idx):
    out = {}
    for k in ["scales", "shapes", "translations", "rotations", "control_levels", "tapering", "bending"]:
        if k in arrays:
            out[k] = arrays[k][idx]
    return out


def _low_control_bbox(arrays, low_idx, margin=LOW_CONTROL_BBOX_MARGIN):
    """Port of sq_ui buildLowControlBoundingBoxPrimitive (padded AABB of the low-control primitives)."""
    verts = []
    for i in low_idx:
        tap = arrays["tapering"][i] if "tapering" in arrays else None
        bend = arrays["bending"][i] if "bending" in arrays else None
        v, _ = rlt.add_superquadric_compact_rot_mat(
            arrays["scales"][i], arrays["shapes"][i], arrays["translations"][i], arrays["rotations"][i],
            tap, bend, resolution=32)
        verts.append(v)
    v = np.concatenate(verts, 0)
    mn, mx = v.min(0), v.max(0)
    pad = np.linalg.norm(mx - mn) * margin
    mn, mx = mn - pad, mx + pad
    center = (mn + mx) / 2
    half = np.maximum((mx - mn) / 2, 1e-4)
    rot = np.array([[1.0, 0.0, 0.0], [0.0, 6.123234e-17, -1.0], [0.0, 1.0, 6.123234e-17]])
    return {
        "scales": np.array([[half[0], half[2], half[1]]]),
        "shapes": np.array([[0.05, 0.05]]),
        "translations": center[None],
        "rotations": rot[None],
        "control_levels": np.array([0.0]),
    }


def _parse_local_prompts(text, n_sq):
    """'6: pink painted wood' lines (1-based primitive index) -> per-primitive list."""
    prompts = [""] * n_sq
    for line in (text or "").replace(";", "\n").splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        idx, prompt = line.split(":", 1)
        idx = idx.strip().lower().replace("sq", "").strip()
        if not idx.isdigit():
            continue
        k = int(idx) - 1
        if 0 <= k < n_sq:
            prompts[k] = prompt.strip()
    return prompts


def _blender_normalized_ply(glb_path, ply_path):
    """Python port of TRELLIS' Blender `--save_mesh` export used by SpaceFlow.

    Blender imports glTF (Y-up) into its Z-up frame, scales the scene so the
    largest bbox side is 1, re-centres the bbox at the origin and writes a PLY.
    """
    mesh = trimesh.load(glb_path, force="mesh", process=False)
    v = np.asarray(mesh.vertices, dtype=np.float64)
    v = np.stack([v[:, 0], -v[:, 2], v[:, 1]], axis=1)
    mn, mx = v.min(0), v.max(0)
    v = v * (1.0 / (mx - mn).max())
    mn, mx = v.min(0), v.max(0)
    v = v - (mn + mx) / 2
    trimesh.Trimesh(vertices=v, faces=np.asarray(mesh.faces), process=False).export(ply_path)


# SpaceFlow / TRELLIS assets are Z-up (the authors' editor uses camera.up = +Z), while glTF and
# gr.Model3D are Y-up. Rotate (x, y, z) -> (x, z, -y) for display only.
_ZUP_TO_YUP = np.array([[1.0, 0.0, 0.0, 0.0],
                        [0.0, 0.0, 1.0, 0.0],
                        [0.0, -1.0, 0.0, 0.0],
                        [0.0, 0.0, 0.0, 1.0]])


def _zup_glb_to_yup(src, dst):
    """Bake a Z-up -> Y-up rotation into every mesh of a GLB (textures/UVs are kept)."""
    scene = trimesh.load(src, force="scene", process=False)
    out = trimesh.Scene()
    for node in scene.graph.nodes_geometry:
        transform, geom_name = scene.graph[node]
        geom = scene.geometry[geom_name].copy()
        geom.apply_transform(_ZUP_TO_YUP @ transform)
        out.add_geometry(geom, node_name=node, geom_name=geom_name)
    out.export(dst)
    return dst


def _partfield_planes(mesh_path, out_dir):
    """Equivalent of run_local_tau.predict_part without the Lightning Trainer."""
    _seed_all(int(PF_CFG.seed))
    item = Demo_Dataset(mesh_path)[0]
    batch = {"pc": item["pc"][None].to("cuda"), "uid": [item["uid"]]}
    with torch.autocast("cuda", dtype=torch.float16):
        part_planes, uid = PARTFIELD.predict_step(batch, 0)
    np.save(os.path.join(out_dir, f"part_feat_{uid}_batch_part_plane.npy"), part_planes)


def preview_primitives(primitives_file):
    """Render the uploaded primitives (orange = high control, white = low control) and read saved prompts."""
    if primitives_file is None:
        return None, gr.update(), gr.update(), gr.update(), ""
    path = primitives_file if isinstance(primitives_file, str) else primitives_file.name
    out = tempfile.mkdtemp()
    glb = os.path.join(out, "input_superquadrics.glb")
    rlt.export_colored_superquadrics_glb(path, glb)
    # Show in a Y-up viewer the same way the result is oriented.
    glb = _zup_glb_to_yup(glb, os.path.join(out, "input_superquadrics_yup.glb"))
    meta = _read_metadata(path)
    arrays = _load_arrays(path)
    n = arrays["scales"].shape[0]
    names = meta.get("primitiveNames") or [f"SQ {i + 1}" for i in range(n)]
    levels = ["high" if float(c) >= 0.5 else "low" for c in arrays["control_levels"]]
    local = meta.get("localTextureTexts") or [""] * n
    rows = "\n".join(
        f"| {i + 1} | {names[i] if i < len(names) else ''} | {levels[i]} | {local[i] if i < len(local) else ''} |"
        for i in range(n))
    table = "| # | name | control | local texture |\n|---|---|---|---|\n" + rows
    shape = meta.get("textPrompt")
    tex = meta.get("globalTextureText")
    local_txt = "\n".join(f"{i + 1}: {t}" for i, t in enumerate(local) if str(t or "").strip())
    return (
        glb,
        gr.update(value=shape) if shape else gr.update(),
        gr.update(value=tex) if tex else gr.update(),
        gr.update(value=local_txt) if meta else gr.update(),
        table,
    )


def _duration(primitives_file, shape_prompt, texture_prompt, local_prompts="", low_tau=3.0, high_tau=10.0,
              polyak_tau=0.18, repaint_steps=10, texture_steps=300, seed=1, attn_boost=2.5, *args, **kwargs):
    # measured on ZeroGPU: ~35 s fixed (structure, PartField, decode/bake) + ~0.58 s per texture step; +15% margin
    return int(min(300, math.ceil((35 + 0.58 * int(texture_steps)) * 1.15)))


@spaces.GPU(duration=_duration)
def generate(
    primitives_file,
    shape_prompt: str,
    texture_prompt: str,
    local_prompts: str = "",
    low_tau: float = 3.0,
    high_tau: float = 10.0,
    polyak_tau: float = 0.18,
    repaint_steps: int = 10,
    texture_steps: int = 300,
    seed: int = 1,
    attn_boost: float = 2.5,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a textured 3D asset from superquadric primitives with SpaceFlow.

    Args:
        primitives_file: SpaceFlow primitives `.npz` (scales, shapes, translations, rotations, control_levels).
        shape_prompt: Text describing the object (e.g. "toy elephant").
        texture_prompt: Global appearance text (e.g. "blue painted wood").
        local_prompts: Optional per-primitive appearance overrides, one per line as "<index>: <text>" (1-based).
        low_tau: Structure-control strength for low-control primitives (lower = more freedom).
        high_tau: Structure-control strength for high-control primitives.
        polyak_tau: Polyak blending factor between the low- and high-control trajectories.
        repaint_steps: Repaint resampling steps used to blend the two control levels.
        texture_steps: Number of self-similarity guided texture flow steps.
        seed: Seed for sparse-structure sampling.
        attn_boost: Region-boosted self-attention factor for local texture routing (1.0 = off).

    Returns:
        Path to the textured GLB and a short run log.
    """
    if primitives_file is None:
        raise gr.Error("Upload a primitives .npz file (or pick an example).")
    shape_prompt = (shape_prompt or "").strip()
    if not shape_prompt:
        raise gr.Error("Enter a shape prompt.")
    texture_prompt = (texture_prompt or "").strip() or shape_prompt
    t0 = time.perf_counter()
    path = primitives_file if isinstance(primitives_file, str) else primitives_file.name

    out_dir = tempfile.mkdtemp(prefix="spaceflow_")
    arrays = _load_arrays(path)
    n_sq = arrays["scales"].shape[0]
    all_path = os.path.join(out_dir, "all.npz")
    np.savez(all_path, **arrays)
    high_idx = [i for i in range(n_sq) if float(arrays["control_levels"][i]) >= 0.5]
    low_idx = [i for i in range(n_sq) if float(arrays["control_levels"][i]) < 0.5]
    local_mode = bool(high_idx) and bool(low_idx)
    if local_mode and float(high_tau) < float(low_tau):
        raise gr.Error("High-control tau must be >= low-control tau.")

    cfg = copy.deepcopy(CFG)
    cfg.sim_guidance.steps = max(2, int(texture_steps))
    cfg.sim_guidance.mixed_self_attn_boost = max(1.0, float(attn_boost))

    _seed_all(0)
    _ensure_null_cond()
    log = []

    # --- spatial-control meshes from the superquadrics -------------------------------------
    sc_mesh = os.path.join(out_dir, "spatial_control_mesh.ply")
    aabb, center, scale = rlt.load_superquadrics(all_path, sc_mesh)
    sampler = {
        "steps": rlt.STEPS_SHAPE_GEN,
        "cfg_strength": rlt.CFG_SHAPE_GEN,
        "t0_idx_value": float(low_tau),
        "spatial_control_mesh_path": sc_mesh,
        "high_control_spatial_control_mesh_path": None,
        "low_control_superquadric_mask_path": None,
        "t0_idx_value_high_control": None,
        "polyak_update_tau": float(polyak_tau) if local_mode else 0.0,
        "local_tau_mode": "guidance",
        "n_repaint_steps": int(repaint_steps),
    }
    if local_mode:
        hi_npz = os.path.join(out_dir, "high_control.npz")
        lo_npz = os.path.join(out_dir, "low_control_bbox.npz")
        np.savez(hi_npz, **_subset(arrays, high_idx))
        np.savez(lo_npz, **_low_control_bbox(arrays, low_idx))
        hi_mesh = os.path.join(out_dir, "high_control_spatial_control_mesh.ply")
        lo_mesh = os.path.join(out_dir, "low_control_superquadric_mask.ply")
        rlt.load_superquadrics(hi_npz, hi_mesh, aabb=aabb, center=center, scale=scale)
        rlt.load_superquadrics(lo_npz, lo_mesh, aabb=aabb, center=center, scale=scale)
        sampler.update({
            "high_control_spatial_control_mesh_path": hi_mesh,
            "low_control_superquadric_mask_path": lo_mesh,
            "t0_idx_value_high_control": float(high_tau),
            "local_tau_mode": "low_control_mask",
        })
        log.append(f"Local control: {len(high_idx)} high / {len(low_idx)} low primitives "
                   f"(tau {low_tau} / {high_tau}, polyak {polyak_tau}).")
    else:
        log.append(f"Global control: all {n_sq} primitives at tau {low_tau}.")

    # --- 1. SpaceControl sparse structure ---------------------------------------------------
    coords = PIPE.gen_structure_v2(shape_prompt, seed=int(seed), vis_output_dir=None,
                                   sparse_structure_sampler_params=sampler)
    voxels = coords.detach().cpu().numpy()[:, 1:]
    log.append(f"Structure: {voxels.shape[0]} voxels ({time.perf_counter() - t0:.1f}s).")
    rlt.sparse_voxels_to_glb(voxels, grid_size=64, output_filename=os.path.join(out_dir, "sample.glb"))
    struct_mesh = trimesh.load(os.path.join(out_dir, "sample.glb"), force="mesh")
    struct_mesh = pointcloud.convert_mesh_yup_to_zup(struct_mesh)  # --convert_yup_to_zup
    struct_zup = os.path.join(out_dir, "struct_mesh_zup.glb")
    struct_mesh.export(struct_zup)

    render_dir = os.path.join(out_dir, "struct_renders")
    os.makedirs(render_dir, exist_ok=True)
    struct_ply = os.path.join(render_dir, "mesh.ply")
    _blender_normalized_ply(struct_zup, struct_ply)
    voxel_dir = os.path.join(out_dir, "voxels")
    os.makedirs(voxel_dir, exist_ok=True)
    pointcloud.voxelize_mesh(struct_ply, save_path=os.path.join(voxel_dir, "struct_voxels.ply"))

    # --- 2. PartField features on the generated structure -----------------------------------
    pf_dir = os.path.join(out_dir, "partfield")
    os.makedirs(pf_dir, exist_ok=True)
    _partfield_planes(struct_ply, pf_dir)

    # --- 3. self-similarity guided appearance flow with local texture routing ---------------
    local_list = _parse_local_prompts(local_prompts, n_sq)
    active = [i + 1 for i, p in enumerate(local_list) if p]
    local_list = local_list if active else None
    individual = rlt.build_individual_sq_meshes_normalized(all_path) if local_list else None
    if active:
        log.append("Local texture overrides on primitives " + ", ".join(map(str, active)) + ".")
    self_similarity.optimize_self_similarity(
        cfg, texture_prompt, "text", out_dir,
        local_prompts=local_list,
        local_prompt_type="text" if local_list else None,
        individual_sq_meshes=individual,
        generation_pipeline=PIPE,
        decode_texture=True,
    )
    out_glb = _zup_glb_to_yup(os.path.join(out_dir, "out_sim.glb"), os.path.join(out_dir, "out_sim_yup.glb"))
    for wp in glob.glob(os.path.join(out_dir, "*warning*.json")):
        try:
            w = json.load(open(wp))
            items = w if isinstance(w, list) else w.get("warnings", [w])
            for it in items:
                if isinstance(it, dict) and it.get("message"):
                    log.append("Warning: " + it["message"])
        except Exception:
            pass
    log.append(f"Done in {time.perf_counter() - t0:.1f}s.")
    print("[SpaceFlow] " + " ".join(log))
    return out_glb, "\n".join(f"- {line}" for line in log)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
EXAMPLES = [
    ["examples/toy_elephant.npz", "toy elephant", "blue painted wood", "6: pink painted wood\n9: pink painted wood"],
    ["examples/sailboat.npz", "Sailboat", "white", None],
    ["examples/chair.npz", "A chair", "light wood", None],
    ["examples/hot_air_balloon.npz", "a vintage hot air balloon with a wicker basket",
     "warm vintage canvas, leather trim, and soft sunlit shading", None],
    ["examples/toy_rocket.npz", "a toy rocket", "white metal", None],
    ["examples/windmill.npz", "a rustic windmill with four wooden blades",
     "weathered white plaster walls with dark wooden beams", None],
]
# Fill local prompts for the examples from the authors' saved metadata.
for _ex in EXAMPLES:
    if _ex[3] is None:
        _m = _read_metadata(_ex[0])
        _ex[3] = "\n".join(f"{i + 1}: {t}" for i, t in enumerate(_m.get("localTextureTexts", [])) if str(t).strip())

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="SpaceFlow") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# SpaceFlow: Locally Controllable 3D Generation\n"
            "Training-free 3D generation from **superquadric primitives**: each primitive is marked as "
            "*high* (stick to it) or *low* (only a loose hint) geometry control, and can carry its own "
            "**local texture prompt**. Built on TRELLIS-text-xlarge + PartField.  \n"
            "[Project page](https://spaceflow3d.github.io/) · [Paper](https://arxiv.org/abs/2610.12399) · "
            "[Code](https://github.com/SpaceFlow3D/spaceflow) · Author primitives in the "
            "[SpaceFlow editor](https://github.com/SpaceFlow3D/spaceflow#interactive-editor) and upload its `all.npz` here."
        )
        with gr.Row():
            with gr.Column(scale=1):
                primitives = gr.File(label="Primitives (.npz from the SpaceFlow editor)", file_types=[".npz"],
                                     type="filepath")
                input_view = gr.Model3D(label="Input primitives (orange = high control, white = low control)",
                                        height=300, interactive=False)
                shape_prompt = gr.Textbox(label="Shape prompt", placeholder="toy elephant")
                texture_prompt = gr.Textbox(label="Global texture prompt", placeholder="blue painted wood")
                local_prompts = gr.Textbox(
                    label="Local texture prompts (optional)", lines=3,
                    placeholder="One per line, '<primitive #>: <texture>', e.g.\n6: pink painted wood")
                run_btn = gr.Button("Generate", variant="primary")
                with gr.Accordion("Primitive list", open=False):
                    prim_table = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    with gr.Row():
                        low_tau = gr.Slider(0, 12, value=3.0, step=0.5, label="Low-control tau")
                        high_tau = gr.Slider(0, 12, value=10.0, step=0.5, label="High-control tau")
                    with gr.Row():
                        polyak_tau = gr.Slider(0, 1, value=0.18, step=0.01, label="Polyak blend")
                        repaint_steps = gr.Slider(0, 20, value=10, step=1, label="Repaint steps")
                    with gr.Row():
                        texture_steps = gr.Slider(50, 300, value=300, step=10, label="Texture flow steps")
                        attn_boost = gr.Slider(1.0, 4.0, value=2.5, step=0.1, label="Local attention boost")
                    seed = gr.Number(value=1, precision=0, label="Structure seed")
            with gr.Column(scale=1):
                output = gr.Model3D(label="Generated textured mesh (GLB)", height=520)
                info = gr.Markdown()

        inputs = [primitives, shape_prompt, texture_prompt, local_prompts, low_tau, high_tau, polyak_tau,
                  repaint_steps, texture_steps, seed, attn_boost]
        primitives.change(preview_primitives, inputs=[primitives],
                          outputs=[input_view, shape_prompt, texture_prompt, local_prompts, prim_table],
                          api_visibility="private")
        run_btn.click(generate, inputs=inputs, outputs=[output, info], api_name="generate")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[primitives, shape_prompt, texture_prompt, local_prompts],
            outputs=[output, info],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)