from __future__ import annotations

import gradio as gr

from ai_runtime import MODEL_OMNI, MODEL_TEXT3D, iter_hunyuan_blueprint_session, finalize_ai_mesh_session
from fallback_generator import iter_blueprint_session as iter_fallback_blueprint
from fallback_generator import iter_meshify_session as iter_fallback_mesh
from viewer import empty_viewer_html, load_points_from_cloud_file, point_cloud_viewer_html

TITLE = "Particle Blueprint 3D"
TAGLINE = "Prompt → AI blueprint → inspect on iPhone → export mesh"
SUBTITLE = (
    "Rebuilt around Hugging Face-hosted Tencent Hunyuan 3D models. "
    "The first stage uses Hunyuan3D-1 for prompt-driven 3D generation, then converts that AI output into a particle blueprint for review."
)

PROMPTS = {
    "Cargo hauler": "small cargo hauler with a boxy hull, rear ramp, cargo bay, 4 engines and landing gear",
    "Compact fighter": "compact fighter with a sleek hull, twin engines, short wings and a small cockpit",
    "Industrial shuttle": "industrial shuttle with a rounded hull, cargo hold, fin tail and landing gear",
}

CSS = """
footer {display:none !important}
.gradio-container {max-width: 1080px !important; margin: 0 auto; padding: 0 12px 28px !important}
.hero, .panel, .status-card {border:1px solid rgba(255,255,255,.08); border-radius:22px; background:rgba(255,255,255,.03); padding:14px 16px}
.hero-title {font-size:1.5rem; font-weight:800; margin:0 0 6px 0}
.hero-sub {opacity:.88; margin:0}
.flow {display:grid; grid-template-columns:repeat(4, minmax(0,1fr)); gap:10px; margin-top:10px}
.flow div {border:1px solid rgba(255,255,255,.08); border-radius:16px; padding:12px; background:rgba(255,255,255,.02)}
.flow strong {display:block; margin-bottom:4px}
.cta-row button {min-height:52px !important; border-radius:16px !important; font-size:1rem !important}
.preset-row button {min-height:44px !important; border-radius:14px !important}
#status-box p {margin:0}
.small-note {font-size:.88rem; opacity:.8}
@media (max-width: 820px) {
  .gradio-container {padding:0 10px 22px !important}
  .hero,.panel,.status-card {padding:12px}
  .flow {grid-template-columns:1fr}
  .cta-row {position:sticky; bottom:10px; z-index:12; background:rgba(10,13,22,.92); backdrop-filter:blur(12px); padding:8px; border:1px solid rgba(255,255,255,.08); border-radius:18px}
}
"""


def _status_md(text: str) -> str:
    return f"**Status**\n\n{text}"


def _blueprint_html_from_ply(path: str | None) -> str:
    if not path:
        return empty_viewer_html()
    try:
        points = load_points_from_cloud_file(path)
        return point_cloud_viewer_html(points, status=f"Blueprint • {len(points)} points")
    except Exception as exc:
        return empty_viewer_html(f"Blueprint file exists but the viewer could not read it: {exc}")


def stream_blueprint(prompt: str, mode: str, save_memory: bool, max_faces: int, detail: int):
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Enter a prompt first.")

    html = empty_viewer_html("Starting…")
    summary = None
    state = None
    blueprint_file = None
    mesh_preview = None

    yield (
        _status_md("Starting blueprint generation…"),
        html,
        None,
        None,
        None,
        gr.update(interactive=False),
        None,
    )

    if mode == "hunyuan":
        try:
            for update in iter_hunyuan_blueprint_session(
                prompt=prompt,
                save_memory=save_memory,
                max_faces_num=int(max_faces),
                preview_points=max(2200, int(detail) * 140),
            ):
                html = update.get("viewer_html", html)
                summary = update.get("summary", summary)
                blueprint_file = update.get("blueprint_path", blueprint_file)
                mesh_preview = update.get("mesh_preview", mesh_preview)
                state = update.get("state", state)
                yield (
                    _status_md(update.get("status", "Working…")),
                    html,
                    summary,
                    blueprint_file,
                    state,
                    gr.update(interactive=bool(state)),
                    mesh_preview,
                )
            return
        except Exception as exc:
            fallback_notice = f"AI model path failed, so the app is falling back to the local scaffold builder. Reason: {exc}"
            yield (
                _status_md(fallback_notice),
                html,
                {"fallback_reason": str(exc), "requested_mode": mode, "prompt": prompt},
                blueprint_file,
                state,
                gr.update(interactive=False),
                mesh_preview,
            )

    for update in iter_fallback_blueprint(prompt=prompt, detail=detail, parser_mode="heuristic"):
        blueprint_path = update.get("blueprint_path")
        if blueprint_path:
            html = _blueprint_html_from_ply(blueprint_path)
            blueprint_file = blueprint_path
        summary = update.get("summary", summary)
        state = update.get("state", state)
        yield (
            _status_md(update.get("status", "Working…")),
            html,
            summary,
            blueprint_file,
            state,
            gr.update(interactive=bool(state)),
            mesh_preview,
        )


