"""
Portrait to 3D Pipeline — Hugging Face Space
Step 1 : FLUX.2 [dev] API → repositions portrait to 3/4 face + white studio lighting
Step 2 : Hunyuan3D-2.1 (hy3dshape) → generates a 3D model (GLB)
"""

import os
import sys
import io
import time
import base64
import random
import traceback
from pathlib import Path

import spaces
import gradio as gr
import torch
import numpy as np
from PIL import Image

OUTPUT_DIR = Path("/tmp/face2mesh_outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

MAX_SEED = np.iinfo(np.int32).max

# ──────────────────────────────────────────────────────────────
# Default prompts
# ──────────────────────────────────────────────────────────────
EDIT_PROMPT = (
    "Turn the person faces for a 3/4 face portrait, "
    "maintaining exact body proportions and facial identity. "
    "Apply clean white studio lighting coming from the upper right, "
    "with soft shadows on the left side. "
    "Change the background to light grey. "
    "Do not alter clothing, hairstyle or skin tone. "
    "In a 3d volumetric style."
)


# ──────────────────────────────────────────────────────────────
# FLUX.2 via Gradio API
# ──────────────────────────────────────────────────────────────
def run_flux_edit(
    input_image: Image.Image,
    prompt: str,
    seed: int,
    guidance_scale: float,
    num_steps: int,
    prompt_upsampling: bool,
    progress,
) -> Image.Image:
    """
    Call the black-forest-labs/FLUX.2-dev Gradio Space API.
    Replaces the previous local pipeline + remote text encoder approach.
    """
    from gradio_client import Client, handle_file

    # Save the input PIL image to a temp file so handle_file can reference it
    ts = int(time.time())
    tmp_input_path = str(OUTPUT_DIR / f"flux_input_{ts}.png")
    input_image.save(tmp_input_path)

    # Compute output dimensions matching input aspect ratio
    orig_w, orig_h = input_image.size
    aspect = orig_w / orig_h
    if aspect >= 1:
        out_w = 1024
        out_h = int(1024 / aspect)
    else:
        out_h = 1024
        out_w = int(1024 * aspect)
    out_w = max(256, min(1024, round(out_w / 8) * 8))
    out_h = max(256, min(1024, round(out_h / 8) * 8))

    progress(0.20, desc="[Step 1/2] Calling FLUX.2 API…")

    client = Client("black-forest-labs/FLUX.2-dev")
    result = client.predict(
        prompt=prompt,
        input_images=[
            {
                "image": handle_file(tmp_input_path),
                "caption": None,
            }
        ],
        seed=int(seed),
        randomize_seed=False,
        width=out_w,
        height=out_h,
        num_inference_steps=int(num_steps),
        guidance_scale=float(guidance_scale),
        prompt_upsampling=prompt_upsampling,
        api_name="/infer",
    )

    # result[0] is the image dict; result[1] is the used seed
    image_info = result[0]

    # The API can return a local path or a URL
    if image_info.get("path"):
        edited_image = Image.open(image_info["path"]).convert("RGB")
    elif image_info.get("url"):
        import urllib.request
        tmp_out_path = str(OUTPUT_DIR / f"flux_output_{ts}.png")
        urllib.request.urlretrieve(image_info["url"], tmp_out_path)
        edited_image = Image.open(tmp_out_path).convert("RGB")
    else:
        raise ValueError(f"FLUX.2 API returned unexpected image info: {image_info}")

    return edited_image


# ──────────────────────────────────────────────────────────────
# Hunyuan3D-2.1 globals
# ──────────────────────────────────────────────────────────────
_hunyuan_pipeline = None


def get_hunyuan_pipeline():
    global _hunyuan_pipeline
    if _hunyuan_pipeline is None:
        repo_root = Path(__file__).parent
        for sub in ("hy3dshape", "hy3dpaint"):
            p = str(repo_root / sub)
            if p not in sys.path:
                sys.path.insert(0, p)

        from hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline

        print("[INFO] Loading Hunyuan3D-2.1…")
        _hunyuan_pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
            "tencent/Hunyuan3D-2.1",
            subfolder="hunyuan3d-dit-v2-1",
            use_safetensors=False,
            device="cuda",
        )
        print("[INFO] Hunyuan3D-2.1 ready.")
    return _hunyuan_pipeline


# ──────────────────────────────────────────────────────────────
# Step 2 — Hunyuan3D-2.1: image → GLB mesh
# ──────────────────────────────────────────────────────────────
@spaces.GPU
def run_hunyuan(
    edited_image: Image.Image,
    num_steps: int,
    guidance_scale: float,
    octree_resolution: int,
    seed: int,
) -> str:
    from hy3dshape.rembg import BackgroundRemover
    rmbg_worker = BackgroundRemover()
    pipe = get_hunyuan_pipeline()
    ts = int(time.time())
    out_dir = OUTPUT_DIR / str(ts)
    out_dir.mkdir(parents=True, exist_ok=True)
    glb_path = str(out_dir / "model.glb")

    with torch.inference_mode():
        mesh = pipe(
            image=rmbg_worker(edited_image),
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            octree_resolution=octree_resolution,
            generator=torch.manual_seed(seed),
            output_type="trimesh",
        )[0]

    mesh.export(glb_path)
    return glb_path


