app.py

from __future__ import annotations

import tempfile
import math
from pathlib import Path
from typing import Any

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import trimesh

from gcode_viewer import build_toolpath_figure, parse_gcode_path
from stl_slicer import SliceStack, load_mesh, slice_stl_to_tiffs
from tiff_to_gcode import generate_snake_path_gcode


ViewerState = dict[str, Any]
SAMPLE_STL_FILENAMES = ("Hollow_Pyramid.stl", "Rounded_Cube_Through_Holes.stl", "halfsphere.stl")
SAMPLE_STL_DIR = Path(__file__).resolve().parent / "sample_stls"
FRONT_CAMERA = (90, 80, None)
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;
}

.gcode-shape-card {
    border: 1px solid var(--border-color-primary);
    border-radius: 0.5rem;
    padding: 0.5rem !important;
    min-height: 220px;
}

.gcode-shape-card .prose {
    margin-bottom: 0.25rem !important;
}

.gcode-param-label {
    font-size: 0.8rem;
    font-weight: 600;
    line-height: 1.15;
    margin-bottom: 0.2rem !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:
        preview = image.copy()

    # Upscale low-resolution TIFF previews so they fill the viewer area better.
    min_display_side = 480
    width, height = preview.size
    max_dim = max(width, height)
    if max_dim > 0 and max_dim < min_display_side:
        scale = min_display_side / max_dim
        new_size = (
            max(1, int(round(width * scale))),
            max(1, int(round(height * scale))),
        )
        preview = preview.resize(new_size, resample=Image.Resampling.NEAREST)

    return preview


def _empty_state() -> ViewerState:
    return {
        "tiff_paths": [],
        "z_values": [],
        "pixel_size": 0.0,
        "x_min": 0.0,
        "y_min": 0.0,
        "image_width": 0,
        "image_height": 0,
    }


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:
    (x_min, y_min, _z_min), (_x_max, _y_max, _z_max) = stack.bounds
    return {
        "tiff_paths": [str(path) for path in stack.tiff_paths],
        "z_values": stack.z_values,
        "pixel_size": stack.pixel_size,
        "x_min": x_min,
        "y_min": y_min,
        "image_width": stack.image_size[0],
        "image_height": stack.image_size[1],
    }


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 _annotate_preview(
    image: Image.Image,
    pixel_size: float,
    x_min: float,
    y_min: float,
    orig_width: int,
    orig_height: int,
) -> Image.Image:
    """Draw a blue origin crosshair with axis labels and a scale bar."""
    rgb = image.convert("RGB")
    draw = ImageDraw.Draw(rgb)

    preview_w, preview_h = rgb.size
    scale_x = preview_w / orig_width if orig_width else 1.0
    scale_y = preview_h / orig_height if orig_height else 1.0

    BLUE = (50, 120, 255)

    try:
        font = ImageFont.load_default(size=14)
    except TypeError:
        font = ImageFont.load_default()
    try:
        small_font = ImageFont.load_default(size=12)
    except TypeError:
        small_font = font

    # --- Origin crosshair & axis indicators ---
    origin_px = (0.0 - x_min) / pixel_size
    origin_py_from_bottom = (0.0 - y_min) / pixel_size
    origin_img_y = orig_height - 1 - origin_py_from_bottom

    ox = int(round(origin_px * scale_x))
    oy = int(round(origin_img_y * scale_y))

    arm = 20
    margin_edge = 8  # inset from image border for off-screen indicators
    on_screen = 0 <= ox < preview_w and 0 <= oy < preview_h

    if on_screen:
        # +X axis (rightward)
        x_start = max(0, ox)
        x_end = min(preview_w - 1, ox + arm)
        if x_end > x_start:
            draw.line([(x_start, oy), (x_end, oy)], fill=BLUE, width=2)
            draw.polygon(
                [(x_end, oy), (x_end - 5, oy - 4), (x_end - 5, oy + 4)],
                fill=BLUE,
            )
            if x_end + 4 < preview_w:
                draw.text((x_end + 4, oy - 7), "X", fill=BLUE, font=small_font)

        # +Y axis (upward in world = upward in image)
        y_end = max(0, oy - arm)
        y_start = min(preview_h - 1, oy)
        if y_start > y_end:
            draw.line([(ox, y_start), (ox, y_end)], fill=BLUE, width=2)
            draw.polygon(
                [(ox, y_end), (ox - 4, y_end + 5), (ox + 4, y_end + 5)],
                fill=BLUE,
            )
            if y_end - 16 >= 0:
                draw.text((ox + 5, y_end - 16), "Y", fill=BLUE, font=small_font)

        # -X stub (leftward from origin)
        stub = min(8, max(0, ox))
        if stub > 0:
            draw.line([(ox - stub, oy), (ox, oy)], fill=BLUE, width=2)

        # -Y stub (downward from origin in image)
        stub_y = min(8, max(0, preview_h - 1 - oy))
        if stub_y > 0:
            draw.line([(ox, oy), (ox, oy + stub_y)], fill=BLUE, width=2)

        # Origin label
        lx = ox + arm + 4 if ox + arm + 40 < preview_w else ox - 45
        ly = oy + 6
        if 0 <= ly < preview_h:
            draw.text((max(0, lx), ly), "(0, 0)", fill=BLUE, font=small_font)

    else:
        # Origin is off-screen — draw edge indicator(s) pointing toward it.
        arrow_len = 14
        arrow_half = 5

        # Compute direction label text showing approximate origin coordinates
        origin_x_mm = x_min
        origin_y_mm = y_min
        coord_text = f"Origin ({-origin_x_mm:+.1f}, {-origin_y_mm:+.1f})"

        if ox < 0:
            # Origin is to the LEFT — draw left-pointing arrow on left edge
            ay = max(margin_edge + arrow_half, min(preview_h - margin_edge - arrow_half, oy))
            draw.polygon(
                [(margin_edge, ay), (margin_edge + arrow_len, ay - arrow_half), (margin_edge + arrow_len, ay + arrow_half)],
                fill=BLUE,
            )
            draw.text((margin_edge + arrow_len + 4, ay - 7), coord_text, fill=BLUE, font=small_font)
        elif ox >= preview_w:
            # Origin is to the RIGHT
            ay = max(margin_edge + arrow_half, min(preview_h - margin_edge - arrow_half, oy))
            rx = preview_w - margin_edge
            draw.polygon(
                [(rx, ay), (rx - arrow_len, ay - arrow_half), (rx - arrow_len, ay + arrow_half)],
                fill=BLUE,
            )
            tw = len(coord_text) * 7
            draw.text((max(0, rx - arrow_len - tw - 4), ay - 7), coord_text, fill=BLUE, font=small_font)

        if oy < 0:
            # Origin is ABOVE — draw upward-pointing arrow on top edge
            ax = max(margin_edge + arrow_half, min(preview_w - margin_edge - arrow_half, ox))
            draw.polygon(
                [(ax, margin_edge), (ax - arrow_half, margin_edge + arrow_len), (ax + arrow_half, margin_edge + arrow_len)],
                fill=BLUE,
            )
        elif oy >= preview_h:
            # Origin is BELOW — draw downward-pointing arrow on bottom edge
            ax = max(margin_edge + arrow_half, min(preview_w - margin_edge - arrow_half, ox))
            by = preview_h - margin_edge
            draw.polygon(
                [(ax, by), (ax - arrow_half, by - arrow_len), (ax + arrow_half, by - arrow_len)],
                fill=BLUE,
            )
            # If we didn't already draw a left/right label, label here
            if 0 <= ox < preview_w:
                draw.text((ax + arrow_half + 4, by - arrow_len - 2), coord_text, fill=BLUE, font=small_font)

    # --- Scale bar (bottom-left) ---
    image_width_mm = orig_width * pixel_size
    target_bar_mm = image_width_mm * 0.2
    nice = [0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100, 200, 500]
    bar_mm = min(nice, key=lambda v: abs(v - target_bar_mm))

    bar_px = (bar_mm / pixel_size) * scale_x
    margin = 12
    bar_y = preview_h - margin
    bar_x0 = margin
    bar_x1 = bar_x0 + bar_px
    cap = 5

    draw.line([(int(bar_x0), int(bar_y)), (int(bar_x1), int(bar_y))], fill=BLUE, width=3)
    draw.line([(int(bar_x0), int(bar_y - cap)), (int(bar_x0), int(bar_y + cap))], fill=BLUE, width=2)
    draw.line([(int(bar_x1), int(bar_y - cap)), (int(bar_x1), int(bar_y + cap))], fill=BLUE, width=2)

    bar_label = f"{bar_mm:g} mm"
    draw.text((int(bar_x0), int(bar_y - 20)), bar_label, fill=BLUE, font=font)

    return rgb


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]
    preview = _read_slice_preview(selected_path)

    pixel_size = state.get("pixel_size", 0.0)
    if pixel_size and pixel_size > 0:
        preview = _annotate_preview(
            preview,
            pixel_size=pixel_size,
            x_min=state.get("x_min", 0.0),
            y_min=state.get("y_min", 0.0),
            orig_width=state.get("image_width", 0) or preview.size[0],
            orig_height=state.get("image_height", 0) or preview.size[1],
        )

    return (
        _slice_label(state, bounded_index),
        preview,
    )


