path_utils.py

from pathlib import Path

from config import (
    ALL_EXAMPLE_IMAGES,
    REAL_WORLD_VIDEO_MASK_DIRS,
    DIGITAL_TWIN_VIDEO_MASK_DIRS,
    REAL_WORLD_EXAMPLE_VIDEOS,
    DIGITAL_TWIN_EXAMPLE_VIDEOS
)

ALL_EXAMPLE_VIDEOS: set[str] = set(REAL_WORLD_EXAMPLE_VIDEOS.values()) | set(DIGITAL_TWIN_EXAMPLE_VIDEOS.values())


def _build_example_image_lookup() -> dict[str, str]:
    """
    Build a lookup table mapping filenames to their original example image paths.

    This allows matching example images even when Gradio transforms their paths
    (e.g., copying to /tmp/gradio/...).

    :return: Dictionary mapping filename -> original path
    """
    lookup = {}
    for path in ALL_EXAMPLE_IMAGES:
        filename = Path(path).name
        # If there's a collision, prefer the first one (shouldn't happen with unique names)
        if filename not in lookup:
            lookup[filename] = path
    return lookup


_EXAMPLE_IMAGE_LOOKUP = _build_example_image_lookup()


def get_original_example_path(image_path: str | Path) -> str | None:
    """
    Get the original example image path for a given image path.

    This handles the case where Gradio has copied the image to a temporary location.
    Matching is done by filename.

    :param image_path: Path to check (may be a temp path from Gradio)
    :return: Original example image path if found, None otherwise
    """
    image_path = Path(image_path)
    filename = image_path.name

    str_path = str(image_path)
    if str_path in ALL_EXAMPLE_IMAGES:
        return str_path

    try:
        resolved = str(image_path.resolve())
        for example_path in ALL_EXAMPLE_IMAGES:
            if str(Path(example_path).resolve()) == resolved:
                return example_path
    except (OSError, RuntimeError):
        pass

    return _EXAMPLE_IMAGE_LOOKUP.get(filename)


def get_mask_path_for_image(image_path: str | Path) -> Path | None:
    """
    Get the corresponding segmentation mask path for an example image.

    The mask is expected to be in the 'segmented' subdirectory with .png extension.

    This function handles:
        1. Direct paths to example images
        2. Gradio-transformed paths (matched by filename)

    :param image_path: Path to the original image
    :return: Path to the mask file if it exists, None otherwise
    """
    image_path = Path(image_path)

    original_path = get_original_example_path(image_path)
    if original_path:
        original_path = Path(original_path)
        parent = original_path.parent
        segmented_dir = parent / "segmented"
        mask_path = segmented_dir / f"{original_path.stem}.png"

        if mask_path.exists():
            return mask_path

    parent = image_path.parent
    segmented_dir = parent / "segmented"
    mask_path = segmented_dir / f"{image_path.stem}.png"

    if mask_path.exists():
        return mask_path

    return None


def is_example_image(image_path: str | Path) -> bool:
    """
    Check if the given image path is one of the predefined example images.

    :param image_path: Path to check
    :return: True if it's an example image, False otherwise
    """
    result = get_original_example_path(image_path) is not None
    return result


def is_example_video(video_path: str | Path) -> bool:
    """
    Check if the given video path is one of the predefined example videos.

    This handles the case where Gradio may have copied the video to a temporary location.
    Matching is done by both full path comparison and filename matching.

    :param video_path: Path to check (may be a temp path from Gradio)
    :return: True if it's an example video, False otherwise
    """
    video_path = Path(video_path)
    str_path = str(video_path)

    if str_path in ALL_EXAMPLE_VIDEOS:
        return True

    try:
        resolved = str(video_path.resolve())
        for example_path in ALL_EXAMPLE_VIDEOS:
            if str(Path(example_path).resolve()) == resolved:
                return True
    except (OSError, RuntimeError):
        pass

    filename = video_path.name
    for example_path in ALL_EXAMPLE_VIDEOS:
        if Path(example_path).name == filename:
            return True

    return False


def get_video_mask_dir(video_path: str | Path, video_label: str | None = None) -> Path | None:
    """
    Get the directory containing frame masks for a video.

    Only returns mask directories for predefined example videos.
    User-uploaded videos will not match, even if they have similar names.

    :param video_path: Path to the video file
    :param video_label: Optional label (Down-Left, Front, Right) if known
    :return: Path to mask directory if it exists, None otherwise
    """
    video_path = Path(video_path)
    if not is_example_video(video_path):
        return None

    video_name = video_path.stem
    parent_name = video_path.parent.name

    if "real-world" in parent_name or "real-world" in video_name:
        mask_dirs = REAL_WORLD_VIDEO_MASK_DIRS
    elif "digital-twin" in parent_name or "digital-twin" in video_name:
        mask_dirs = DIGITAL_TWIN_VIDEO_MASK_DIRS
    else:
        return None

    if video_label and video_label in mask_dirs:
        mask_dir = mask_dirs[video_label]
        if mask_dir.exists():
            return mask_dir

    for label, mask_dir in mask_dirs.items():
        if label.lower().replace("-", "") in video_name.lower().replace("-", ""):
            if mask_dir.exists():
                return mask_dir

    return None


def get_frame_mask_path(mask_dir: Path, frame_idx: int) -> Path | None:
    """
    Get the mask file path for a specific video frame.

    The mask files follow the naming pattern: frame_{frame_idx:06d}_{timestamp}.png

    :param mask_dir: Directory containing the mask files
    :param frame_idx: Frame index (0-based)
    :return: Path to the mask file if found, None otherwise
    """
    if not mask_dir.exists():
        return None

    frame_prefix = f"frame_{frame_idx:06d}_"

    for mask_file in mask_dir.glob(f"{frame_prefix}*.png"):
        return mask_file

    return None


def build_frame_mask_index(mask_dir: Path) -> dict[int, Path]:
    """
    Build an index mapping frame numbers to their mask file paths.

    Masks are sorted by their original frame number and re-indexed
    starting from 0 to match cut video frame indices.

    :param mask_dir: Directory containing mask files
    :return: Dictionary mapping frame_idx -> mask_path
    """
    if not mask_dir.exists():
        return {}

    masks_with_original_idx = []
    for mask_file in mask_dir.glob("frame_*.png"):
        parts = mask_file.stem.split("_")
        if len(parts) >= 2:
            try:
                original_frame_idx = int(parts[1])
                masks_with_original_idx.append((original_frame_idx, mask_file))
            except ValueError:
                continue

    masks_with_original_idx.sort(key=lambda x: x[0])

    return {i: mask_file for i, (_, mask_file) in enumerate(masks_with_original_idx)}