def stream_mesh(state: dict, prepare_omni: bool, voxel_pitch: float):
    if not state:
        raise gr.Error("Generate a blueprint first.")

    yield _status_md("Starting mesh generation…"), None, None, None

    if state.get("raw_ai_mesh_path"):
        for update in finalize_ai_mesh_session(state, prepare_omni=prepare_omni):
            yield (
                _status_md(update.get("status", "Meshing…")),
                update.get("mesh_path"),
                update.get("mesh_file"),
                update.get("summary"),
            )
        return

    for update in iter_fallback_mesh(state=state, voxel_pitch=voxel_pitch, use_target_model_cache=prepare_omni):
        yield (
            _status_md(update.get("status", "Meshing…")),
            update.get("mesh_path"),
            update.get("mesh_file"),
            update.get("summary"),
        )


with gr.Blocks(title=TITLE, fill_width=True) as demo:
    session_state = gr.State(value=None)

    gr.HTML(
        f"""
        <div class='hero'>
          <div class='hero-title'>{TITLE}</div>
          <p class='hero-sub'><strong>{TAGLINE}</strong><br>{SUBTITLE}</p>
          <div class='flow'>
            <div><strong>1. Describe</strong>Write what you want.</div>
            <div><strong>2. Generate blueprint</strong>{MODEL_TEXT3D} runs, then its AI mesh is turned into a particle blueprint.</div>
            <div><strong>3. Inspect</strong>Rotate, zoom and pan the blueprint on iPhone before committing.</div>
            <div><strong>4. Make mesh</strong>Export the AI mesh as GLB, with {MODEL_OMNI} preloaded for later refinement work.</div>
          </div>
        </div>
        """
    )

    with gr.Column(elem_classes=["panel"]):
        prompt = gr.Textbox(
            label="Describe the model",
            lines=4,
            max_lines=7,
            placeholder="Example: small cargo hauler with a boxy hull, cargo bay, rear ramp, 4 engines and landing gear",
        )

        with gr.Row(elem_classes=["preset-row"]):
            p1 = gr.Button("Cargo hauler")
            p2 = gr.Button("Compact fighter")
            p3 = gr.Button("Industrial shuttle")

        with gr.Accordion("Generation settings", open=False):
            mode = gr.Radio(
                choices=[("Hunyuan3D-1 AI", "hunyuan"), ("Local fallback scaffold", "fallback")],
                value="hunyuan",
                label="Blueprint generation mode",
                info="Use Hunyuan first. If its repo dependencies are not fully ready in the Space yet, the app will fall back automatically.",
            )
            save_memory = gr.Checkbox(
                value=True,
                label="Use save-memory mode for Hunyuan3D-1",
                info="Useful on tighter ZeroGPU runs.",
            )
            max_faces = gr.Slider(20000, 90000, value=70000, step=5000, label="Max faces for the AI mesh")
            detail = gr.Slider(14, 34, value=22, step=2, label="Blueprint preview density")
            voxel_pitch = gr.Slider(0.055, 0.12, value=0.085, step=0.005, label="Fallback mesh density")
            prepare_omni = gr.Checkbox(
                value=True,
                label=f"Preload {MODEL_OMNI} during mesh step",
                info="This caches the controllable refinement model in the Space even though this build still exports the Hunyuan3D-1 mesh on step two.",
            )

        with gr.Row(elem_classes=["cta-row"]):
            blueprint_btn = gr.Button("Generate blueprint", variant="primary")
            mesh_btn = gr.Button("Make mesh", interactive=False)
            clear_btn = gr.Button("Clear")

    with gr.Column(elem_classes=["status-card"]):
        status = gr.Markdown(_status_md("Ready."), elem_id="status-box")
        gr.Markdown("<span class='small-note'>On iPhone: one finger orbits. Two fingers pan and zoom.</span>")

    with gr.Tabs():
        with gr.TabItem("Blueprint"):
            blueprint_view = gr.HTML(value=empty_viewer_html(), label="Blueprint viewer")
        with gr.TabItem("Mesh"):
            mesh_view = gr.Model3D(
                label="Mesh preview (.glb)",
                display_mode="solid",
                clear_color=(0.02, 0.02, 0.03, 1.0),
                camera_position=(35, 65, 6),
                zoom_speed=1.15,
                pan_speed=0.95,
                height=560,
            )
        with gr.TabItem("Summary and files"):
            summary = gr.JSON(label="Session summary")
            blueprint_file = gr.File(label="Particle blueprint (.ply)")
            mesh_file = gr.File(label="Mesh export (.glb)")

    p1.click(lambda: PROMPTS["Cargo hauler"], outputs=prompt)
    p2.click(lambda: PROMPTS["Compact fighter"], outputs=prompt)
    p3.click(lambda: PROMPTS["Industrial shuttle"], outputs=prompt)

    blueprint_btn.click(
        fn=stream_blueprint,
        inputs=[prompt, mode, save_memory, max_faces, detail],
        outputs=[status, blueprint_view, summary, blueprint_file, session_state, mesh_btn, mesh_view],
    )

    mesh_btn.click(
        fn=stream_mesh,
        inputs=[session_state, prepare_omni, voxel_pitch],
        outputs=[status, mesh_view, mesh_file, summary],
    )

    clear_btn.click(
        lambda: (
            "",
            _status_md("Ready."),
            empty_viewer_html(),
            None,
            None,
            None,
            None,
            gr.update(interactive=False),
            None,
        ),
        outputs=[prompt, status, blueprint_view, mesh_view, summary, blueprint_file, mesh_file, mesh_btn, session_state],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(theme=gr.themes.Soft(), css=CSS)