def _opacity_to_alpha(opacity: float) -> int:
    bounded = max(0.05, min(float(opacity), 1.0))
    return int(round(255 * bounded))


def _resolve_model_opacity(setting: float | bool | None) -> float:
    if isinstance(setting, bool):
        return 0.75 if setting else 1.0
    if setting is None:
        return 1.0
    return max(0.05, min(float(setting), 1.0))


def _viewer_update(model_path: str | None) -> dict[str, Any]:
    return gr.update(value=model_path, camera_position=FRONT_CAMERA)


def _build_annotated_scene(mesh: trimesh.Trimesh, opacity: float = 1.0) -> str:
    """Export a GLB containing the mesh, origin axes, and a Z=0 grid plane."""
    scene = trimesh.Scene()
    display_transform = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])

    # --- Model (muted orange to match the Gradio theme accent) ---
    model_copy = mesh.copy()
    model_copy.apply_transform(display_transform)
    bounded_opacity = _resolve_model_opacity(opacity)
    mat = trimesh.visual.material.PBRMaterial(
        baseColorFactor=[230, 150, 90, _opacity_to_alpha(bounded_opacity)],
        alphaMode="OPAQUE" if bounded_opacity >= 0.999 else "BLEND",
        metallicFactor=0.0,
        roughnessFactor=0.6,
    )
    model_copy.visual = trimesh.visual.TextureVisuals(material=mat)
    scene.add_geometry(model_copy, geom_name="model")

    bounds = mesh.bounds
    (x_min, y_min, z_min), (x_max, y_max, z_max) = bounds
    extent = max(x_max - x_min, y_max - y_min, z_max - z_min)

    # --- Origin axes (coloured cylinders + cones) ---
    axis_len = extent * 0.4
    axis_radius = extent * 0.008
    cone_radius = axis_radius * 3.5
    cone_height = axis_len * 0.12

    axis_defs = [
        ("X", [1, 0, 0], [255, 50, 50, 255]),
        ("Y", [0, 1, 0], [50, 200, 50, 255]),
        ("Z", [0, 0, 1], [50, 120, 255, 255]),
    ]

    for name, direction, color in axis_defs:
        d = np.array(direction, dtype=float)

        # Cylinder from origin along axis
        cyl = trimesh.creation.cylinder(
            radius=axis_radius, height=axis_len, sections=12
        )
        # Default cylinder is along Z; rotate to desired axis
        midpoint = d * axis_len / 2
        if name == "X":
            cyl.apply_transform(trimesh.transformations.rotation_matrix(
                np.pi / 2, [0, 1, 0]
            ))
        elif name == "Y":
            cyl.apply_transform(trimesh.transformations.rotation_matrix(
                -np.pi / 2, [1, 0, 0]
            ))
        cyl.apply_translation(midpoint)
        cyl.apply_transform(display_transform)
        cyl.visual = trimesh.visual.ColorVisuals(
            mesh=cyl,
            face_colors=np.tile(color, (len(cyl.faces), 1)),
        )
        scene.add_geometry(cyl, geom_name=f"axis_{name}")

        # Cone arrowhead at tip
        cone = trimesh.creation.cone(
            radius=cone_radius, height=cone_height, sections=12
        )
        if name == "X":
            cone.apply_transform(trimesh.transformations.rotation_matrix(
                np.pi / 2, [0, 1, 0]
            ))
        elif name == "Y":
            cone.apply_transform(trimesh.transformations.rotation_matrix(
                -np.pi / 2, [1, 0, 0]
            ))
        cone.apply_translation(d * (axis_len + cone_height / 2))
        cone.apply_transform(display_transform)
        cone.visual = trimesh.visual.ColorVisuals(
            mesh=cone,
            face_colors=np.tile(color, (len(cone.faces), 1)),
        )
        scene.add_geometry(cone, geom_name=f"cone_{name}")

    # --- Grid plane at z=0 ---
    nice_spacings = [0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100]
    target_spacing = extent * 0.1
    grid_spacing = min(nice_spacings, key=lambda v: abs(v - target_spacing))

    # Grid extends to cover model footprint plus some margin
    margin = grid_spacing * 2
    gx_min = math.floor((x_min - margin) / grid_spacing) * grid_spacing
    gx_max = math.ceil((x_max + margin) / grid_spacing) * grid_spacing
    gy_min = math.floor((y_min - margin) / grid_spacing) * grid_spacing
    gy_max = math.ceil((y_max + margin) / grid_spacing) * grid_spacing

    grid_color = [160, 160, 160, 100]
    grid_segments: list[list[list[float]]] = []

    # Lines parallel to Y
    x = gx_min
    while x <= gx_max:
        grid_segments.append([[x, gy_min, 0], [x, gy_max, 0]])
        x += grid_spacing
    # Lines parallel to X
    y = gy_min
    while y <= gy_max:
        grid_segments.append([[gx_min, y, 0], [gx_max, y, 0]])
        y += grid_spacing

    if grid_segments:
        grid_path = trimesh.load_path(grid_segments)
        grid_path.apply_transform(display_transform)
        grid_path.colors = np.tile(grid_color, (len(grid_path.entities), 1))
        scene.add_geometry(grid_path, geom_name="grid")

    # Export to GLB (camera angle is set via gr.Model3D camera_position)
    out_path = Path(tempfile.mkdtemp(prefix="model3d_")) / "scene.glb"
    scene.export(str(out_path), file_type="glb")
    return str(out_path)


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


