app.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import gradio as gr
from PIL import Image

from stl_slicer import SliceStack, load_mesh, slice_stl_to_tiffs


ViewerState = dict[str, Any]
SAMPLE_STL_FILENAMES = ("Hollow_Pyramid.stl", "balanced_die.stl", "halfsphere.stl")
SAMPLE_STL_DIR = Path(__file__).resolve().parent / "sample_stls"
APP_CSS = """
.gradio-container {
    font-size: 90%;
    padding-top: 0.5rem !important;
    padding-bottom: 0.5rem !important;
}

.gradio-container .gr-row {
    gap: 0.5rem !important;
}

.gradio-container .gr-form,
.gradio-container .gr-box,
.gradio-container .block {
    padding: 0.4rem !important;
}

.gradio-container .prose {
    margin-bottom: 0.4rem !important;
}

.model3D button[aria-label="Undo"] {
    color: var(--block-label-text-color) !important;
    cursor: pointer !important;
    opacity: 1 !important;
}
"""

# Gradio 6.10's gr.Model3D leaves the Undo (reset view) button permanently
# disabled when the value is supplied programmatically — its `has_change_history`
# state only flips on uploads through Model3D's own upload widget. This script
# strips the disabled attribute so clicks reach Svelte's handle_undo, which
# calls reset_camera_position on the underlying canvas.
APP_HEAD = """
<script>
(function () {
    function enableUndoButtons(root) {
        (root || document).querySelectorAll('.model3D button[aria-label="Undo"]').forEach(function (btn) {
            if (btn.disabled) {
                btn.disabled = false;
            }
        });
    }
    function start() {
        enableUndoButtons();
        var observer = new MutationObserver(function (mutations) {
            for (var i = 0; i < mutations.length; i++) {
                var m = mutations[i];
                if (m.type === 'attributes' && m.target && m.target.matches && m.target.matches('.model3D button[aria-label="Undo"]')) {
                    if (m.target.disabled) m.target.disabled = false;
                } else if (m.type === 'childList') {
                    enableUndoButtons(m.target);
                }
            }
        });
        observer.observe(document.body, {
            childList: true,
            subtree: true,
            attributes: true,
            attributeFilter: ['disabled']
        });
    }
    if (document.readyState === 'loading') {
        document.addEventListener('DOMContentLoaded', start);
    } else {
        start();
    }
})();
</script>
"""


def _read_slice_preview(path: str) -> Image.Image:
    with Image.open(path) as image:
        return image.copy()


def _empty_state() -> ViewerState:
    return {"tiff_paths": [], "z_values": []}


def _reset_slider() -> dict[str, Any]:
    return gr.update(minimum=0, maximum=0, value=0, step=1, interactive=False)


def _stack_to_state(stack: SliceStack) -> ViewerState:
    return {
        "tiff_paths": [str(path) for path in stack.tiff_paths],
        "z_values": stack.z_values,
    }


def _format_model_details(source_name: str, mesh) -> str:
    extents = mesh.extents
    return "\n".join(
        [
            "### Model Details",
            f"- Source: `{source_name}`",
            f"- Extents: `{extents[0]:.3f} x {extents[1]:.3f} x {extents[2]:.3f}`",
            f"- Faces: `{len(mesh.faces)}`",
            f"- Vertices: `{len(mesh.vertices)}`",
            f"- Watertight: `{'yes' if mesh.is_watertight else 'no'}`",
        ]
    )


def _slice_label(state: ViewerState, index: int) -> str:
    path = Path(state["tiff_paths"][index]).name
    z_value = state["z_values"][index]
    total = len(state["tiff_paths"])
    return f"Slice {index + 1} / {total} | z = {z_value:.4f} | {path}"


def _render_selected_slice(state: ViewerState, index: int) -> tuple[str, Image.Image | None]:
    tiff_paths = state.get("tiff_paths", [])
    if not tiff_paths:
        return "No slice stack loaded yet.", None

    bounded_index = max(0, min(int(index), len(tiff_paths) - 1))
    selected_path = tiff_paths[bounded_index]
    return (
        _slice_label(state, bounded_index),
        _read_slice_preview(selected_path),
    )


