processing.py

import csv
import logging
import shutil
import tempfile
import uuid
import zipfile
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
from PIL import Image

from config import SEGMENTATION_COLOR
from managers import managed_temp_dir, video_capture, video_writer
from mask_splitter.nn.infer import MaskSplitterInference
from mask_splitter.yolo_model import YoloSegmentation
from path_utils import get_mask_path_for_image, is_example_image, get_video_mask_dir, build_frame_mask_index
from validation import validate_image, validate_video
from vision_utils import (
    AnnotationState, colorize_masks, create_overlay, add_text_overlay, ensure_binary_mask,
    geometric_split_mask, pil_to_rgb_array, resize_mask_to_image, create_segmentation_overlay,
    create_annotation_visualization
)

logger = logging.getLogger(__name__)


def load_preexisting_mask(image_path: str | Path) -> np.ndarray | None:
    """
    Load a pre-existing segmentation mask for an example image.

    :param image_path: Path to the image file
    :return: Binary mask array (H, W) with values 0 or 255, or None if not found
    """
    mask_path = get_mask_path_for_image(image_path)
    if mask_path is None:
        return None

    try:
        mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
        if mask is not None:
            mask = ensure_binary_mask(mask)
            return mask
    except Exception as e:
        logger.warning("Failed to load pre-existing mask from %s: %s", mask_path, e)

    return None


def load_preexisting_mask_pil(image_path: str | Path) -> Image.Image | None:
    """
    Load a pre-existing segmentation mask for an example image as PIL Image.

    :param image_path: Path to the image file
    :return: PIL Image in grayscale mode, or None if not found
    """
    mask_path = get_mask_path_for_image(image_path)
    if mask_path is None:
        return None

    try:
        mask = Image.open(mask_path).convert("L")
        return mask
    except Exception as e:
        logger.warning("Failed to load pre-existing mask from %s: %s", mask_path, e)

    return None