def preload_sample_models(opacity: float = 1.0) -> tuple:
    outputs: list[Any] = []
    resolved_opacity = _resolve_model_opacity(opacity)

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

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

        outputs.extend([
            str(stl_path),
            _viewer_update(_build_annotated_scene(mesh, opacity=resolved_opacity)),
            _format_model_details(stl_path.name, mesh),
        ])

    return tuple(outputs)


def refresh_all_model_viewers(
    stl1: str | None,
    stl2: str | None,
    stl3: str | None,
    opacity: float,
) -> tuple:
    outputs: list[Any] = []
    resolved_opacity = _resolve_model_opacity(opacity)
    for stl_file in (stl1, stl2, stl3):
        if not stl_file:
            outputs.extend([_viewer_update(None), "No model loaded."])
            continue
        outputs.extend(load_single_model(stl_file, resolved_opacity))
    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 run_all_tiff_to_gcode(
    zip1: str | None,
    zip2: str | None,
    zip3: str | None,
    pressure1: float,
    valve1: float,
    port1: float,
    pressure2: float,
    valve2: float,
    port2: float,
    pressure3: float,
    valve3: float,
    port3: float,
    layer_height: float = 0.8,
    pixel_size: float = 0.8,
) -> tuple[str | None, str | None, str | None, str]:
    specs = [
        (1, zip1, pressure1, valve1, port1),
        (2, zip2, pressure2, valve2, port2),
        (3, zip3, pressure3, valve3, port3),
    ]

    outputs: list[str | None] = [None, None, None]
    messages: list[str] = []

    for idx, zip_path, pressure, valve, port in specs:
        if not zip_path:
            messages.append(f"Shape {idx}: skipped (no TIFF ZIP available).")
            continue

        zip_name = Path(zip_path).stem
        default_shape_name = f"shape{idx}"
        shape_name = zip_name.replace("_tiff_slices", "") or default_shape_name

        try:
            gcode_path = generate_snake_path_gcode(
                zip_path=zip_path,
                shape_name=shape_name,
                pressure=float(pressure),
                valve=int(valve),
                port=int(port),
                layer_height=float(layer_height),
                fil_width=float(pixel_size),
            )
            outputs[idx - 1] = str(gcode_path)
            messages.append(f"Shape {idx}: wrote `{gcode_path.name}`.")
        except Exception as exc:  # surface errors in the UI
            outputs[idx - 1] = None
            messages.append(f"Shape {idx}: failed ({exc}).")

    return outputs[0], outputs[1], outputs[2], "\n".join(messages)


