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]


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_summary(stack: SliceStack, source_name: str) -> str:
    (x_min, y_min, z_min), (x_max, y_max, z_max) = stack.bounds
    return "\n".join(
        [
            "### Slice Stack Ready",
            f"- Source: `{source_name}`",
            f"- TIFF count: `{len(stack.tiff_paths)}`",
            f"- Image size: `{stack.image_size[0]} x {stack.image_size[1]}` pixels",
            f"- Layer height: `{stack.layer_height}`",
            f"- Pixel size: `{stack.pixel_size}`",
            f"- Bounds: `x={x_min:.3f}..{x_max:.3f}`, `y={y_min:.3f}..{y_max:.3f}`, `z={z_min:.3f}..{z_max:.3f}`",
        ]
    )


def _format_model_status(source_name: str) -> str:
    return "\n".join(
        [
            "### Model Loaded",
            f"- Source: `{source_name}`",
            "- Rotate the model in the 3D viewer, then choose slice settings and generate the TIFF stack.",
        ]
    )


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, str | None]:
    tiff_paths = state.get("tiff_paths", [])
    if not tiff_paths:
        return "No slice stack loaded yet.", None, 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),
        selected_path,
    )


def load_model_assets(stl_file: str | None):
    if not stl_file:
        return (
            "Upload an STL file to begin.",
            _empty_state(),
            _reset_slider(),
            "No slice stack loaded yet.",
            None,
            None,
            None,
            None,
            "No model loaded yet.",
        )

    mesh = load_mesh(stl_file)
    source_name = Path(stl_file).name

    return (
        _format_model_status(source_name),
        _empty_state(),
        _reset_slider(),
        "No slice stack loaded yet.",
        None,
        None,
        None,
        stl_file,
        _format_model_details(source_name, mesh),
    )


def generate_stack(
    stl_file: str | None,
    layer_height: float,
    pixel_size: float,
    progress: gr.Progress = gr.Progress(),
):
    if not stl_file:
        raise gr.Error("Upload an STL file before generating slices.")

    def report_progress(current: int, total: int) -> None:
        progress(current / total, desc=f"Rendering slice {current} of {total}")

    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, selected_path = _render_selected_slice(state, 0)
    slider_update = gr.update(
        minimum=0,
        maximum=max(0, len(stack.tiff_paths) - 1),
        value=0,
        step=1,
        interactive=len(stack.tiff_paths) > 1,
    )

    return (
        _format_summary(stack, Path(stl_file).name),
        state,
        slider_update,
        label,
        preview,
        str(stack.zip_path),
        selected_path,
    )


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


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

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


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="STL TIFF Slicer") as demo:
        gr.Markdown(
            """
            # STL to TIFF Slicer
            Upload an STL to inspect it in a rotatable 3D viewer.
            Then choose a layer height and XY pixel size, generate the TIFF stack, and browse the slices below.
            """
        )

        state = gr.State(_empty_state())

        with gr.Row():
            with gr.Column(scale=1):
                stl_file = gr.File(
                    label="STL File",
                    file_types=[".stl"],
                    type="filepath",
                )
                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 Stack", variant="primary")
                download_zip = gr.File(label="Download All TIFFs (ZIP)")
                current_tiff = gr.File(label="Current TIFF Slice")

            with gr.Column(scale=2):
                summary = gr.Markdown("Upload an STL file to begin.")
                model_details = gr.Markdown("No model loaded yet.")
                model_viewer = gr.Model3D(
                    label="Interactive 3D Viewer",
                    display_mode="solid",
                    clear_color=(0.94, 0.95, 0.97, 1.0),
                    height=360,
                )
                slice_label = gr.Markdown("No slice stack loaded yet.")
                slice_preview = gr.Image(
                    label="Slice Preview",
                    type="pil",
                    image_mode="L",
                    height=420,
                )
                with gr.Row():
                    prev_button = gr.Button("Previous Slice")
                    next_button = gr.Button("Next Slice")
                slice_slider = gr.Slider(
                    label="Slice Index",
                    minimum=0,
                    maximum=0,
                    value=0,
                    step=1,
                    interactive=False,
                )

        stl_file.change(
            fn=load_model_assets,
            inputs=stl_file,
            outputs=[
                summary,
                state,
                slice_slider,
                slice_label,
                slice_preview,
                download_zip,
                current_tiff,
                model_viewer,
                model_details,
            ],
        )

        generate_button.click(
            fn=generate_stack,
            inputs=[stl_file, layer_height, pixel_size],
            outputs=[
                summary,
                state,
                slice_slider,
                slice_label,
                slice_preview,
                download_zip,
                current_tiff,
            ],
        )

        slice_slider.release(
            fn=jump_to_slice,
            inputs=[state, slice_slider],
            outputs=[slice_label, slice_preview, current_tiff],
            queue=False,
        )

        prev_button.click(
            fn=lambda state_value, index: shift_slice(state_value, index, -1),
            inputs=[state, slice_slider],
            outputs=[slice_slider, slice_label, slice_preview, current_tiff],
            queue=False,
        )

        next_button.click(
            fn=lambda state_value, index: shift_slice(state_value, index, 1),
            inputs=[state, slice_slider],
            outputs=[slice_slider, slice_label, slice_preview, current_tiff],
            queue=False,
        )

    return demo


demo = build_demo()


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