#!/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 base64
import string
import tempfile
from collections.abc import Sequence
from io import BytesIO
from pathlib import Path

from PIL import Image

# 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 save_base64_image(base64_data: str) -> str:
    """Save a base64-encoded image to a temporary file.

    Args:
        base64_data: Base64-encoded image data (with or without data URI prefix)

    Returns:
        Path to the temporary file containing the decoded image

    Raises:
        ValueError: If the base64 data is invalid or cannot be decoded

    """
    # Extract base64 data from data URI if present
    if base64_data.startswith("data:image"):
        if "," not in base64_data:
            raise ValueError("Invalid data URI format")
        base64_data = base64_data.split(",", 1)[1]

    try:
        # Decode base64 data
        image_bytes = base64.b64decode(base64_data)

        # Load image with PIL to validate and get format
        image = Image.open(BytesIO(image_bytes))

        # Create temporary file
        with tempfile.NamedTemporaryFile(
            delete=False, suffix=f".{image.format.lower() if image.format else 'png'}"
        ) as temp_file:
            image.save(temp_file.name)

        return temp_file.name

    except Exception as e:
        raise ValueError(f"Failed to decode base64 image: {e}") from e


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

    Performs automatic path resolution including:
    - Converting relative paths to absolute paths
    - Expanding user paths (~)
    - Expanding environment variables
    - Checking if the file exists
    - Handling base64-encoded images by saving them as temporary files

    Args:
        path: File path as string or Path object, or base64-encoded image string.

    Returns:
        Resolved absolute Path object.

    Raises:
        FileNotFoundError: If the file does not exist at the resolved path.
        ValueError: If the path is empty or invalid.

    Example:
        ```python
        # Relative path
        path = resolve_and_validate_path("./image.jpg")

        # User path
        path = resolve_and_validate_path("~/pictures/image.jpg")

        # Already absolute
        path = resolve_and_validate_path("/home/user/image.jpg")

        # Base64 image
        path = resolve_and_validate_path(
            "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD..."
        )
        ```

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

    path_str = str(path)

    # Check if it's a base64-encoded image
    if is_base64_image(path_str):
        try:
            temp_path = save_base64_image(path_str)
            return Path(temp_path)
        except ValueError as e:
            raise ValueError(f"Invalid base64 image data: {e}") from e

    # 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_and_validate_paths(
    paths: Sequence[str | Path],
) -> list[Path]:
    """Resolve and validate multiple file paths.

    Applies path resolution and validation to a list of paths. If any path
    is invalid or doesn't exist, provides detailed error information.

    Args:
        paths: Sequence of file paths as strings or Path objects.

    Returns:
        List of resolved absolute Path objects.

    Raises:
        FileNotFoundError: If any file does not exist.
        ValueError: If any path is empty or invalid.

    Example:
        ```python
        # Batch of paths
        paths = resolve_and_validate_paths(
            ["./image1.jpg", "~/pictures/image2.jpg", "/absolute/path/image3.jpg"]
        )

        for path in paths:
            print(f"Resolved: {path}")
        ```

    """
    if not paths:
        raise ValueError("Paths sequence cannot be empty")

    resolved_paths: list[Path] = []
    errors: list[tuple[int, str | Path, Exception]] = []

    for idx, path in enumerate(paths):
        try:
            resolved_path = resolve_and_validate_path(path)
            resolved_paths.append(resolved_path)
        except (FileNotFoundError, ValueError) as e:
            errors.append((idx, path, e))

    # If there are any errors, raise with detailed information
    if errors:
        error_messages = [
            f"  [{idx}] '{path}': {type(error).__name__}: {error}"
            for idx, path, error in errors
        ]
        raise FileNotFoundError(
            f"Failed to resolve {len(errors)} of {len(paths)} paths:\n"
            + "\n".join(error_messages)
        )

    return resolved_paths


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