def load_single_model(stl_file: str | None) -> tuple[str | None, str]:
    if not stl_file:
        return None, "No model loaded."
    mesh = load_mesh(stl_file)
    return stl_file, _format_model_details(Path(stl_file).name, mesh)


def preload_sample_models() -> tuple:
    outputs: list[Any] = []

    for filename in SAMPLE_STL_FILENAMES:
        stl_path = SAMPLE_STL_DIR / filename
        if not stl_path.exists():
            outputs.extend([
                None,
                None,
                f"Sample file not found: {stl_path}",
            ])
            continue

        try:
            mesh = load_mesh(stl_path)
        except Exception as exc:
            outputs.extend([
                str(stl_path),
                None,
                f"Failed to load sample model: {stl_path.name} ({exc})",
            ])
            continue

        outputs.extend([
            str(stl_path),
            str(stl_path),
            _format_model_details(stl_path.name, mesh),
        ])

    return tuple(outputs)


def generate_all_stacks(
    stl1: str | None,
    stl2: str | None,
    stl3: str | None,
    layer_height: float,
    pixel_size: float,
    progress: gr.Progress = gr.Progress(),
):
    files = [stl1, stl2, stl3]
    valid_count = max(1, sum(1 for f in files if f))
    results: list = []
    completed = 0

    for stl_file in files:
        if not stl_file:
            results.extend([
                _empty_state(),
                _reset_slider(),
                "No slice stack loaded yet.",
                None,
                None,
            ])
            continue

        slot_offset = completed

        def report_progress(cur: int, tot: int, offset: int = slot_offset) -> None:
            progress(
                (offset + cur / tot) / valid_count,
                desc=f"Slicing object {offset + 1} of {valid_count}\u2026",
            )

        stack = slice_stl_to_tiffs(
            stl_file,
            layer_height=layer_height,
            pixel_size=pixel_size,
            progress_callback=report_progress,
        )
        state = _stack_to_state(stack)
        label, preview = _render_selected_slice(state, 0)
        slider = gr.update(
            minimum=0,
            maximum=max(0, len(stack.tiff_paths) - 1),
            value=0,
            step=1,
            interactive=len(stack.tiff_paths) > 1,
        )
        results.extend([
            state,
            slider,
            label,
            preview,
            str(stack.zip_path),
        ])
        completed += 1

    return tuple(results)


def jump_to_slice(state: ViewerState, index: float) -> tuple[str, Image.Image | None]:
    return _render_selected_slice(state, int(index))


