import re
from typing import Dict, Any, Literal, List, NamedTuple, Union

from . import attachment, attachment_context, decoder_base64


class ExtractionResult(NamedTuple):
    attachments: List[attachment.Attachment]
    sanitized_data: Any


class AttachmentsExtractor:
    """
    Extracts and processes attachments embedded as Base64 strings within data structures.

    This class is designed to identify and decode Base64-encoded attachments located
    within the provided data. It uses a regular expression pattern to search for
    Base64 strings that meet a specified minimum length. Extracted attachments are
    decoded and replaced with sanitized placeholders in the original data.
    """

    def __init__(self, min_attachment_size: int):
        """
        Initializes the class with a minimum attachment size and configures the base64
        pattern for decoding attachments based on its length.

        Args:
            min_attachment_size: The minimum size of the attachment in characters
                for it to be considered valid. This ensures that only large enough
                base64 strings are matched to minimize false positives.
        """
        self._min_attachment_size = min_attachment_size
        self.decoder = decoder_base64.Base64AttachmentDecoder()

        # Pattern to match base64 strings (can be embedded in text)
        # Requires at least min_attachment_size characters to reduce false positives
        min_base64_groups = int(min_attachment_size / 4)
        BASE64_PATTERN = (
            r"(?:[A-Za-z0-9+/]{4}){"
            + str(min_base64_groups)
            + ",}(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?"
        )
        self.pattern = re.compile(BASE64_PATTERN)

    def extract_and_replace(
        self,
        data: Union[Dict[str, Any], List[Any]],
        entity_type: Literal["span", "trace"],
        entity_id: str,
        project_name: str,
        context: Literal["input", "output", "metadata"],
    ) -> List[attachment_context.AttachmentWithContext]:
        """
        Extract attachments from data and replace with placeholders.

        Handles both dict and list at the top level, recursively processing
        nested structures to find and extract base64-encoded attachments.

        Args:
            data: The data structure to process (dict or list)
            entity_type: Type of entity (span or trace)
            entity_id: ID of the entity
            project_name: Name of the project
            context: Context where data is located (input, output, or metadata)

        Returns:
            List of extracted attachments with context
        """
        attachments: List[attachment_context.AttachmentWithContext] = []

        if isinstance(data, dict):
            # For dicts, iterate over items and extract attachments from values
            for key, value in data.items():
                extraction_result = self._try_extract_attachments(value, context)
                if extraction_result.attachments:
                    # replace the original value with the sanitized one and collect attachments
                    data[key] = extraction_result.sanitized_data
                    for extracted_attachment in extraction_result.attachments:
                        attachments.append(
                            attachment_context.AttachmentWithContext(
                                attachment_data=extracted_attachment,
                                entity_type=entity_type,
                                entity_id=entity_id,
                                project_name=project_name,
                                context=context,
                            )
                        )
        elif isinstance(data, list):
            # For lists, extract attachments from each item and replace in place
            extraction_result = self._try_extract_attachments(data, context)
            # Replace list contents with sanitized version
            data.clear()
            data.extend(extraction_result.sanitized_data)
            # Convert extracted attachments to AttachmentWithContext
            for extracted_attachment in extraction_result.attachments:
                attachments.append(
                    attachment_context.AttachmentWithContext(
                        attachment_data=extracted_attachment,
                        entity_type=entity_type,
                        entity_id=entity_id,
                        project_name=project_name,
                        context=context,
                    )
                )
        else:
            # For other types (str, int, bool, None, etc.), try to extract but don't mutate
            extraction_result = self._try_extract_attachments(data, context)
            for extracted_attachment in extraction_result.attachments:
                attachments.append(
                    attachment_context.AttachmentWithContext(
                        attachment_data=extracted_attachment,
                        entity_type=entity_type,
                        entity_id=entity_id,
                        project_name=project_name,
                        context=context,
                    )
                )

        return attachments

    def _try_extract_attachments(
        self, data: Any, context: Literal["input", "output", "metadata"]
    ) -> ExtractionResult:
        """
        Recursively extract attachments from data that can be a string, dict, list, or other type.

        Args:
            data: The data to process (can be str, dict, list, or other types)
            context: The context where the data is located (input, output, or metadata)

        Returns:
            ExtractionResult with extracted attachments and sanitized data
        """
        # Handle string data - check for base64 attachments
        if isinstance(data, str):
            return self._extract_from_string(data, context)

        # Handle dictionary data - recursively process each value
        elif isinstance(data, dict):
            return self._extract_from_dict(data, context)

        # Handle list data - recursively process each element
        elif isinstance(data, list):
            return self._extract_from_list(data, context)

        # For other types (int, bool, None, etc.), return as-is
        else:
            return ExtractionResult(attachments=[], sanitized_data=data)

    def _extract_from_string(
        self, data: str, context: Literal["input", "output", "metadata"]
    ) -> ExtractionResult:
        """Extract attachments from a string value."""
        if len(data) < self._min_attachment_size:
            # skip short strings
            return ExtractionResult(attachments=[], sanitized_data=data)

        attachments: List[attachment.Attachment] = []
        sanitized_data = data
        for match in self.pattern.finditer(data):
            to_decode = match.group()
            decoded_attachment = self.decoder.decode(to_decode, context)
            if decoded_attachment is not None:
                attachments.append(decoded_attachment)
                sanitized_data = sanitized_data.replace(
                    to_decode, f"[{decoded_attachment.file_name}]"
                )

        return ExtractionResult(attachments=attachments, sanitized_data=sanitized_data)

    def _extract_from_dict(
        self, data: Dict[str, Any], context: Literal["input", "output", "metadata"]
    ) -> ExtractionResult:
        """Recursively extract attachments from a dictionary."""
        all_attachments: List[attachment.Attachment] = []
        sanitized_dict = {}

        for key, value in data.items():
            result = self._try_extract_attachments(value, context)
            sanitized_dict[key] = result.sanitized_data
            all_attachments.extend(result.attachments)

        return ExtractionResult(
            attachments=all_attachments, sanitized_data=sanitized_dict
        )

    def _extract_from_list(
        self, data: List[Any], context: Literal["input", "output", "metadata"]
    ) -> ExtractionResult:
        """Recursively extract attachments from a list."""
        all_attachments: List[attachment.Attachment] = []
        sanitized_list = []

        for item in data:
            result = self._try_extract_attachments(item, context)
            sanitized_list.append(result.sanitized_data)
            all_attachments.extend(result.attachments)

        return ExtractionResult(
            attachments=all_attachments, sanitized_data=sanitized_list
        )