def resolve_sources_to_paths(
    sources: str | Path | Sequence[str | Path],
    recursive: bool = False,
) -> list[Path]:
    """Resolve source(s) to a list of validated image paths.

    Handles three input types:
    1. Single file path → Returns [resolved_path]
    2. Directory path → Returns list of all images in directory
    3. List of paths (files/directories) → Returns list of all resolved images

    Args:
        sources: Can be:
            - Single file path (str or Path)
            - Directory path (str or Path)
            - List of file/directory paths
        recursive: If True, recursively search subdirectories when
            processing directory paths. Defaults to False.

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

    Raises:
        FileNotFoundError: If any file/directory does not exist.
        ValueError: If any path is invalid or no images found.

    Example:
        ```python
        # Single file
        paths = resolve_sources_to_paths("./image.jpg")
        # Returns: [Path("/abs/path/to/image.jpg")]

        # Directory
        paths = resolve_sources_to_paths("./my_images/")
        # Returns: [Path("/abs/.../img1.jpg"), Path("/abs/.../img2.jpg"), ...]

        # Mixed list
        paths = resolve_sources_to_paths(["./image.jpg", "./folder/", "~/pic.png"])
        # Returns: All resolved image paths

        # Recursive directory search
        paths = resolve_sources_to_paths("./dataset/", recursive=True)
        # Returns: All images including subdirectories
        ```

    """
    if not sources:
        raise ValueError("Sources cannot be empty")

    resolved_paths: list[Path] = []

    # Handle single source (string or Path)
    if isinstance(sources, (str, Path)):
        source_path = Path(sources).expanduser().resolve(strict=False)

        # Check if it exists
        if not source_path.exists():
            raise FileNotFoundError(
                f"Source not found: '{sources}'\nResolved to: '{source_path}'"
            )

        # If it's a directory, get all images
        if source_path.is_dir():
            resolved_paths = resolve_directory_to_images(
                source_path, recursive=recursive
            )
        # If it's a file, validate and add
        else:
            resolved_paths = [resolve_and_validate_path(sources)]

    # Handle sequence of sources
    else:
        errors: list[tuple[int, str | Path, Exception]] = []

        for idx, source in enumerate(sources):
            try:
                source_path = Path(source).expanduser().resolve(strict=False)

                if not source_path.exists():
                    raise FileNotFoundError(f"Source not found: '{source}'")

                # If directory, expand to images
                if source_path.is_dir():
                    dir_images = resolve_directory_to_images(
                        source_path, recursive=recursive
                    )
                    resolved_paths.extend(dir_images)
                # If file, validate and add
                else:
                    resolved_paths.append(resolve_and_validate_path(source))

            except (FileNotFoundError, ValueError) as e:
                errors.append((idx, source, e))

        # If there are any errors, raise with detailed information
        if errors:
            error_messages = [
                f"  [{idx}] '{path}': {type(error).__name__}: {error}"
                for idx, path, error in errors
            ]
            raise FileNotFoundError(
                f"Failed to resolve {len(errors)} of {len(sources)} sources:\n"
                + "\n".join(error_messages)
            )

    # Remove duplicates while preserving order
    seen = set()
    unique_paths = []
    for path in resolved_paths:
        if path not in seen:
            seen.add(path)
            unique_paths.append(path)

    return unique_paths


def resolve_video_sources_to_paths(
    sources: str | Path | Sequence[str | Path],
    recursive: bool = False,
) -> list[Path]:
    """Resolve source(s) to a list of validated video paths.

    Handles three input types:
    1. Single file path → Returns [resolved_path]
    2. Directory path → Returns list of all videos in directory
    3. List of paths (files/directories) → Returns list of all resolved videos

    Args:
        sources: Can be:
            - Single file path (str or Path)
            - Directory path (str or Path)
            - List of file/directory paths
        recursive: If True, recursively search subdirectories when
            processing directory paths. Defaults to False.

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

    Raises:
        FileNotFoundError: If any file/directory does not exist.
        ValueError: If any path is invalid or no videos found.

    Example:
        ```python
        # Single file
        paths = resolve_video_sources_to_paths("./video.mp4")
        # Returns: [Path("/abs/path/to/video.mp4")]

        # Directory
        paths = resolve_video_sources_to_paths("./my_videos/")
        # Returns: [Path("/abs/.../vid1.mp4"), Path("/abs/.../vid2.mp4"), ...]

        # Mixed list
        paths = resolve_video_sources_to_paths(
            ["./video.mp4", "./folder/", "~/vid.mov"]
        )
        # Returns: All resolved video paths

        # Recursive directory search
        paths = resolve_video_sources_to_paths("./dataset/", recursive=True)
        # Returns: All videos including subdirectories
        ```

    """
    if not sources:
        raise ValueError("Sources cannot be empty")

    resolved_paths: list[Path] = []

    # Handle single source (string or Path)
    if isinstance(sources, (str, Path)):
        source_path = Path(sources).expanduser().resolve(strict=False)

        # Check if it exists
        if not source_path.exists():
            raise FileNotFoundError(
                f"Source not found: '{sources}'\nResolved to: '{source_path}'"
            )

        # If it's a directory, get all videos
        if source_path.is_dir():
            resolved_paths = resolve_directory_to_videos(
                source_path, recursive=recursive
            )
        # If it's a file, validate and add
        else:
            resolved_paths = [resolve_and_validate_path(sources)]

    # Handle sequence of sources
    else:
        errors: list[tuple[int, str | Path, Exception]] = []

        for idx, source in enumerate(sources):
            try:
                source_path = Path(source).expanduser().resolve(strict=False)

                if not source_path.exists():
                    raise FileNotFoundError(f"Source not found: '{source}'")

                # If directory, expand to videos
                if source_path.is_dir():
                    dir_videos = resolve_directory_to_videos(
                        source_path, recursive=recursive
                    )
                    resolved_paths.extend(dir_videos)
                # If file, validate and add
                else:
                    resolved_paths.append(resolve_and_validate_path(source))

            except (FileNotFoundError, ValueError) as e:
                errors.append((idx, source, e))

        # If there are any errors, raise with detailed information
        if errors:
            error_messages = [
                f"  [{idx}] '{path}': {type(error).__name__}: {error}"
                for idx, path, error in errors
            ]
            raise FileNotFoundError(
                f"Failed to resolve {len(errors)} of {len(sources)} sources:\n"
                + "\n".join(error_messages)
            )

    # Remove duplicates while preserving order
    seen = set()
    unique_paths = []
    for path in resolved_paths:
        if path not in seen:
            seen.add(path)
            unique_paths.append(path)

    return unique_paths