def run_image_inference(
        image: Image.Image | None,
        mask_image: Image.Image | None,
        mask_splitter_model_name: str,
        use_yolo: bool,
        mask_splitter_engines: dict[str, MaskSplitterInference],
        yolo_engines: dict[str, YoloSegmentation],
        image_path: str | None = None,
        mask_source: str | None = None,
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None, str]:
    """
    Run mask splitting on a single image.

    Returns:
        - Original image with YOLO/input mask overlay
        - Visualization with front/back overlay
        - Front mask
        - Back mask
        - Status message

    :param image: Input PIL image
    :param mask_image: Optional segmentation mask (PIL image)
    :param mask_splitter_model_name: Name of the mask splitter model to use
    :param use_yolo: Whether to use YOLO for segmentation
    :param mask_splitter_engines: Dictionary of loaded mask splitter engines
    :param yolo_engines: Dictionary of loaded YOLO engines
    :param image_path: Optional path to original image (for loading pre-existing masks)
    :param mask_source: Optional pre-determined mask source ("preexisting", "yolo", "uploaded")
    :return: Tuple of (input_vis, output_vis, front_mask, back_mask, status_message)
    """
    is_valid, error_msg = validate_image(image)
    if not is_valid:
        if image is None:
            return None, None, None, None, "Please upload an image."
        raise gr.Error(error_msg)

    if mask_splitter_model_name not in mask_splitter_engines:
        available = list(mask_splitter_engines.keys())
        raise gr.Error(f"Invalid model: {mask_splitter_model_name}. Available: {available}")

    image_rgb = pil_to_rgb_array(image)
    original_h, original_w = image_rgb.shape[:2]
    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    mask, determined_mask_source = _get_segmentation_mask(
        image_bgr, mask_image, mask_splitter_model_name, use_yolo, yolo_engines, image_path, mask_source
    )

    if mask is None:
        return None, None, None, None, determined_mask_source

    if mask.sum() == 0:
        return None, None, None, None, "No object detected in mask. Try a different image."

    engine = mask_splitter_engines[mask_splitter_model_name]
    front_mask, back_mask = engine.infer(image_bgr, mask)

    front_mask = cv2.resize(front_mask, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
    back_mask = cv2.resize(back_mask, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
    mask = resize_mask_to_image(mask, image_bgr)

    # Create visualizations (in RGB for Gradio)
    # 1. Input visualization (image + segmentation mask)
    input_vis = create_segmentation_overlay(image_rgb, mask, alpha=0.6)
    input_vis = add_text_overlay(input_vis, f"Input: {determined_mask_source}", position="bottom-right")

    # 2. Output visualization (image + front/back masks)
    output_vis = create_overlay(image_rgb, front_mask, back_mask, alpha=0.5)
    output_vis = add_text_overlay(output_vis, "Front (Red) | Back (Blue)", position="bottom-right")

    front_pixels = (front_mask > 0).sum()
    back_pixels = (back_mask > 0).sum()
    total_pixels = front_pixels + back_pixels
    front_pct = (front_pixels / total_pixels * 100) if total_pixels > 0 else 0
    back_pct = (back_pixels / total_pixels * 100) if total_pixels > 0 else 0

    status = (
        f"Inference complete | Model: {mask_splitter_model_name}\n"
        f"Front: {front_pixels:,} px ({front_pct:.1f}%) | Back: {back_pixels:,} px ({back_pct:.1f}%)"
    )

    return input_vis, output_vis, front_mask, back_mask, status


def run_video_inference(
        video_path: str | None,
        mask_splitter_model_name: str,
        use_yolo: bool,
        export_csv: bool,
        mask_splitter_engines: dict[str, MaskSplitterInference],
        yolo_engines: dict[str, YoloSegmentation],
        progress: gr.Progress = gr.Progress(),
        video_label: str | None = None,
        export_frames_zip: bool = False,
) -> tuple[str | None, str | None, str | None]:
    """
    Run mask splitting on a video file.

    Returns:
        - Path to output video
        - Path to CSV file (if export_csv is True)
        - Path to ZIP file with frames and masks (if export_frames_zip is True)


    :param video_path: Path to the input video
    :param mask_splitter_model_name: Name of the mask splitter model to use
    :param use_yolo: Whether to use YOLO for segmentation (fallback if no pre-existing masks)
    :param export_csv: Whether to export frame annotations to CSV
    :param mask_splitter_engines: Dictionary of loaded mask splitter engines
    :param yolo_engines: Dictionary of loaded YOLO engines
    :param progress: Gradio progress tracker
    :param video_label: Optional video label (Down-Left, Front, Right) for mask lookup
    :param export_frames_zip: Whether to export frames and masks as a ZIP file
    :return: Tuple of (output_video_path, csv_path, frames_zip_path)
    """
    if video_path is None:
        return None, None, None

    if mask_splitter_model_name not in mask_splitter_engines:
        raise gr.Error(f"Invalid mask-splitter model: {mask_splitter_model_name}")

    if use_yolo and (len(yolo_engines) == 0 or mask_splitter_model_name not in yolo_engines):
        raise gr.Error("YOLO model not available. Cannot process video without segmentation.")

    is_valid, error_msg, video_info = validate_video(video_path)
    if not is_valid:
        raise gr.Error(error_msg)

    mask_dir = get_video_mask_dir(video_path, video_label)
    use_preexisting_masks = mask_dir is not None

    if use_preexisting_masks:
        frame_mask_index = build_frame_mask_index(mask_dir)
        logger.info("Using pre-existing masks (%d frames indexed)", len(frame_mask_index))
    else:
        frame_mask_index = {}
        if not use_yolo:
            logger.info("No pre-existing masks found. Enabling YOLO segmentation automatically.")
            use_yolo = True
        if len(yolo_engines) == 0 or mask_splitter_model_name not in yolo_engines:
            raise gr.Error("YOLO model not available and no pre-existing masks found. Cannot process video.")

    engine = mask_splitter_engines[mask_splitter_model_name]
    yolo_model = yolo_engines.get(mask_splitter_model_name) if use_yolo else None
    fps = video_info.fps
    total_frames = video_info.total_frames
    width = video_info.width
    height = video_info.height

    mask_source = "Pre-existing Masks" if use_preexisting_masks else "YOLO Segmentation"
    logger.info(
        "Processing video: %d frames at %.1f FPS (%dx%d), Mask source: %s",
        total_frames, fps, width, height, mask_source
    )

    with managed_temp_dir() as temp_dir:
        output_path = str(Path(temp_dir) / "output.mp4")
        csv_path = str(Path(temp_dir) / "annotations.csv") if export_csv else None
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        output_width = width * 2

        frames_dir = None
        segmented_dir = None
        front_dir = None
        back_dir = None
        if export_frames_zip:
            frames_dir = Path(temp_dir) / "frames"
            segmented_dir = Path(temp_dir) / "segmented"
            front_dir = Path(temp_dir) / "front"
            back_dir = Path(temp_dir) / "back"
            frames_dir.mkdir(exist_ok=True)
            segmented_dir.mkdir(exist_ok=True)
            front_dir.mkdir(exist_ok=True)
            back_dir.mkdir(exist_ok=True)

        with video_capture(video_path) as capture, \
                video_writer(output_path, fourcc, fps, (output_width, height)) as writer:

            annotations = []
            frame_idx = 0
            progress(0, desc="Processing video...")

            while True:
                ret, frame = capture.read()
                if not ret:
                    break

                preexisting_mask = None
                if use_preexisting_masks and frame_idx in frame_mask_index:
                    mask_path = frame_mask_index[frame_idx]
                    preexisting_mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
                    if preexisting_mask is not None:
                        preexisting_mask = ensure_binary_mask(preexisting_mask)
                        if preexisting_mask.shape[:2] != (height, width):
                            preexisting_mask = cv2.resize(
                                preexisting_mask, (width, height), interpolation=cv2.INTER_NEAREST
                            )
                front_mask, back_mask, mask = _process_video_frame(
                    frame, engine, yolo_model, width, height, preexisting_mask
                )
                combined = _create_video_frame_visualization(frame, front_mask, back_mask, mask, mask_source)
                writer.write(combined)

                if export_frames_zip:
                    frame_filename = f"frame_{frame_idx:06d}.png"
                    cv2.imwrite(str(frames_dir / frame_filename), frame)
                    cv2.imwrite(str(segmented_dir / frame_filename), mask)
                    cv2.imwrite(str(front_dir / frame_filename), front_mask)
                    cv2.imwrite(str(back_dir / frame_filename), back_mask)

                if export_csv:
                    annotations.append(_create_frame_annotation(frame_idx, fps, front_mask, back_mask, mask))

                frame_idx += 1
                if frame_idx % 10 == 0:
                    progress(frame_idx / total_frames, desc=f"Processing frame {frame_idx}/{total_frames}")

        if export_csv and annotations:
            _write_annotations_csv(csv_path, annotations)

        frames_zip_path = None
        if export_frames_zip:
            progress(0.95, desc="Creating ZIP file...")
            frames_zip_path = _create_frames_zip(temp_dir, video_path)

        progress(1.0, desc="Done!")
        logger.info("Processed %d frames successfully", frame_idx)

        return _save_video_outputs_persistent_temp_location(output_path, csv_path, frames_zip_path)


def _get_segmentation_mask(
        image_bgr: np.ndarray,
        mask_image: Image.Image | None,
        model_name: str | None,
        use_yolo: bool,
        yolo_engines: dict,
        image_path: str | None = None,
        mask_source: str | None = None
) -> tuple[np.ndarray | None, str]:
    """
    Get segmentation mask from pre-existing file, YOLO, or uploaded mask.

    Priority:
    1. Pre-existing mask (for example images)
    2. Uploaded mask
    3. YOLO segmentation

    :return: Tuple of (mask, source_description) or (None, error_message)
    """
    source_display_map = {
        "preexisting": "Pre-existing Mask",
        "yolo": "YOLO Segmentation",
        "uploaded": "Uploaded Mask",
    }
    # 1. Try to load pre-existing mask for example images (only if no mask_image provided)
    if mask_image is None and image_path and is_example_image(image_path):
        preexisting_mask = load_preexisting_mask(image_path)
        if preexisting_mask is not None:
            img_h, img_w = image_bgr.shape[:2]
            if preexisting_mask.shape[:2] != (img_h, img_w):
                preexisting_mask = cv2.resize(preexisting_mask, (img_w, img_h), interpolation=cv2.INTER_NEAREST)
            return preexisting_mask, "Pre-existing Mask"

    # 2. Use uploaded mask if provided
    if mask_image is not None:
        display_source = source_display_map.get(mask_source, "Uploaded Mask")
        mask = np.array(mask_image.convert("L"))
        mask = ensure_binary_mask(mask)
        return mask, display_source

    # 3. Fall back to YOLO segmentation
    if use_yolo and model_name and len(yolo_engines) > 0 and model_name in yolo_engines:
        yolo_model = yolo_engines[model_name]
        _, mask = yolo_model.segment_image(image_bgr)
        return mask, "YOLO Segmentation"

    logger.info("Invalid Mask or YOLO segmentation...")
    return None, "Please provide a mask or enable YOLO segmentation."


def generate_yolo_mask(
        image: Image.Image | None,
        model_name: str,
        use_yolo: bool,
        yolo_engines: dict[str, YoloSegmentation],
        image_path: str | None = None
) -> Image.Image | None:
    """
    Generate a YOLO segmentation mask for the given image.

    For example images, attempts to load pre-existing masks first.
    Falls back to YOLO segmentation if no pre-existing mask is available.

    :return: PIL Image (grayscale mask) or None
    """
    if image is None:
        return None

    if image_path and is_example_image(image_path):
        preexisting_mask = load_preexisting_mask_pil(image_path)
        if preexisting_mask is not None:
            return preexisting_mask

    if not use_yolo:
        return None

    if model_name not in yolo_engines:
        return None

    try:
        image_rgb = pil_to_rgb_array(image)
        image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

        yolo_model = yolo_engines[model_name]
        _, mask = yolo_model.segment_image(image_bgr)

        if mask is not None and mask.sum() > 0:
            return Image.fromarray(mask)
        return None
    except Exception as e:
        logger.warning("Failed to generate YOLO mask: %s", e)
        return None


def _process_video_frame(
        frame: np.ndarray,
        engine,
        yolo_model,
        width: int,
        height: int,
        preexisting_mask: np.ndarray | None = None
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Process a single video frame through YOLO and mask-splitter."""
    if preexisting_mask is not None:
        mask = preexisting_mask
    elif yolo_model is not None:
        _, mask = yolo_model.segment_image(frame)
    else:
        mask = np.zeros((height, width), dtype=np.uint8)

    if mask.sum() > 0:
        front_mask, back_mask = engine.infer(frame, mask)
        front_mask = cv2.resize(front_mask, (width, height), interpolation=cv2.INTER_NEAREST)
        back_mask = cv2.resize(back_mask, (width, height), interpolation=cv2.INTER_NEAREST)
    else:
        front_mask = np.zeros((height, width), dtype=np.uint8)
        back_mask = np.zeros((height, width), dtype=np.uint8)

    if mask.shape[:2] != (height, width):
        mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)

    return front_mask, back_mask, mask


def _create_video_frame_visualization(
        frame: np.ndarray,
        front_mask: np.ndarray,
        back_mask: np.ndarray,
        mask: np.ndarray,
        mask_source: str = "YOLO Segmentation"
) -> np.ndarray:
    """Create side-by-side visualization for video frame."""
    # Left panel: input + segmentation overlay
    left_panel = frame.copy()
    seg_color = np.zeros_like(frame)
    seg_color[mask > 0] = SEGMENTATION_COLOR
    left_panel = cv2.addWeighted(left_panel, 0.6, seg_color, 0.4, 0)

    # Right panel: input + front/back overlay
    right_panel = frame.copy()
    pred_color = colorize_masks(front_mask, back_mask, use_bgr=True)
    pred_mask_area = (front_mask == 255) | (back_mask == 255)
    if pred_mask_area.any():
        right_panel[pred_mask_area] = cv2.addWeighted(
            frame, 0.5, pred_color, 0.5, 0
        )[pred_mask_area]

    left_panel = add_text_overlay(left_panel, mask_source, position="top-right")
    right_panel = add_text_overlay(right_panel, "Front (Red) | Back (Blue)")

    return np.hstack((left_panel, right_panel))


def _create_frame_annotation(
        frame_idx: int,
        fps: float,
        front_mask: np.ndarray,
        back_mask: np.ndarray,
        mask: np.ndarray,
) -> dict:
    """Create annotation dictionary for a single frame."""
    timestamp = frame_idx / fps if fps > 0 else 0
    front_pixels = (front_mask > 0).sum()
    back_pixels = (back_mask > 0).sum()
    return {
        "frame": frame_idx,
        "timestamp_sec": round(timestamp, 4),
        "front_pixels": int(front_pixels),
        "back_pixels": int(back_pixels),
        "total_mask_pixels": int(front_pixels + back_pixels),
        "has_detection": int(mask.sum() > 0),
    }


def _write_annotations_csv(csv_path: str, annotations: list[dict]) -> None:
    """Write annotations to CSV file."""
    fieldnames = ["frame", "timestamp_sec", "front_pixels", "back_pixels", "total_mask_pixels", "has_detection"]
    with open(csv_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(annotations)


def _save_video_outputs_persistent_temp_location(
        output_path: str, csv_path: str | None, frames_zip_path: str | None = None
) -> tuple[str, str | None, str | None]:
    """Copy video outputs to persistent temp location."""
    persistent_temp = tempfile.gettempdir()
    final_id = uuid.uuid4()
    final_output = shutil.copy(output_path, Path(persistent_temp) / f"output_{final_id}.mp4")
    final_csv = None
    if csv_path:
        final_csv = shutil.copy(csv_path, Path(persistent_temp) / f"annotations_{final_id}.csv")
    final_zip = None
    if frames_zip_path:
        final_zip = shutil.copy(frames_zip_path, Path(persistent_temp) / f"video_data_{final_id}.zip")
    return str(final_output), str(final_csv) if final_csv else None, str(final_zip) if final_zip else None


def _create_frames_zip(temp_dir: str, video_path: str) -> str:
    """
    Create a ZIP file containing frames and masks directories.

    Structure:
        video_data/
        ├── frames/      # Original video frames
        ├── segmented/   # Segmentation masks
        ├── front/       # Front mask splits
        └── back/        # Back mask splits

    :param temp_dir: Temporary directory containing the subdirectories
    :param video_path: Original video path (used for naming)
    :return: Path to the created ZIP file
    """
    temp_dir = Path(temp_dir)
    video_name = Path(video_path).stem

    zip_path = temp_dir / f"{video_name}_data.zip"

    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for subdir in ["frames", "segmented", "front", "back"]:
            subdir_path = temp_dir / subdir
            if subdir_path.exists():
                for file in sorted(subdir_path.glob("*.png")):
                    archive_name = f"{video_name}_data/{subdir}/{file.name}"
                    zipf.write(file, archive_name)

    logger.info("Created frames ZIP: %s", zip_path)
    return str(zip_path)


def load_annotation_image(
        image: Image.Image | None,
        mask_image: Image.Image | None,
        use_yolo: bool,
        yolo_model_name: str | None,
        state: AnnotationState,
        yolo_engines: dict[str, YoloSegmentation],
        source_filename: str | None = None,
        image_path: str | None = None
) -> tuple[np.ndarray | None, np.ndarray | None, str, AnnotationState]:
    """Load image and mask for annotation."""
    if image is None:
        return None, None, "Please upload an image.", state

    image_rgb = pil_to_rgb_array(image)
    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    mask, mask_source = _get_segmentation_mask(
        image_bgr, mask_image, yolo_model_name, use_yolo, yolo_engines, image_path
    )
    if mask is None:
        return None, None, mask_source, state

    mask = resize_mask_to_image(mask, image_rgb)

    if mask.sum() == 0:
        return None, mask, "No object detected in mask.", state

    state.original_image = image_rgb
    state.mask = mask
    state.source_filename = source_filename
    state.reset_split()

    vis = create_annotation_visualization(image_rgb, mask, show_instructions=True, is_split_preview=False)
    status = f"Image loaded ({mask_source}). Click on the FRONT portion of the object."
    return vis, mask, status, state


def annotate_click(
        image: np.ndarray, state: AnnotationState, evt: gr.SelectData
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, str, AnnotationState]:
    """
    Handle click on annotation canvas.

    This implements the click-to-split interaction similar to CarMaskSplitter's
    get_user_click() and geometric_split_mask() workflow. In Gradio, we get
    immediate visual feedback without the confirmation loop (K/Enter to confirm,
    R to redo) that the mask-splitter repo tool provides.
    """
    if not state.is_loaded:
        return None, None, None, "Please load an image first.", state

    x, y = evt.index[0], evt.index[1]

    if y >= state.mask.shape[0] or x >= state.mask.shape[1]:
        return image, state.front_mask, state.back_mask, "Click inside the image bounds.", state

    if state.mask[y, x] == 0:
        status_hint = " (adjusted to nearest mask point)"
    else:
        status_hint = ""

    front_mask, back_mask = geometric_split_mask(state.mask, (x, y))
    state.front_mask = front_mask
    state.back_mask = back_mask
    state.last_click_point = (x, y)

    vis = create_annotation_visualization(
        state.original_image,
        state.mask,
        front_mask=front_mask,
        back_mask=back_mask,
        click_point=(x, y),
        show_instructions=True,
        is_split_preview=True,
    )

    front_pixels = (front_mask > 0).sum()
    back_pixels = (back_mask > 0).sum()
    total_pixels = front_pixels + back_pixels
    front_pct = (front_pixels / total_pixels * 100) if total_pixels > 0 else 0
    back_pct = (back_pixels / total_pixels * 100) if total_pixels > 0 else 0
    status = (
        f"Split complete{status_hint}! "
        f"Front: {front_pixels:,} px ({front_pct:.1f}%) | Back: {back_pixels:,} px ({back_pct:.1f}%)\n"
        f"Click again to adjust, or Reset to start over."
    )

    return vis, front_mask, back_mask, status, state


def reset_annotation(
        state: AnnotationState
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, str, AnnotationState]:
    """
    Reset annotation state.

    This is equivalent to pressing 'R' (redo) in CarMaskSplitter's confirmation
    loop, returning to the initial state where the user can click again.
    """
    state.reset_split()

    if not state.is_loaded:
        return None, None, None, "Please load an image first.", state

    vis = create_annotation_visualization(
        state.original_image,
        state.mask,
        show_instructions=True,
        is_split_preview=False,
    )
    return vis, None, None, "Reset. Click on the FRONT portion of the object.", state


def download_split_masks(state: AnnotationState) -> str | None:
    """
    Create a downloadable ZIP file containing the front and back masks.

    The masks are saved with the original frame name as prefix:
    - {frame_name}_front.png
    - {frame_name}_back.png

    Returns:
        Path to the ZIP file, or None if no split has been performed.
    """
    if not state.has_split:
        gr.Warning("No split masks available. Please perform a split first.")
        return None

    base_name = state.get_base_filename()

    persistent_temp = tempfile.gettempdir()
    unique_id = uuid.uuid4()
    zip_filename = f"{base_name}_masks_{unique_id}.zip"
    zip_path = Path(persistent_temp) / zip_filename

    front_filename = f"{base_name}_front.png"
    back_filename = f"{base_name}_back.png"

    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        front_pil = Image.fromarray(state.front_mask)
        front_temp = Path(persistent_temp) / front_filename
        front_pil.save(front_temp)
        zipf.write(front_temp, front_filename)
        front_temp.unlink()

        back_pil = Image.fromarray(state.back_mask)
        back_temp = Path(persistent_temp) / back_filename
        back_pil.save(back_temp)
        zipf.write(back_temp, back_filename)
        back_temp.unlink()

    return str(zip_path)


def get_source_filename_from_upload(image_data) -> str | None:
    """
    Extract the original filename from a Gradio image upload.

    Gradio's Image component with type="pil" doesn't preserve filename,
    but we can extract it from the component's internal data if available.
    """
    if image_data is None:
        return None

    # If it's a file path (from gallery selection), extract the filename
    if isinstance(image_data, str):
        return Path(image_data).name

    # For PIL images, try to get filename from the info dict
    if hasattr(image_data, 'filename'):
        return Path(image_data.filename).name

    return None