# ──────────────────────────────────────────────────────────────
# Full pipeline orchestration
# ──────────────────────────────────────────────────────────────
def full_pipeline(
    input_image,
    edit_prompt,
    flux_seed,
    flux_guidance,
    flux_steps,
    flux_prompt_upsampling,
    hy_steps,
    hy_guidance,
    hy_octree_res,
    hy_seed,
    skip_flux,
    progress=gr.Progress(track_tqdm=True),
):
    if input_image is None:
        raise gr.Error("Please upload a portrait image.")

    logs = []

    try:
        # ── Step 1: FLUX.2 API edit ────────────────────────────
        if skip_flux:
            progress(0.20, desc="[Step 1 skipped] Using original image")
            edited_image = input_image
            logs.append("⏭  Step 1 skipped — original image passed to Hunyuan3D.")
        else:
            progress(0.05, desc="[Step 1/2] FLUX.2 API editing…")
            logs.append("🎨  FLUX.2 [dev] API — repositioning portrait + studio lighting…")
            edited_image = run_flux_edit(
                input_image=input_image,
                prompt=edit_prompt,
                seed=int(flux_seed),
                guidance_scale=float(flux_guidance),
                num_steps=int(flux_steps),
                prompt_upsampling=flux_prompt_upsampling,
                progress=progress,
            )
            logs.append("✅  Step 1 done.")
            progress(0.45, desc="[Step 1/2] Editing complete")

        # Save intermediate result
        ts = int(time.time())
        edited_path = str(OUTPUT_DIR / f"edited_{ts}.png")
        edited_image.save(edited_path)

        # ── Step 2: Hunyuan3D-2.1 ─────────────────────────────
        progress(0.50, desc="[Step 2/2] Generating 3D model…")
        logs.append("🧊  Hunyuan3D-2.1 — generating 3D mesh…")
        glb_path = run_hunyuan(
            edited_image=edited_image,
            num_steps=int(hy_steps),
            guidance_scale=float(hy_guidance),
            octree_resolution=int(hy_octree_res),
            seed=int(hy_seed),
        )
        logs.append("✅  Step 2 done.")
        logs.append(f"   • GLB: {glb_path}")
        progress(1.0, desc="Pipeline complete ✓")

        return edited_image, glb_path, "\n".join(logs)

    except Exception as e:
        logs.append(f"❌ Error: {e}\n\n{traceback.format_exc()}")
        raise gr.Error(str(e))


# ──────────────────────────────────────────────────────────────
# Gradio UI
# ──────────────────────────────────────────────────────────────
with gr.Blocks(
    title="Portrait → 3D Studio",
    theme=gr.themes.Soft(primary_hue="violet"),
) as demo:

    gr.Markdown(
        """
        ## 🧑‍🎨 Portrait → 3D Studio
        **Step 1** — FLUX.2 [dev] reshapes the portrait to a 3/4 angle with studio lighting.  
        **Step 2** — Hunyuan3D-2.1 converts the edited image into a 3D GLB mesh.
        """
    )

    with gr.Row():
        # ── Left column: inputs ───────────────────────────────
        with gr.Column(scale=1):
            gr.Markdown("### 📷 Source image")
            input_image = gr.Image(type="pil", label="Portrait photo", height=320)

            with gr.Accordion("⚙️ FLUX.2 settings (Step 1)", open=False):
                skip_flux = gr.Checkbox(
                    label="Skip FLUX.2 step (use image as-is)",
                    value=False,
                )
                edit_prompt = gr.Textbox(
                    label="Edit prompt",
                    value=EDIT_PROMPT,
                    lines=5,
                )
                flux_prompt_upsampling = gr.Checkbox(
                    label="Prompt upsampling (built-in FLUX.2 refinement)",
                    value=True,
                    info="Lets the FLUX.2 API refine the prompt internally before generation.",
                )
                with gr.Row():
                    flux_seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                    flux_steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
                flux_guidance = gr.Slider(
                    0.0, 10.0, value=4.0, step=0.1, label="Guidance scale"
                )

            with gr.Accordion("⚙️ Hunyuan3D-2.1 settings (Step 2)", open=False):
                with gr.Row():
                    hy_steps = gr.Slider(10, 50, value=30, step=1, label="DiT steps")
                    hy_guidance = gr.Slider(
                        1.0, 10.0, value=5.5, step=0.5, label="Guidance scale"
                    )
                with gr.Row():
                    hy_octree_res = gr.Slider(
                        256, 512, value=380, step=1, label="Octree resolution"
                    )
                    hy_seed = gr.Slider(0, 9999, value=0, step=1, label="3D seed")

            run_btn = gr.Button("🚀 Run pipeline", variant="primary", size="lg")

        # ── Right column: outputs ─────────────────────────────
        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ Edited image (FLUX.2)")
            edited_out = gr.Image(
                label="3/4-face portrait — studio lighting", height=300
            )

            gr.Markdown("### 🧊 3D model (Hunyuan3D-2.1)")
            glb_out = gr.File(label="GLB mesh (untextured)")

            log_out = gr.Textbox(
                label="📋 Execution log", lines=8, interactive=False
            )

    run_btn.click(
        fn=full_pipeline,
        inputs=[
            input_image,
            edit_prompt,
            flux_seed,
            flux_guidance,
            flux_steps,
            flux_prompt_upsampling,
            hy_steps,
            hy_guidance,
            hy_octree_res,
            hy_seed,
            skip_flux,
        ],
        outputs=[edited_out, glb_out, log_out],
    )

demo.launch()