#!/usr/bin/env python

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

@File    :   source_resolution.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Source resolution utilities that handle paths, URLs, and data URIs.
"""

from collections.abc import Sequence
from pathlib import Path

from vi.inference.types import SourceType
from vi.inference.utils.path_utils import (
    is_data_uri,
    is_url,
    resolve_directory_to_images,
    resolve_file_path,
)


class ResolvedSource:
    """Represents a resolved source with type information.

    Attributes:
        value: The source value (file path, URL, or data URI)
        source_type: Type of source (SourceType enum)

    """

    def __init__(self, value: str, source_type: SourceType):
        """Initialize resolved source.

        Args:
            value: The source value
            source_type: Type of source (SourceType enum)

        """
        self.value = value
        self.source_type = source_type

    def __str__(self) -> str:
        """Return string representation."""
        return self.value

    def __repr__(self) -> str:
        """Return detailed representation."""
        return f"ResolvedSource(value='{self.value}', type={self.source_type})"


def resolve_single_source(source: str | Path) -> ResolvedSource:
    """Resolve a single source to its appropriate type.

    Args:
        source: File path, URL, or data URI

    Returns:
        ResolvedSource object with value and type

    Raises:
        FileNotFoundError: If file path doesn't exist
        ValueError: If source is invalid

    Example:
        ```python
        # URL
        resolved = resolve_single_source("https://example.com/image.jpg")
        assert resolved.source_type == SourceType.URL

        # Data URI
        resolved = resolve_single_source("data:image/jpeg;base64,...")
        assert resolved.source_type == SourceType.DATA_URI

        # File path
        resolved = resolve_single_source("/path/to/image.jpg")
        assert resolved.source_type == SourceType.PATH
        ```

    """
    source_str = str(source)

    # Check if URL
    if is_url(source_str):
        return ResolvedSource(source_str, SourceType.URL)

    # Check if data URI
    if is_data_uri(source_str):
        return ResolvedSource(source_str, SourceType.DATA_URI)

    # Otherwise it's a file path - validate it
    resolved_path = resolve_file_path(source)
    return ResolvedSource(str(resolved_path), SourceType.PATH)


def resolve_sources(
    sources: str | Path | Sequence[str | Path],
    recursive: bool = False,
) -> list[ResolvedSource]:
    """Resolve source(s) to a list of ResolvedSource objects.

    Handles three input types:
    1. Single source (file/URL/data URI) → Returns [resolved_source]
    2. Directory path → Returns list of all images in directory
    3. List of sources → Returns list of all resolved sources

    Args:
        sources: Can be:
            - Single file path (str or Path)
            - Single URL (http:// or https://)
            - Single data URI (data:image/...)
            - Directory path (str or Path)
            - List of any combination of the above
        recursive: If True, recursively search subdirectories when
            processing directory paths. Defaults to False.

    Returns:
        List of ResolvedSource objects

    Raises:
        FileNotFoundError: If any file/directory does not exist
        ValueError: If any source is invalid or no files found

    Example:
        ```python
        # Single file
        sources = resolve_sources("./image.jpg")
        # Returns: [ResolvedSource("/abs/path/to/image.jpg", "path")]

        # URL
        sources = resolve_sources("https://example.com/image.jpg")
        # Returns: [ResolvedSource("https://...", "url")]

        # Mixed list
        sources = resolve_sources(
            [
                "./image.jpg",
                "https://example.com/image2.jpg",
                "data:image/jpeg;base64,...",
            ]
        )
        # Returns: [ResolvedSource(...), ResolvedSource(...), ResolvedSource(...)]

        # Directory
        sources = resolve_sources("./my_images/")
        # Returns: List of ResolvedSource objects for all images
        ```

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

    resolved_sources: list[ResolvedSource] = []

    # Handle single source
    if isinstance(sources, (str, Path)):
        source_str = str(sources)

        # URLs and data URIs are never directories
        if is_url(source_str) or is_data_uri(source_str):
            return [resolve_single_source(source_str)]

        # Check if it's a directory
        source_path = Path(source_str).expanduser().resolve(strict=False)

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

        if source_path.is_dir():
            # Get all image paths from directory
            image_paths = resolve_directory_to_images(source_path, recursive=recursive)
            resolved_sources = [
                ResolvedSource(str(p), SourceType.PATH) for p in image_paths
            ]
        else:
            # Single file
            resolved_sources = [resolve_single_source(sources)]

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

        for idx, source in enumerate(sources):
            try:
                source_str = str(source)

                # URLs and data URIs are never directories
                if is_url(source_str) or is_data_uri(source_str):
                    resolved_sources.append(resolve_single_source(source))
                    continue

                # Check if it's a directory
                source_path = Path(source_str).expanduser().resolve(strict=False)

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

                if source_path.is_dir():
                    # Get all images from directory
                    dir_images = resolve_directory_to_images(
                        source_path, recursive=recursive
                    )
                    resolved_sources.extend(
                        [ResolvedSource(str(p), SourceType.PATH) for p in dir_images]
                    )
                else:
                    # Single file
                    resolved_sources.append(resolve_single_source(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_sources = []
    for source in resolved_sources:
        if source.value not in seen:
            seen.add(source.value)
            unique_sources.append(source)

    return unique_sources
