"""Media conversion pipeline for transforming messages with arbitrary media types.

Provides the :class:`MediaConverter` class which maintains an ordered list of
:class:`BaseConverter` instances and applies them to message blocks, converting
unsupported media types into allowed ones.
"""

from typing import List, Optional, Sequence, Set, Type

from arbol import aprint

from litemind.agent.messages.message import Message
from litemind.agent.messages.message_block import MessageBlock
from litemind.media.conversion.converters.base_converter import BaseConverter
from litemind.media.media_base import MediaBase
from litemind.media.types.media_action import Action
from litemind.media.types.media_types import all_media_types


class MediaConverter:
    """Pipeline for converting messages with arbitrary media to a restricted set of allowed types.

    Maintains an ordered list of ``BaseConverter`` instances. When
    converting, the first converter that can handle a given media takes
    precedence. Conversion is applied recursively by default so that
    intermediate types (e.g. Table -> Text) are fully resolved.
    """

    def __init__(self):
        """Initialise the media converter pipeline with an empty converter list."""
        # List of media converters
        self.media_converters: List[BaseConverter] = []

    def add_default_converters(
        self,
        convert_audio: bool = True,
        convert_videos: bool = True,
        convert_documents: bool = True,
    ) -> List[BaseConverter]:
        """Register the standard set of converters.

        Adds converters for NdImage, Table, Object, Code, Json, and
        optionally Video (requires ffmpeg), Document (requires pymupdf
        and/or docling), and Audio (requires local Whisper). Converters
        are only added when their dependencies are available.

        Parameters
        ----------
        convert_audio : bool, optional
            Include the Whisper-based audio converter. Default is True.
        convert_videos : bool, optional
            Include the ffmpeg-based video converter. Default is True.
        convert_documents : bool, optional
            Include document converters (pymupdf, docling, txt fallback).
            Default is True.

        Returns
        -------
        List[BaseConverter]
            The converters that were actually added.
        """

        # Create a list of default converters
        default_converters = []

        # Import the default converters
        from litemind.media.conversion.converters.code_converter import CodeConverter
        from litemind.media.conversion.converters.document_converter_docling import (
            is_docling_available,
        )
        from litemind.media.conversion.converters.document_converter_pymupdf import (
            is_pymupdf_available,
        )
        from litemind.media.conversion.converters.document_converter_txt import (
            DocumentConverterTxt,
        )
        from litemind.media.conversion.converters.file_converter import FileConverter
        from litemind.media.conversion.converters.json_converter import JsonConverter
        from litemind.media.conversion.converters.ndimage_converter import (
            NdImageConverter,
        )
        from litemind.media.conversion.converters.object_converter import (
            ObjectConverter,
        )
        from litemind.media.conversion.converters.table_converter import TableConverter

        # Import the is_* functions:
        from litemind.utils.ffmpeg_utils import is_ffmpeg_available
        from litemind.utils.whisper_transcribe_audio import is_local_whisper_available

        # Add default converters to the list
        default_converters.append(NdImageConverter())
        default_converters.append(TableConverter())
        default_converters.append(ObjectConverter())
        default_converters.append(CodeConverter())
        default_converters.append(JsonConverter())

        if convert_videos and is_ffmpeg_available():
            from litemind.media.conversion.converters.video_converter_ffmpeg import (
                VideoConverterFfmpeg,
            )

            # Add the video converter to the list of default converters
            # This converter is only available if ffmpeg is installed
            default_converters.append(VideoConverterFfmpeg())

        if convert_documents and is_pymupdf_available():
            from litemind.media.conversion.converters.document_converter_pymupdf import (
                DocumentConverterPymupdf,
            )

            # Add the document converter to the list of default converters
            # This converter is only available if pymupdf is installed
            default_converters.append(DocumentConverterPymupdf())

        if convert_documents and is_docling_available():
            from litemind.media.conversion.converters.document_converter_docling import (
                DocumentConverterDocling,
            )

            # Add the document converter to the list of default converters
            # This converter is only available if docling is installed
            default_converters.append(DocumentConverterDocling())

        if convert_audio and is_local_whisper_available():
            from litemind.media.conversion.converters.audio_converter_whisper_local import (
                AudioConverterWhisperLocal,
            )

            # Add the audio converter to the list of default converters
            # This converter is only available if local whisper is installed
            default_converters.append(AudioConverterWhisperLocal())

        # This converter is at the end of the list to have the lowest priority
        if convert_documents:
            default_converters.append(DocumentConverterTxt())

        # Add FileConverter to the list of default converters
        default_converters.append(FileConverter())

        # Add default converters to the list of media converters
        self.media_converters.extend(default_converters)

        return default_converters

    def add_media_converter(
        self, media_converter: BaseConverter, highest_priority: bool = False
    ):
        """Add a converter to the pipeline.

        Parameters
        ----------
        media_converter : BaseConverter
            The converter to add.
        highest_priority : bool, optional
            If True, insert at the front of the list so it takes precedence
            over existing converters. Default is False (append).
        """
        if highest_priority:
            # Add the media converter to the beginning of the list
            self.media_converters.insert(0, media_converter)
        else:
            # Add the media converter to the end of the list
            self.media_converters.append(media_converter)

    def remove_media_converter(self, media_converter: BaseConverter):
        """Remove a converter from the pipeline.

        Parameters
        ----------
        media_converter : BaseConverter
            The converter instance to remove.

        Raises
        ------
        ValueError
            If the converter is not in the pipeline.
        """
        self.media_converters.remove(media_converter)

    def can_convert_within(
        self,
        source_media_type: Type[MediaBase],
        allowed_media_types: Set[Type[MediaBase]],
    ) -> bool:
        """Check whether a source type can be fully converted to allowed types.

        Performs a BFS over the conversion graph to verify that every
        reachable type from *source_media_type* is contained in
        *allowed_media_types*.

        Parameters
        ----------
        source_media_type : Type[MediaBase]
            The media type to start from.
        allowed_media_types : Set[Type[MediaBase]]
            The set of acceptable target types.

        Returns
        -------
        bool
            True if all conversion outputs are within the allowed set.
        """

        # If there are no converters, then we cannot convert anything, source type must be in allowed types:
        if not self.media_converters:
            return source_media_type in allowed_media_types

        # Build a conversion graph from all converter rules
        conversion_graph = {}

        # Process all converter rules to build graph
        for converter in self.media_converters:
            for src_type, target_types in converter.rule():
                if src_type not in conversion_graph:
                    conversion_graph[src_type] = set()
                # Add all possible conversion targets
                conversion_graph[src_type].update(target_types)

        # Queue for BFS traversal
        queue = [source_media_type]
        # Keep track of visited types to avoid cycles
        visited = set()

        # BFS to find all possible conversion results
        while queue:
            current_type = queue.pop(0)

            # Skip if already visited
            if current_type in visited:
                continue

            visited.add(current_type)

            # Check if current type can be converted
            if current_type in conversion_graph:
                for next_type in conversion_graph[current_type]:
                    # If any next type is not in allowed types, return False
                    if next_type not in allowed_media_types:
                        return False

                    if next_type not in visited:
                        queue.append(next_type)

        # If we get here, all conversions are within allowed types
        return True

    def get_convertible_media_types(
        self, allowed_media_types: Set[Type[MediaBase]]
    ) -> Set[Type[MediaBase]]:
        """Find all media types that can be converted to the allowed set.

        Parameters
        ----------
        allowed_media_types : Set[Type[MediaBase]]
            The set of acceptable target types.

        Returns
        -------
        Set[Type[MediaBase]]
            All media types for which ``can_convert_within`` returns True.
        """
        convertable_media_types = set()

        # iterate through all media types that derive from MediaBase:
        for media_type in all_media_types():

            # Check if the converter can convert this media type:
            if self.can_convert_within(media_type, allowed_media_types):
                # If it can, add it to the set of convertable media types
                convertable_media_types.add(media_type)

        return convertable_media_types

    def convert(
        self,
        messages: List[Message],
        allowed_media_types: Sequence[Type[MediaBase]],
        exclude_extensions: Optional[Sequence[str]] = None,
        recursive: bool = True,
    ) -> List[Message]:
        """Convert messages so they contain only allowed media types.

        Iterates over each message block and applies the first matching
        converter. The conversion is **not** performed in place; new
        messages are returned.

        Parameters
        ----------
        messages : List[Message]
            The messages to convert.
        allowed_media_types : Sequence[Type[MediaBase]]
            Media types that should be kept as-is.
        exclude_extensions : Sequence[str], optional
            File extensions to skip during conversion (media with these
            extensions are kept unchanged). Default is None (no exclusions).
        recursive : bool, optional
            If True, reapply conversion until no further changes occur.
            Default is True.

        Returns
        -------
        List[Message]
            New messages containing only allowed (or unconvertible) media.
        """

        # List that will hold the converted messages:
        converted_messages = []

        # By default no conversion happened:
        conversion_happened = False

        # Iterate over the messages:
        for message in messages:

            # Create a new message with the same role as the original message:
            new_message = Message(role=message.role)

            # Iterate over the blocks in the message:
            for block in message.blocks:

                # extract attributes from block:
                attributes = block.attributes

                # get the media from the block:
                media = block.media

                # Attempt to get the extension if defined:
                extension = getattr(media, "extension", None)

                if any(
                    isinstance(media, allowed_type)
                    for allowed_type in allowed_media_types
                ):
                    # If the media type is allowed, keep it as is:
                    new_message.append_block(block)
                elif (
                    extension is not None
                    and exclude_extensions is not None
                    and extension in exclude_extensions
                ):
                    # If the media type is not allowed and the extension is in the exclude list, skip conversion:
                    # aprint("Warning: Skipping media with excluded extension:", extension)
                    new_message.append_block(block)
                else:
                    # Otherwise, attempt to convert it using the available converters
                    converted = None

                    # Iterate over the media converters:
                    for converter in self.media_converters:
                        try:
                            # Check if the converter can convert the media:
                            if converter.can_convert(media):
                                # If the converter can convert the media, convert it:
                                converted = converter.convert(media)
                                # Break after the first successful conversion
                                break
                        except Exception as e:
                            aprint("Error during conversion:", e)
                            new_message.append_text(
                                f"Could not convert {type(media).__name__} to allowed types because of error {e}."
                            )
                            import traceback

                            traceback.print_exc()

                    if converted is not None:
                        conversion_happened = True
                        # If a converter was found, append the converted media
                        for converted_item in converted:
                            new_message.append_block(
                                MessageBlock(
                                    media=converted_item, attributes=attributes
                                )
                            )
                    else:
                        # If no converter is found, keep the original media:
                        new_message.append_block(block)

                        # if the original media is not allowed, log a warning:
                        if not any(
                            isinstance(media, allowed_type)
                            for allowed_type in allowed_media_types
                        ) and not isinstance(media, Action):
                            aprint(
                                f"Warning: No converter found for non-allowed media of type: {type(media)}"
                            )

            converted_messages.append(new_message)

        # If conversion happened, call convert recursively
        if recursive and conversion_happened:
            return self.convert(converted_messages, allowed_media_types, recursive=True)

        return converted_messages
