#!/usr/bin/env python

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

@File    :   media_utils.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Utility functions for media processing (images and videos).
"""

import base64
import io

import av
import httpx
import numpy as np
from PIL import Image
from transformers.video_utils import VideoMetadata


def decode_base64_image(base64_image: str) -> Image.Image:
    """Decode a base64-encoded image string to a PIL Image.

    Args:
        base64_image: Base64-encoded image string. Can be raw base64
            (e.g., "/9j/4AAQSkZJRgABAQEAYABgAAD...") or a data URI
            (e.g., "data:image/jpeg;base64,/9j/...").

    Returns:
        PIL Image object.

    Raises:
        ValueError: If the base64 string cannot be decoded.

    """
    # Extract base64 data from data URI if present
    if base64_image.startswith("data:image"):
        if "," in base64_image:
            base64_image = base64_image.split(",", 1)[1]

    # Decode base64 and convert to PIL Image
    try:
        image_bytes = base64.b64decode(base64_image)
        return Image.open(io.BytesIO(image_bytes))
    except Exception as e:
        raise ValueError(f"Failed to decode base64 image: {e}") from e


def decode_base64_video(base64_video: str) -> np.ndarray:
    """Decode a base64-encoded video string to a video frames array.

    Uses PyAV to decode video in-memory without temporary files.

    Args:
        base64_video: Base64-encoded video string. Can be raw base64 or a data URI
            (e.g., "data:video/mp4;base64,...").

    Returns:
        Numpy array with shape (num_frames, height, width, channels) in RGB format.

    Raises:
        ValueError: If the base64 string cannot be decoded or video decoding fails.

    """
    # Extract base64 data from data URI if present
    if base64_video.startswith("data:video"):
        if "," in base64_video:
            base64_video = base64_video.split(",", 1)[1]

    # Decode base64
    try:
        video_bytes = base64.b64decode(base64_video)
    except Exception as e:
        raise ValueError(f"Failed to decode base64 video: {e}") from e

    # Decode video using PyAV
    try:
        container = av.open(io.BytesIO(video_bytes))
        frames = []
        for frame in container.decode(video=0):
            frames.append(frame.to_ndarray(format="rgb24"))
        container.close()

        if not frames:
            raise ValueError("No video frames found in the provided data")

        # Stack into (N, H, W, C) array in RGB format
        return np.stack(frames)
    except ValueError:
        raise
    except Exception as e:
        raise ValueError(f"Failed to decode video with PyAV: {e}") from e


def download_image_from_url(url: str, timeout: float = 30.0) -> Image.Image:
    """Download an image from URL and return as PIL Image.

    Args:
        url: URL to download from (http:// or https://)
        timeout: Request timeout in seconds. Defaults to 30.0.

    Returns:
        PIL Image object.

    Raises:
        ValueError: If download fails or URL is invalid
        httpx.HTTPError: If the HTTP request fails

    Example:
        ```python
        # Download image from URL
        image = download_image_from_url("https://example.com/image.jpg")
        print(f"Image size: {image.size}")
        ```

    """
    try:
        # Download with httpx
        with httpx.Client(timeout=timeout, follow_redirects=True) as client:
            response = client.get(url)
            response.raise_for_status()

        # Decode image from response content
        return Image.open(io.BytesIO(response.content))

    except httpx.HTTPError as e:
        raise ValueError(f"Failed to download from URL '{url}': {e}") from e
    except Exception as e:
        raise ValueError(f"Error downloading image from URL '{url}': {e}") from e


def download_video_from_url(url: str, timeout: float = 30.0) -> np.ndarray:
    """Download a video from URL and return as frames array.

    Uses PyAV to decode video in-memory without temporary files.

    Args:
        url: URL to download from (http:// or https://)
        timeout: Request timeout in seconds. Defaults to 30.0.

    Returns:
        Numpy array with shape (num_frames, height, width, channels) in RGB format.

    Raises:
        ValueError: If download fails or video decoding fails
        httpx.HTTPError: If the HTTP request fails

    Example:
        ```python
        # Download video from URL
        frames = download_video_from_url("https://example.com/video.mp4")
        print(f"Video frames shape: {frames.shape}")
        ```

    """
    try:
        # Download with httpx
        with httpx.Client(timeout=timeout, follow_redirects=True) as client:
            response = client.get(url)
            response.raise_for_status()

        video_bytes = response.content

    except httpx.HTTPError as e:
        raise ValueError(f"Failed to download from URL '{url}': {e}") from e
    except Exception as e:
        raise ValueError(f"Error downloading from URL '{url}': {e}") from e

    # Decode video using PyAV
    try:
        container = av.open(io.BytesIO(video_bytes))
        frames = []
        for frame in container.decode(video=0):
            frames.append(frame.to_ndarray(format="rgb24"))
        container.close()

        if not frames:
            raise ValueError("No video frames found in the downloaded video")

        # Stack into (N, H, W, C) array in RGB format
        return np.stack(frames)
    except ValueError:
        raise
    except Exception as e:
        raise ValueError(f"Failed to decode video with PyAV: {e}") from e


def extract_video_metadata(source: str | bytes) -> VideoMetadata:
    """Extract video metadata from a file path, URL, or bytes.

    Uses PyAV to efficiently extract metadata without decoding all frames.

    Args:
        source: Video source - can be file path, URL, or bytes.

    Returns:
        VideoMetadata object containing video information (fps, num_frames, duration, etc.)

    Raises:
        ValueError: If metadata extraction fails.

    Example:
        ```python
        # Extract metadata from file
        metadata = extract_video_metadata("video.mp4")
        print(f"FPS: {metadata.fps}, Frames: {metadata.total_num_frames}")

        # Extract metadata from URL
        metadata = extract_video_metadata("https://example.com/video.mp4")
        ```

    """
    try:
        # Prepare video data based on source type
        video_data: str | io.BytesIO = source  # type: ignore[assignment]

        if isinstance(source, bytes):
            video_data = io.BytesIO(source)
        elif isinstance(source, str):
            # Handle data URI - decode base64 to bytes
            if source.startswith("data:video"):
                if "," in source:
                    base64_data = source.split(",", 1)[1]
                    video_data = io.BytesIO(base64.b64decode(base64_data))
            # Validate file path exists
            elif not source.startswith(("http://", "https://")):
                import os

                if not os.path.exists(source):
                    raise ValueError(
                        f"Video file not found: {source}. "
                        "Please check that the file path is correct."
                    )
                if not os.path.isfile(source):
                    raise ValueError(
                        f"Path is not a file: {source}. "
                        "Please provide a valid video file path."
                    )

        # Open container to extract metadata (efficient - doesn't decode frames)
        # PyAV can handle file paths, URLs, and BytesIO directly
        container = av.open(video_data)
        stream = container.streams.video[0]

        # Extract FPS with fallbacks
        fps = None
        if stream.average_rate is not None:
            fps = float(stream.average_rate)
        elif stream.guessed_rate is not None:
            fps = float(stream.guessed_rate)

        # Extract frame count with fallbacks
        total_num_frames = stream.frames or None
        if not total_num_frames and container.duration and fps:
            # Estimate from duration and fps
            total_num_frames = int(float(container.duration) / av.time_base * fps)

        # Extract duration with fallbacks
        duration = None
        if stream.duration is not None and stream.time_base is not None:
            duration = float(stream.duration * stream.time_base)
        elif container.duration is not None:
            duration = float(container.duration) / av.time_base

        # Extract dimensions
        width = stream.codec_context.width
        height = stream.codec_context.height

        container.close()

        return VideoMetadata(
            total_num_frames=total_num_frames,
            fps=fps,
            width=width,
            height=height,
            duration=duration,
            video_backend="pyav",
        )
    except ValueError:
        # Re-raise ValueError as-is (contains our custom error messages)
        raise
    except Exception as e:
        # Add more context to other errors
        error_msg = f"Failed to extract video metadata from source: {e}"
        if "Invalid data" in str(e) or "No such file" in str(e):
            error_msg += (
                "\n\nPossible causes:"
                "\n  - Video file is corrupted or incomplete"
                "\n  - Video format is not supported"
                "\n  - File is not a valid video file"
                "\n\nSupported formats: MP4, AVI, MOV, MKV, and other common video formats."
            )
        raise ValueError(error_msg) from e