GCODE_SOURCE_SHAPE1 = "Use Shape 1 G-Code"
GCODE_SOURCE_UPLOAD = "Upload G-Code file"


def toggle_gcode_source(source: str) -> dict[str, Any]:
    return gr.update(interactive=(source == GCODE_SOURCE_UPLOAD))


def render_toolpath(
    source: str,
    uploaded_path: str | None,
    shape1_path: str | None,
    travel_opacity: float = 0.55,
    print_opacity: float = 1.0,
    travel_color: str = "#969696",
    print_color: str = "#1f77b4",
) -> tuple[Any, str, dict]:
    if source == GCODE_SOURCE_UPLOAD:
        path = uploaded_path
        if not path:
            return None, "No G-code file uploaded yet.", {}
    else:
        path = shape1_path
        if not path:
            return None, "No Shape 1 G-code available yet. Generate it on the TIFF Slices to GCode tab first.", {}

    try:
        text = Path(path).read_text()
    except OSError as exc:
        return None, f"Failed to read G-code file: {exc}", {}

    parsed = parse_gcode_path(text)
    if parsed["point_count"] == 0:
        return None, "No G0/G1 movement lines found in the file.", {}

    figure = build_toolpath_figure(parsed, travel_opacity=travel_opacity, print_opacity=print_opacity, travel_color=travel_color, print_color=print_color)
    (x_min, y_min, z_min), (x_max, y_max, z_max) = parsed["bounds"]
    summary = (
        f"**{parsed['point_count']} moves parsed** — "
        f"{len(parsed['print_segments'])} print segment(s), "
        f"{len(parsed['travel_segments'])} travel segment(s).  \n"
        f"Bounds: X ∈ [{x_min:.2f}, {x_max:.2f}], "
        f"Y ∈ [{y_min:.2f}, {y_max:.2f}], "
        f"Z ∈ [{z_min:.2f}, {z_max:.2f}] mm."
    )
    return figure, summary, parsed


