"""
Media response handler for Cliver client.

This module provides functionality for handling multimedia responses from LLMs,
including saving media content to local files and displaying text responses.
"""

import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from langchain_core.messages import BaseMessage

from cliver.media import MediaContent, MediaType, get_file_extension
from cliver.llm.base import LLMInferenceEngine

logger = logging.getLogger(__name__)


class MultimediaResponse:
    """
    Represents a multimedia response from an LLM.

    A MultimediaResponse can contain both text content and multimedia content
    such as images, audio, or video generated by the LLM.
    """

    def __init__(
        self,
        text_content: str = "",
        media_content: List[MediaContent] = None,
        raw_response: Optional[BaseMessage] = None,
    ):
        """
        Initialize a MultimediaResponse.

        Args:
            text_content: Text content of the response
            media_content: List of MediaContent objects
            raw_response: Raw response from the LLM
        """
        self.text_content = text_content
        self.media_content = media_content or []
        self.raw_response = raw_response

    def has_text(self) -> bool:
        """Check if the response has text content."""
        return bool(self.text_content and self.text_content.strip())

    def has_media(self) -> bool:
        """Check if the response has media content."""
        return len(self.media_content) > 0

    def get_media_by_type(self, media_type: MediaType) -> List[MediaContent]:
        """Get media content by type."""
        return [media for media in self.media_content if media.type == media_type]


class MultimediaResponseHandler:
    """
    Handler for processing multimedia responses from LLMs.

    This class provides methods to:
    1. Parse multimedia responses from LLMs
    2. Save media content to local files
    3. Display text responses appropriately
    4. Work both as CLI tool and library
    """

    def __init__(self, save_directory: Optional[str] = None):
        """
        Initialize the MultimediaResponseHandler.

        Args:
            save_directory: Directory to save media files. If None, uses current directory.
        """
        self.save_directory = Path(save_directory) if save_directory else Path.cwd()
        self.save_directory.mkdir(parents=True, exist_ok=True)

    def process_response(
        self, response: Union[str, BaseMessage], llm_engine: Optional['LLMInferenceEngine'] = None, auto_save_media: bool = False
    ) -> MultimediaResponse:
        """
        Process a response from an LLM, handling both text and multimedia content.

        Args:
            response: Response from the LLM (string or BaseMessage)
            llm_engine: LLM engine used to generate the response (for media extraction)
            auto_save_media: Whether to automatically save media content to files

        Returns:
            MultimediaResponse object containing parsed content
        """
        if isinstance(response, str):
            return MultimediaResponse(text_content=response)

        # Handle BaseMessage response
        text_content = ""

        if hasattr(response, 'content'):
            if isinstance(response.content, str):
                text_content = response.content
            elif isinstance(response.content, list):
                # Handle structured content (like OpenAI's multimodal format)
                for item in response.content:
                    if isinstance(item, dict):
                        if item.get('type') == 'text':
                            text_content += item.get('text', '')
                        # TODO: Handle other content types if needed

        # Extract media content using the LLM engine if provided
        media_content = []
        if llm_engine and response:
            try:
                media_content = llm_engine.extract_media_from_response(response)
            except Exception as e:
                logger.warning(f"Error extracting media from response: {e}")
                raise e

        multimedia_response = MultimediaResponse(
            text_content=text_content,
            media_content=media_content,
            raw_response=response
        )

        # Auto-save media if requested
        if auto_save_media and multimedia_response.has_media():
            saved_files = self.save_media_content(multimedia_response)
            if saved_files:
                logger.info(f"Saved media files: {saved_files}")

        return multimedia_response


    def save_media_content(
        self, response: MultimediaResponse, prefix: str = "cliver_media"
    ) -> List[str]:
        """
        Save media content from a response to local files.

        Args:
            response: MultimediaResponse containing media content
            prefix: Prefix for saved filenames

        Returns:
            List of paths to saved files
        """
        saved_files = []

        for i, media in enumerate(response.media_content):
            try:
                # Generate filename if not provided
                if not media.filename:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    extension = get_file_extension(media.mime_type)
                    media.filename = f"{prefix}_{timestamp}_{i}{extension}"

                # Create full file path
                file_path = self.save_directory / media.filename

                # Save the media content
                media.save(file_path)
                saved_files.append(str(file_path))
                logger.info(f"Saved media to {file_path}")

            except Exception as e:
                logger.error(f"Error saving media {media.filename}: {e}")
                raise e

        return saved_files

    @staticmethod
    def display_response(
            response: MultimediaResponse,
        show_text: bool = True,
        show_media_info: bool = True
    ):
        """
        Display a multimedia response.

        Args:
            response: MultimediaResponse to display
            show_text: Whether to show text content
            show_media_info: Whether to show media information
        """
        output = []

        if show_text and response.has_text():
            output.append(response.text_content)

        if show_media_info and response.has_media():
            output.append(f"\n[Media Content: {len(response.media_content)} items]")
            for i, media in enumerate(response.media_content):
                info = f"  {i+1}. {media.type.value}"
                if media.filename:
                    info += f" ({media.filename})"
                if media.mime_type:
                    info += f" [{media.mime_type}]"
                output.append(info)

        return '\n'.join(output)

    @staticmethod
    def get_response_summary(response: MultimediaResponse) -> Dict[str, Any]:
        """
        Get a summary of the multimedia response.

        Args:
            response: MultimediaResponse to summarize

        Returns:
            Dictionary with response summary
        """
        return {
            'has_text': response.has_text(),
            'text_length': len(response.text_content) if response.text_content else 0,
            'has_media': response.has_media(),
            'media_count': len(response.media_content),
            'media_types': [media.type.value for media in response.media_content],
            'media_filenames': [media.filename for media in response.media_content if media.filename]
        }

def save_response_media(
    response: Union[str, BaseMessage, MultimediaResponse],
    save_directory: Optional[str] = None,
    prefix: str = "cliver_media"
) -> List[str]:
    """
    Convenience function to save media from a response.

    Args:
        response: Response from LLM or MultimediaResponse object
        save_directory: Directory to save media files
        prefix: Prefix for saved filenames

    Returns:
        List of paths to saved files
    """
    handler = MultimediaResponseHandler(save_directory)

    if not isinstance(response, MultimediaResponse):
        response = handler.process_response(response)

    return handler.save_media_content(response, prefix)