def shift_slice(state: ViewerState, index: float, delta: int) -> tuple[int, str, Image.Image | None]:
    tiff_paths = state.get("tiff_paths", [])
    if not tiff_paths:
        return 0, "No slice stack loaded yet.", None

    new_index = max(0, min(int(index) + delta, len(tiff_paths) - 1))
    label, preview = _render_selected_slice(state, new_index)
    return new_index, label, preview


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="STL TIFF Slicer", css=APP_CSS, head=APP_HEAD) as demo:
        with gr.Tab("STL to TIFF Slicer"):
            gr.Markdown(
                """
                # STL to TIFF Slicer
                Upload up to three STL files, choose a shared layer height and XY pixel size, then generate TIFF stacks for all uploaded models.
                """
            )

            with gr.Row():
                load_samples_button = gr.Button(
                    "Load Sample STLs",
                    variant="secondary",
                    size="sm",
                    min_width=140,
                    scale=0,
                )

            # --- Upload + 3D viewer row ---
            stl_files: list[gr.File] = []
            model_viewers: list[gr.Model3D] = []
            model_details_list: list[gr.Markdown] = []

            with gr.Row():
                for i in range(3):
                    with gr.Column(min_width=250):
                        stl_file = gr.File(
                            label=f"STL File {i + 1}",
                            file_types=[".stl"],
                            type="filepath",
                        )
                        model_viewer = gr.Model3D(
                            label=f"3D Viewer {i + 1}",
                            display_mode="solid",
                            clear_color=(0.94, 0.95, 0.97, 1.0),
                            height=270,
                        )
                        model_details = gr.Markdown(f"No model {i + 1} loaded.")
                        stl_files.append(stl_file)
                        model_viewers.append(model_viewer)
                        model_details_list.append(model_details)

            # --- Shared slicing controls ---
            with gr.Row():
                layer_height = gr.Number(label="Layer Height", value=0.1, minimum=0.0001, step=0.01)
                pixel_size = gr.Number(label="Pixel Size", value=0.05, minimum=0.0001, step=0.01)
                generate_button = gr.Button("Generate TIFF Stacks", variant="primary")

            # --- Per-object slice browsers ---
            states: list[gr.State] = []
            sliders: list[gr.Slider] = []
            slice_labels: list[gr.Markdown] = []
            slice_previews: list[gr.Image] = []
            download_zips: list[gr.File] = []

            with gr.Row():
                for i in range(3):
                    with gr.Column(min_width=250):
                        slice_label = gr.Markdown("No slice stack loaded yet.")
                        slice_preview = gr.Image(
                            label=f"Slice Preview {i + 1}",
                            type="pil",
                            image_mode="L",
                            height=270,
                        )
                        with gr.Row():
                            prev_button = gr.Button("\u25c4 Prev", scale=1, min_width=90, size="sm")
                            next_button = gr.Button("Next \u25ba", scale=1, min_width=90, size="sm")
                        slice_slider = gr.Slider(
                            label="Slice Index",
                            minimum=0,
                            maximum=0,
                            value=0,
                            step=1,
                            interactive=False,
                        )
                        download_zip = gr.File(label=f"Download TIFF ZIP {i + 1}")
                        state = gr.State(_empty_state())

                        slice_labels.append(slice_label)
                        slice_previews.append(slice_preview)
                        sliders.append(slice_slider)
                        download_zips.append(download_zip)
                        states.append(state)

                        slice_slider.release(
                            fn=jump_to_slice,
                            inputs=[state, slice_slider],
                            outputs=[slice_label, slice_preview],
                            queue=False,
                        )
                        prev_button.click(
                            fn=lambda sv, idx: shift_slice(sv, idx, -1),
                            inputs=[state, slice_slider],
                            outputs=[slice_slider, slice_label, slice_preview],
                            queue=False,
                        )
                        next_button.click(
                            fn=lambda sv, idx: shift_slice(sv, idx, 1),
                            inputs=[state, slice_slider],
                            outputs=[slice_slider, slice_label, slice_preview],
                            queue=False,
                        )

            # --- File upload handlers ---
            for i in range(3):
                stl_files[i].change(
                    fn=load_single_model,
                    inputs=stl_files[i],
                    outputs=[model_viewers[i], model_details_list[i]],
                )

            # --- Generate button ---
            generate_outputs: list = []
            for i in range(3):
                generate_outputs.extend([
                    states[i],
                    sliders[i],
                    slice_labels[i],
                    slice_previews[i],
                    download_zips[i],
                ])

            preload_outputs: list = []
            for i in range(3):
                preload_outputs.extend([
                    stl_files[i],
                    model_viewers[i],
                    model_details_list[i],
                ])

            load_samples_button.click(
                fn=preload_sample_models,
                outputs=preload_outputs,
            )

            generate_button.click(
                fn=generate_all_stacks,
                inputs=[stl_files[0], stl_files[1], stl_files[2], layer_height, pixel_size],
                outputs=generate_outputs,
            )

        with gr.Tab("TIFF Slices to GCode"):
            gr.Textbox(
                label="Status",
                value="More to come later.",
                interactive=False,
            )

    return demo


demo = build_demo()


if __name__ == "__main__":
    demo.launch(ssr_mode=False)