def update_toolpath_opacity(
    parsed: dict,
    travel_opacity: float,
    print_opacity: float,
) -> Any:
    if not parsed or not parsed.get("point_count"):
        return None
    return build_toolpath_figure(parsed, travel_opacity=travel_opacity, print_opacity=print_opacity)


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 generate_reference_stack(
    state1: ViewerState,
    state2: ViewerState,
    state3: ViewerState,
    progress: gr.Progress = gr.Progress(),
) -> tuple:
    """Combine all available TIFF stacks into a single reference stack.

    For each pixel in each layer the result is black (0) when *any* source
    stack has a black pixel at that position, and white (255) only when *all*
    sources are white.  Images of different sizes are centred on a canvas
    sized to the largest dimensions.
    """
    active_states = [s for s in [state1, state2, state3] if s.get("tiff_paths")]

    if not active_states:
        return (
            _empty_state(),
            _reset_slider(),
            "No TIFF stacks available. Generate TIFF stacks first.",
            None,
        )

    max_layers = max(len(s["tiff_paths"]) for s in active_states)

    # Determine the largest image dimensions across all stacks.
    max_width = 0
    max_height = 0
    source_sizes: list[tuple[int, int]] = []
    for state in active_states:
        w = state.get("image_width", 0)
        h = state.get("image_height", 0)
        if not w or not h:
            with Image.open(state["tiff_paths"][0]) as img:
                w, h = img.size
        source_sizes.append((w, h))
        max_width = max(max_width, w)
        max_height = max(max_height, h)

    # Compute annotation metadata from the first active state, accounting for
    # the centering offset applied to its image on the larger canvas.
    first = active_states[0]
    first_w, first_h = source_sizes[0]
    ref_pixel_size = first.get("pixel_size", 0.0)
    x_off_first = (max_width - first_w) // 2
    y_off_first = (max_height - first_h) // 2
    ref_x_min = first.get("x_min", 0.0) - x_off_first * ref_pixel_size
    ref_y_min = first.get("y_min", 0.0) - y_off_first * ref_pixel_size

    output_dir = Path(tempfile.mkdtemp(prefix="reference_stack_"))
    slices_dir = output_dir / "tiff_slices"
    slices_dir.mkdir(parents=True, exist_ok=True)

    tiff_paths: list[Path] = []
    z_values: list[float] = []

    for layer_idx in range(max_layers):
        progress(
            layer_idx / max_layers,
            desc=f"Compositing reference layer {layer_idx + 1}/{max_layers}",
        )

        # Start with an all-white canvas.
        ref_array = np.full((max_height, max_width), 255, dtype=np.uint8)

        for state in active_states:
            paths = state["tiff_paths"]
            if layer_idx >= len(paths):
                continue  # Stack exhausted – contributes white.

            with Image.open(paths[layer_idx]) as img:
                arr = np.asarray(img)

            h, w = arr.shape[:2]
            y_off = (max_height - h) // 2
            x_off = (max_width - w) // 2

            # Black (0) wins: pixel-wise minimum keeps any black pixel.
            region = ref_array[y_off : y_off + h, x_off : x_off + w]
            ref_array[y_off : y_off + h, x_off : x_off + w] = np.minimum(region, arr)

        ref_image = Image.fromarray(ref_array, mode="L")
        tiff_path = slices_dir / f"ref_slice_{layer_idx:04d}.tif"
        ref_image.save(tiff_path, compression="tiff_deflate")
        tiff_paths.append(tiff_path)

        # Use z-value from the first active state that covers this layer.
        z_val = 0.0
        for state in active_states:
            if layer_idx < len(state["z_values"]):
                z_val = state["z_values"][layer_idx]
                break
        z_values.append(z_val)

    ref_state: ViewerState = {
        "tiff_paths": [str(p) for p in tiff_paths],
        "z_values": z_values,
        "pixel_size": ref_pixel_size,
        "x_min": ref_x_min,
        "y_min": ref_y_min,
        "image_width": max_width,
        "image_height": max_height,
    }

    label, preview = _render_selected_slice(ref_state, 0)
    slider = gr.update(
        minimum=0,
        maximum=max(0, len(tiff_paths) - 1),
        value=0,
        step=1,
        interactive=len(tiff_paths) > 1,
    )

    return ref_state, slider, 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,
                )
                with gr.Column(scale=0, min_width=240):
                    model_opacity = gr.Checkbox(
                        label="Use 75% 3D Model Opacity",
                        value=False,
                    )

            # --- 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),
                            camera_position=FRONT_CAMERA,
                            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.8, minimum=0.0001, step=0.01)
                pixel_size = gr.Number(
                    label="Pixel Size/Fill Width",
                    value=0.8,
                    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="RGB",
                            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}", interactive=False)
                        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,
                        )

            # --- Reference TIFF Stack ---
            gr.Markdown("---")
            gr.Markdown("### Reference TIFF Stack")

            with gr.Row():
                with gr.Column(scale=1, min_width=200):
                    ref_generate_button = gr.Button(
                        "Generate Reference TIFF Stack",
                        variant="primary",
                    )
                with gr.Column(scale=3, min_width=250):
                    ref_slice_label = gr.Markdown("No reference stack generated yet.")
                    ref_slice_preview = gr.Image(
                        label="Reference Slice Preview",
                        type="pil",
                        image_mode="RGB",
                        height=270,
                    )
                    with gr.Row():
                        ref_prev_button = gr.Button("\u25c4 Prev", scale=1, min_width=90, size="sm")
                        ref_next_button = gr.Button("Next \u25ba", scale=1, min_width=90, size="sm")
                    ref_slice_slider = gr.Slider(
                        label="Slice Index",
                        minimum=0,
                        maximum=0,
                        value=0,
                        step=1,
                        interactive=False,
                    )
            ref_state = gr.State(_empty_state())

            ref_slice_slider.release(
                fn=jump_to_slice,
                inputs=[ref_state, ref_slice_slider],
                outputs=[ref_slice_label, ref_slice_preview],
                queue=False,
            )
            ref_prev_button.click(
                fn=lambda sv, idx: shift_slice(sv, idx, -1),
                inputs=[ref_state, ref_slice_slider],
                outputs=[ref_slice_slider, ref_slice_label, ref_slice_preview],
                queue=False,
            )
            ref_next_button.click(
                fn=lambda sv, idx: shift_slice(sv, idx, 1),
                inputs=[ref_state, ref_slice_slider],
                outputs=[ref_slice_slider, ref_slice_label, ref_slice_preview],
                queue=False,
            )

            # --- File upload handlers ---
            for i in range(3):
                stl_files[i].change(
                    fn=load_single_model,
                    inputs=[stl_files[i], model_opacity],
                    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,
                inputs=[model_opacity],
                outputs=preload_outputs,
            )

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

            model_opacity.change(
                fn=refresh_all_model_viewers,
                inputs=[stl_files[0], stl_files[1], stl_files[2], model_opacity],
                outputs=refresh_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,
            )

            ref_generate_button.click(
                fn=generate_reference_stack,
                inputs=[states[0], states[1], states[2]],
                outputs=[ref_state, ref_slice_slider, ref_slice_label, ref_slice_preview],
            )

        with gr.Tab("TIFF Slices to GCode"):
            gr.Markdown(
                """
                # TIFF Slices to GCode
                Uses TIFF ZIP outputs from the first tab. Set pressure, valve,
                and port for each shape, then generate G-code files in one run.
                """
            )

            with gr.Row():
                with gr.Column(min_width=250):
                    with gr.Group(elem_classes=["gcode-shape-card"]):
                        gr.Markdown("### Shape 1")
                        with gr.Row():
                            with gr.Column(min_width=70):
                                gr.Markdown("Pressure (psi)", elem_classes=["gcode-param-label"])
                                gcode_pressure_1 = gr.Number(
                                    show_label=False,
                                    value=25.0,
                                    minimum=0.0,
                                    step=0.5,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Valve", elem_classes=["gcode-param-label"])
                                gcode_valve_1 = gr.Number(
                                    show_label=False,
                                    value=4,
                                    minimum=0,
                                    step=1,
                                    precision=0,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Port", elem_classes=["gcode-param-label"])
                                gcode_port_1 = gr.Number(
                                    show_label=False,
                                    value=1,
                                    minimum=1,
                                    step=1,
                                    precision=0,
                                )
                with gr.Column(min_width=250):
                    with gr.Group(elem_classes=["gcode-shape-card"]):
                        gr.Markdown("### Shape 2")
                        with gr.Row():
                            with gr.Column(min_width=70):
                                gr.Markdown("Pressure (psi)", elem_classes=["gcode-param-label"])
                                gcode_pressure_2 = gr.Number(
                                    show_label=False,
                                    value=25.0,
                                    minimum=0.0,
                                    step=0.5,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Valve", elem_classes=["gcode-param-label"])
                                gcode_valve_2 = gr.Number(
                                    show_label=False,
                                    value=4,
                                    minimum=0,
                                    step=1,
                                    precision=0,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Port", elem_classes=["gcode-param-label"])
                                gcode_port_2 = gr.Number(
                                    show_label=False,
                                    value=1,
                                    minimum=1,
                                    step=1,
                                    precision=0,
                                )
                with gr.Column(min_width=250):
                    with gr.Group(elem_classes=["gcode-shape-card"]):
                        gr.Markdown("### Shape 3")
                        with gr.Row():
                            with gr.Column(min_width=70):
                                gr.Markdown("Pressure (psi)", elem_classes=["gcode-param-label"])
                                gcode_pressure_3 = gr.Number(
                                    show_label=False,
                                    value=25.0,
                                    minimum=0.0,
                                    step=0.5,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Valve", elem_classes=["gcode-param-label"])
                                gcode_valve_3 = gr.Number(
                                    show_label=False,
                                    value=4,
                                    minimum=0,
                                    step=1,
                                    precision=0,
                                )
                            with gr.Column(min_width=70):
                                gr.Markdown("Port", elem_classes=["gcode-param-label"])
                                gcode_port_3 = gr.Number(
                                    show_label=False,
                                    value=1,
                                    minimum=1,
                                    step=1,
                                    precision=0,
                                )

            gcode_button = gr.Button("Generate G-Code", variant="primary")

            with gr.Row():
                gcode_file_1 = gr.File(label="Download G-Code Shape 1")
                gcode_file_2 = gr.File(label="Download G-Code Shape 2")
                gcode_file_3 = gr.File(label="Download G-Code Shape 3")

            gcode_status = gr.Markdown("")

            gcode_button.click(
                fn=run_all_tiff_to_gcode,
                inputs=[
                    download_zips[0],
                    download_zips[1],
                    download_zips[2],
                    gcode_pressure_1,
                    gcode_valve_1,
                    gcode_port_1,
                    gcode_pressure_2,
                    gcode_valve_2,
                    gcode_port_2,
                    gcode_pressure_3,
                    gcode_valve_3,
                    gcode_port_3,
                    layer_height,
                    pixel_size,
                ],
                outputs=[gcode_file_1, gcode_file_2, gcode_file_3, gcode_status],
            )

        with gr.Tab("G-Code Visualization"):
            gr.Markdown(
                "### 3D Tool-Path Viewer\n"
                "Choose a G-code source, then click **Render Tool Path** to visualize the nozzle path."
            )
            with gr.Row():
                gcode_source = gr.Radio(
                    choices=[GCODE_SOURCE_SHAPE1, GCODE_SOURCE_UPLOAD],
                    value=GCODE_SOURCE_SHAPE1,
                    label="G-Code source",
                )
                gcode_upload = gr.File(
                    label="Upload G-Code",
                    file_types=[".txt", ".gcode", ".nc"],
                    interactive=False,
                )
            render_button = gr.Button("Render Tool Path", variant="primary")
            with gr.Row():
                travel_opacity_slider = gr.Slider(
                    label="Travel (G0) opacity",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.55,
                    step=0.05,
                )
                travel_color_picker = gr.Dropdown(
                    label="Travel (G0) color",
                    choices=[("Grey", "#969696"), ("Orange", "#ff7f0e"), ("Green", "#2ca02c"), ("Red", "#d62728"), ("Purple", "#9467bd"), ("Pink", "#e377c2"), ("Black", "#000000"), ("White", "#ffffff")],
                    value="#969696",
                    allow_custom_value=False,
                )
                print_opacity_slider = gr.Slider(
                    label="Print (G1) opacity",
                    minimum=0.0,
                    maximum=1.0,
                    value=1.0,
                    step=0.05,
                )
                print_color_picker = gr.Dropdown(
                    label="Print (G1) color",
                    choices=[("Blue", "#1f77b4"), ("Orange", "#ff7f0e"), ("Green", "#2ca02c"), ("Red", "#d62728"), ("Purple", "#9467bd"), ("Pink", "#e377c2"), ("Black", "#000000"), ("White", "#ffffff")],
                    value="#1f77b4",
                    allow_custom_value=False,
                )
            toolpath_plot = gr.Plot(label="Tool Path", elem_id="toolpath_plot")
            toolpath_status = gr.Markdown("")
            parsed_state = gr.State({})

            gcode_source.change(
                fn=toggle_gcode_source,
                inputs=[gcode_source],
                outputs=[gcode_upload],
                queue=False,
            )
            render_button.click(
                fn=render_toolpath,
                inputs=[gcode_source, gcode_upload, gcode_file_1, travel_opacity_slider, print_opacity_slider, travel_color_picker, print_color_picker],
                outputs=[toolpath_plot, toolpath_status, parsed_state],
            )
            travel_opacity_slider.release(
                fn=None,
                inputs=[travel_opacity_slider],
                outputs=[],
                js="""(opacity_val) => {
                    const container = document.getElementById("toolpath_plot");
                    if (!container) return [];
                    const plotDiv = container.querySelector(".js-plotly-plot");
                    if (!plotDiv || !plotDiv.data) return [];
                    const indices = plotDiv.data
                        .map((t, i) => t.name === "Travel (G0)" ? i : -1)
                        .filter(i => i >= 0);
                    if (indices.length > 0) Plotly.restyle(plotDiv, {opacity: opacity_val}, indices);
                    return [];
                }"""
            )
            print_opacity_slider.release(
                fn=None,
                inputs=[print_opacity_slider],
                outputs=[],
                js="""(opacity_val) => {
                    const container = document.getElementById("toolpath_plot");
                    if (!container) return [];
                    const plotDiv = container.querySelector(".js-plotly-plot");
                    if (!plotDiv || !plotDiv.data) return [];
                    const indices = plotDiv.data
                        .map((t, i) => t.name === "Print (G1)" ? i : -1)
                        .filter(i => i >= 0);
                    if (indices.length > 0) Plotly.restyle(plotDiv, {opacity: opacity_val}, indices);
                    return [];
                }"""
            )
            travel_color_picker.change(
                fn=None,
                inputs=[travel_color_picker],
                outputs=[],
                js="""(color) => {
                    const container = document.getElementById("toolpath_plot");
                    if (!container) return [];
                    const plotDiv = container.querySelector(".js-plotly-plot");
                    if (!plotDiv || !plotDiv.data) return [];
                    const indices = plotDiv.data
                        .map((t, i) => t.name === "Travel (G0)" ? i : -1)
                        .filter(i => i >= 0);
                    if (indices.length > 0) Plotly.restyle(plotDiv, {"line.color": color}, indices);
                    return [];
                }"""
            )
            print_color_picker.change(
                fn=None,
                inputs=[print_color_picker],
                outputs=[],
                js="""(color) => {
                    const container = document.getElementById("toolpath_plot");
                    if (!container) return [];
                    const plotDiv = container.querySelector(".js-plotly-plot");
                    if (!plotDiv || !plotDiv.data) return [];
                    const indices = plotDiv.data
                        .map((t, i) => t.name === "Print (G1)" ? i : -1)
                        .filter(i => i >= 0);
                    if (indices.length > 0) Plotly.restyle(plotDiv, {"line.color": color}, indices);
                    return [];
                }"""
            )

    return demo


demo = build_demo()


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