#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   path_utils.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Path resolution and validation utilities for inference.
"""

import string
from pathlib import Path

# Supported image file extensions
IMAGE_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".bmp",
    ".gif",
    ".tiff",
    ".tif",
    ".webp",
    ".svg",
    ".ico",
    ".ppm",
    ".pgm",
    ".pbm",
    ".pnm",
    ".jfif",
    ".jp2",
    ".jpx",
    ".j2k",
    ".heic",
    ".heif",
    ".avif",
}

# Supported video file extensions
VIDEO_EXTENSIONS = {
    ".mp4",
    ".avi",
    ".mov",
    ".mkv",
    ".webm",
    ".flv",
    ".wmv",
    ".m4v",
    ".mpeg",
    ".mpg",
    ".mpe",
    ".3gp",
    ".3g2",
    ".ogv",
    ".ogg",
    ".vob",
    ".ts",
    ".mts",
    ".m2ts",
    ".asf",
    ".f4v",
    ".divx",
    ".xvid",
    ".qt",
}


def is_video_file(path: str | Path) -> bool:
    """Check if a file is a video based on its extension.

    Args:
        path: Path to the file (string or Path object).

    Returns:
        True if the file has a video extension, False otherwise.

    Example:
        ```python
        if is_video_file("video.mp4"):
            print("This is a video file")
        ```

    """
    file_path = Path(path)
    return file_path.suffix.lower() in VIDEO_EXTENSIONS


def is_image_file(path: str | Path) -> bool:
    """Check if a file is an image based on its extension.

    Args:
        path: Path to the file (string or Path object).

    Returns:
        True if the file has an image extension, False otherwise.

    Example:
        ```python
        if is_image_file("photo.jpg"):
            print("This is an image file")
        ```

    """
    file_path = Path(path)
    return file_path.suffix.lower() in IMAGE_EXTENSIONS


def is_base64_image(image_str: str) -> bool:
    """Check if a string appears to be a base64-encoded image.

    Args:
        image_str: String to check

    Returns:
        True if the string appears to be base64-encoded image data

    """
    # Check for data URI format
    if image_str.startswith("data:image"):
        return True

    # Check if it's a long base64-like string (rough heuristic)
    # Base64 images are typically much longer than file paths
    if len(image_str) < 100:  # File paths are usually short
        return False

    # Check if it contains only valid base64 characters
    valid_chars = set(string.ascii_letters + string.digits + "+/=")
    return all(c in valid_chars for c in image_str)


def is_url(source: str) -> bool:
    """Check if a string is a URL (http:// or https://).

    Args:
        source: String to check

    Returns:
        True if the string is a URL, False otherwise

    Example:
        ```python
        if is_url("https://example.com/image.jpg"):
            print("This is a URL")
        ```

    """
    return source.startswith(("http://", "https://"))


def is_data_uri(source: str) -> bool:
    """Check if a string is a data URI.

    Args:
        source: String to check

    Returns:
        True if the string is a data URI, False otherwise

    Example:
        ```python
        if is_data_uri("data:image/jpeg;base64,..."):
            print("This is a data URI")
        ```

    """
    return source.startswith("data:")


def is_video_url(url: str) -> bool:
    """Check if a URL points to a video based on its file extension.

    Args:
        url: URL string to check

    Returns:
        True if the URL has a video file extension, False otherwise

    Example:
        ```python
        if is_video_url("https://example.com/video.mp4"):
            print("This is a video URL")
        ```

    """
    # Extract the path component and check extension
    # Handle URLs with query params by splitting on '?'
    url_path = url.split("?")[0]
    return Path(url_path).suffix.lower() in VIDEO_EXTENSIONS


def is_video_data_uri(uri: str) -> bool:
    """Check if a data URI is for video content.

    Args:
        uri: Data URI string to check

    Returns:
        True if the data URI is for video content, False otherwise

    Example:
        ```python
        if is_video_data_uri("data:video/mp4;base64,..."):
            print("This is a video data URI")
        ```

    """
    return uri.startswith("data:video/")


def resolve_file_path(path: str | Path) -> Path:
    """Resolve and validate a local file path.

    Handles only local file paths (no URLs or data URIs).

    Args:
        path: File path as string or Path object

    Returns:
        Resolved absolute Path object

    Raises:
        FileNotFoundError: If the file does not exist
        ValueError: If the path is invalid

    Example:
        ```python
        path = resolve_file_path("./image.jpg")
        ```

    """
    # Convert to Path object if string
    path_obj = Path(path)

    # Expand user path (~) and environment variables
    path_obj = path_obj.expanduser()

    # Resolve to absolute path (also resolves symlinks and '..')
    try:
        resolved_path = path_obj.resolve(strict=False)
    except (OSError, RuntimeError) as e:
        raise ValueError(f"Invalid path '{path}': {e}") from e

    # Check if file exists
    if not resolved_path.exists():
        raise FileNotFoundError(
            f"File not found: '{path}'\n"
            f"Resolved to: '{resolved_path}'\n"
            f"Please check that the file exists and the path is correct."
        )

    # Check if it's a file (not a directory)
    if not resolved_path.is_file():
        raise ValueError(
            f"Path is not a file: '{path}'\n"
            f"Resolved to: '{resolved_path}'\n"
            f"Please provide a path to an image file."
        )

    return resolved_path


def resolve_directory_to_images(
    directory: str | Path,
    recursive: bool = False,
) -> list[Path]:
    """Resolve a directory path to a list of image file paths.

    Scans a directory for image files and returns their absolute paths.
    Supports common image formats (JPEG, PNG, BMP, GIF, TIFF, WebP).

    Args:
        directory: Directory path as string or Path object.
        recursive: If True, recursively search subdirectories. Defaults to False.

    Returns:
        List of resolved absolute Path objects for image files found.

    Raises:
        FileNotFoundError: If the directory does not exist.
        ValueError: If the path is not a directory or no images found.

    Example:
        ```python
        # Get all images in a folder
        images = resolve_directory_to_images("./my_images/")

        # Recursively search subdirectories
        images = resolve_directory_to_images("./dataset/", recursive=True)
        ```

    """
    if not directory:
        raise ValueError("Directory path cannot be empty")

    # Convert to Path object
    dir_path = Path(directory)

    # Expand user path and resolve
    dir_path = dir_path.expanduser()
    try:
        resolved_dir = dir_path.resolve(strict=False)
    except (OSError, RuntimeError) as e:
        raise ValueError(f"Invalid directory path '{directory}': {e}") from e

    # Check if directory exists
    if not resolved_dir.exists():
        raise FileNotFoundError(
            f"Directory not found: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Please check that the directory exists and the path is correct."
        )

    # Check if it's a directory
    if not resolved_dir.is_dir():
        raise ValueError(
            f"Path is not a directory: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Please provide a path to a directory containing images."
        )

    # Find image files
    image_files: list[Path] = []

    if recursive:
        # Recursively search all subdirectories
        for ext in IMAGE_EXTENSIONS:
            image_files.extend(resolved_dir.rglob(f"*{ext}"))
            image_files.extend(resolved_dir.rglob(f"*{ext.upper()}"))
    else:
        # Only search immediate directory
        for ext in IMAGE_EXTENSIONS:
            image_files.extend(resolved_dir.glob(f"*{ext}"))
            image_files.extend(resolved_dir.glob(f"*{ext.upper()}"))

    # Remove duplicates and sort
    image_files = sorted(set(image_files))

    if not image_files:
        raise ValueError(
            f"No image files found in directory: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Supported formats: {', '.join(sorted(IMAGE_EXTENSIONS))}\n"
            f"Recursive search: {recursive}"
        )

    return image_files


def resolve_directory_to_videos(
    directory: str | Path,
    recursive: bool = False,
) -> list[Path]:
    """Resolve a directory path to a list of video file paths.

    Scans a directory for video files and returns their absolute paths.
    Supports common video formats (MP4, AVI, MOV, MKV, WebM, etc.).

    Args:
        directory: Directory path as string or Path object.
        recursive: If True, recursively search subdirectories. Defaults to False.

    Returns:
        List of resolved absolute Path objects for video files found.

    Raises:
        FileNotFoundError: If the directory does not exist.
        ValueError: If the path is not a directory or no videos found.

    Example:
        ```python
        # Get all videos in a folder
        videos = resolve_directory_to_videos("./my_videos/")

        # Recursively search subdirectories
        videos = resolve_directory_to_videos("./dataset/", recursive=True)
        ```

    """
    if not directory:
        raise ValueError("Directory path cannot be empty")

    # Convert to Path object
    dir_path = Path(directory)

    # Expand user path and resolve
    dir_path = dir_path.expanduser()
    try:
        resolved_dir = dir_path.resolve(strict=False)
    except (OSError, RuntimeError) as e:
        raise ValueError(f"Invalid directory path '{directory}': {e}") from e

    # Check if directory exists
    if not resolved_dir.exists():
        raise FileNotFoundError(
            f"Directory not found: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Please check that the directory exists and the path is correct."
        )

    # Check if it's a directory
    if not resolved_dir.is_dir():
        raise ValueError(
            f"Path is not a directory: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Please provide a path to a directory containing videos."
        )

    # Find video files
    video_files: list[Path] = []

    if recursive:
        # Recursively search all subdirectories
        for ext in VIDEO_EXTENSIONS:
            video_files.extend(resolved_dir.rglob(f"*{ext}"))
            video_files.extend(resolved_dir.rglob(f"*{ext.upper()}"))
    else:
        # Only search immediate directory
        for ext in VIDEO_EXTENSIONS:
            video_files.extend(resolved_dir.glob(f"*{ext}"))
            video_files.extend(resolved_dir.glob(f"*{ext.upper()}"))

    # Remove duplicates and sort
    video_files = sorted(set(video_files))

    if not video_files:
        raise ValueError(
            f"No video files found in directory: '{directory}'\n"
            f"Resolved to: '{resolved_dir}'\n"
            f"Supported formats: {', '.join(sorted(VIDEO_EXTENSIONS))}\n"
            f"Recursive search: {recursive}"
        )

    return video_files
