"""Media Asset Module with Lazy Loading

Defines the MediaAsset class for handling media file paths within assets.
Implements lazy loading for URL-based media to prevent unnecessary downloads.
"""

import os
from io import BytesIO
from rapidata.rapidata_client.assets._base_asset import BaseAsset
import requests
import re
from PIL import Image
from tinytag import TinyTag
import tempfile
from pydantic import StrictStr, StrictBytes
from typing import Optional, cast
import logging
from functools import cached_property


class MediaAsset(BaseAsset):
    """MediaAsset Class with Lazy Loading

    Represents a media asset by storing the file path or URL.
    Only downloads URL content when needed.
    Supports local files and URLs for images, MP3, and MP4.

    Args:
        path (str): The file system path to the media asset or URL.

    Raises:
        FileNotFoundError: If the provided file path does not exist.
    """
    _logger = logging.getLogger(__name__ + '.MediaAsset')

    ALLOWED_TYPES = [
        'image/', 
        'audio/mp3',      # MP3
        'video/mp4',       # MP4
    ]

    MIME_TYPES = {
        'jpg': 'image/jpeg',
        'jpeg': 'image/jpeg',
        'png': 'image/png',
        'gif': 'image/gif',
        'webp': 'image/webp',
        'mp3': 'audio/mp3',
        'mp4': 'video/mp4'
    }

    FILE_SIGNATURES = {
        b'\xFF\xD8\xFF': 'image/jpeg',
        b'\x89PNG\r\n\x1a\n': 'image/png',
        b'GIF87a': 'image/gif',
        b'GIF89a': 'image/gif',
        b'RIFF': 'image/webp',
        b'ID3': 'audio/mp3',
        b'\xFF\xFB': 'audio/mp3',
        b'\xFF\xF3': 'audio/mp3',
        b'ftyp': 'video/mp4',
    }

    def __init__(self, path: str):
        """
        Initialize a MediaAsset instance.

        Args:
            path (str): The file system path to the media asset or a URL.

        Raises:
            FileNotFoundError: If the provided file path does not exist.
            ValueError: If path is not a string.
        """
        if not isinstance(path, str):
            raise ValueError("Media must be a string, either a local file path or a URL")
        
        self._url = None
        self._content = None
        
        if re.match(r'^https?://', path):
            self._url = path
            self.name = path.split('/')[-1]
            self.name = self.__check_name_ending(self.name)
            self.path = None  # Will be set when content is downloaded
            return
        
        if not os.path.exists(path):
            raise FileNotFoundError(f"File not found: {path}")
        
        self.path = path
        self.name = path

    @cached_property
    def content(self) -> bytes:
        """
        Lazy loader for URL content. Only downloads when first accessed.
        Uses cached_property to store the result after first download.
        """
        if self._url is None:
            self.path = cast(str, self.path)
            with open(self.path, 'rb') as f:
                return f.read()
            
        return self.__get_media_bytes(self._url)

    def get_duration(self) -> int:
        """
        Get the duration of audio/video files in milliseconds.
        Returns 0 for static images.

        Returns:
            int: Duration in milliseconds for audio/video, 0 for static images

        Raises:
            ValueError: If the duration cannot be determined
        """
        path_to_check = self.name.lower()
        
        # Return 0 for static images
        if any(path_to_check.endswith(ext) for ext in ('.jpg', '.jpeg', '.png', '.webp', '.gif')):
            return 0

        try:
            # Create temporary file from content
            with tempfile.NamedTemporaryFile(suffix=os.path.splitext(self.name)[1], delete=False) as tmp:
                tmp.write(self.content)
                tmp.flush()
                tmp_path = tmp.name
                
                try:
                    tag = TinyTag.get(tmp_path)
                finally:
                    # Clean up the temporary file
                    os.unlink(tmp_path)
            
            if tag.duration is None:
                raise ValueError("Could not read duration from file")
                
            return int(tag.duration * 1000)  # Convert to milliseconds
            
        except Exception as e:
            raise ValueError(f"Could not determine media duration: {str(e)}")

    def get_image_dimension(self) -> tuple[int, int] | None:
        """
        Get the dimensions (width, height) of an image file.
        Returns None for non-image files or if dimensions can't be determined.
        """
        if not any(self.name.lower().endswith(ext) for ext in ('.jpg', '.jpeg', '.png', '.gif', '.webp')):
            return None
            
        try:
            img = Image.open(BytesIO(self.content))
            return img.size
        except Exception:
            return None

    def set_custom_name(self, name: str) -> 'MediaAsset':
        """Set a custom name for the media asset (only works with URLs)."""
        if self._url is not None:
            self.name = self.__check_name_ending(name)
        else:
            raise ValueError("Custom name can only be set for URLs.")
        return self
    
    def __check_name_ending(self, name: str) -> str:
        """Check if the media path is valid."""
        if not name.endswith(('.jpg', '.jpeg', '.png', '.gif', '.mp3', '.mp4', '.webp')):
            print("Warning: Supported file types: jpg, jpeg, png, gif, mp3, mp4. Image might not be displayed correctly.")
            name = name + '.jpg'
        return name

    def __get_media_type_from_extension(self, url: str) -> Optional[str]:
        """
        Determine media type from URL file extension.
        
        Args:
            url: The URL to check
            
        Returns:
            Optional[str]: MIME type if valid extension found, None otherwise
        """
        try:
            ext = url.lower().split('?')[0].split('.')[-1]
            return self.MIME_TYPES.get(ext)
        except IndexError:
            return None
        
    def __validate_image_content(self, content: bytes) -> bool:
        """
        Validate image content using PIL.
        
        Args:
            content: Image bytes to validate
            
        Returns:
            bool: True if valid image, False otherwise
        """
        try:
            img = Image.open(BytesIO(content))
            img.verify()
            return True
        except Exception as e:
            self._logger.debug(f"Image validation failed: {str(e)}")
            return False
        
    def __get_media_type_from_signature(self, content: bytes) -> Optional[str]:
        """
        Determine media type from file signature.
        
        Args:
            content: File content bytes
            
        Returns:
            Optional[str]: MIME type if valid signature found, None otherwise
        """
        file_start = content[:32]
        for signature, mime_type in self.FILE_SIGNATURES.items():
            if file_start.startswith(signature) or (signature in file_start[:10]):
                return mime_type
        return None

    def __get_media_bytes(self, url: str) -> bytes:
        """
        Downloads and validates media files from URL.
        
        Args:
            url: URL of the media file
            
        Returns:
            bytes: Validated media content
            
        Raises:
            ValueError: If media type is unsupported or content validation fails
            requests.exceptions.RequestException: If download fails
        """
        try:
            response = requests.get(url, stream=False)
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            self._logger.error(f"Failed to download media from {url}: {str(e)}")
            raise

        content = response.content
        content_type = response.headers.get('content-type', '').lower()

        # Case 1: Content-type is already allowed
        if any(content_type.startswith(t) for t in self.ALLOWED_TYPES):
            self._logger.debug(f"Content-type {content_type} is allowed")
            return content

        # Case 2: Try to validate based on extension
        mime_type = self.__get_media_type_from_extension(url)
        if mime_type and mime_type.startswith(tuple(self.ALLOWED_TYPES)):
            self._logger.debug(f"Found valid mime type from extension: {mime_type}")
            return content

        # Case 3: Try to validate based on file signature
        mime_type = self.__get_media_type_from_signature(content)
        if mime_type and mime_type.startswith(tuple(self.ALLOWED_TYPES)):
            self._logger.debug(f"Found valid mime type from signature: {mime_type}")
            return content

        # Case 4: Last resort - try direct image validation
        if self.__validate_image_content(content):
            self._logger.debug("Content validated as image through direct validation")
            return content

        # If we get here, validation failed
        error_msg = (
            f'Could not validate media type from content.\n'
            f'Content-Type: {content_type}\n'
            f'URL extension: {url.split("?")[0].split(".")[-1]}\n'
            f'Allowed types: {self.ALLOWED_TYPES}'
        )
        self._logger.error(error_msg)
        raise ValueError(error_msg)
    
    def to_file(self) -> StrictStr | tuple[StrictStr, StrictBytes] | StrictBytes:
        """Convert the media asset to a file representation."""
        if self._url is None:
            self.path = cast(str, self.path)
            return self.path
        else:
            return (self.name, self.content)
