"""
This module contains the `DocumentParsingPipe` class, which is responsible for parsing incoming documents into plaintext.
"""

import asyncio
import logging
import time
import uuid
from typing import AsyncGenerator, Optional, Union

from r2r import parsers
from r2r.base import (
    AsyncState,
    Document,
    DocumentType,
    Extraction,
    ExtractionType,
    KVLoggingSingleton,
    PipeType,
    generate_id_from_label,
)
from r2r.base.abstractions.exception import R2RDocumentProcessingError
from r2r.base.pipes.base_pipe import AsyncPipe
from r2r.parsers.media.audio_parser import AudioParser
from r2r.parsers.media.docx_parser import DOCXParser
from r2r.parsers.media.img_parser import ImageParser
from r2r.parsers.media.movie_parser import MovieParser
from r2r.parsers.media.pdf_parser import PDFParser, PDFParserUnstructured
from r2r.parsers.media.ppt_parser import PPTParser
from r2r.parsers.structured.csv_parser import CSVParser, CSVParserAdvanced
from r2r.parsers.structured.json_parser import JSONParser
from r2r.parsers.structured.xlsx_parser import XLSXParser, XLSXParserAdvanced
from r2r.parsers.text.html_parser import HTMLParser
from r2r.parsers.text.md_parser import MDParser
from r2r.parsers.text.text_parser import TextParser

logger = logging.getLogger(__name__)


class ParsingPipe(AsyncPipe):
    """
    Processes incoming documents into plaintext based on their data type.
    Supports TXT, JSON, HTML, and PDF formats.
    """

    class Input(AsyncPipe.Input):
        message: AsyncGenerator[Document, None]

    AVAILABLE_PARSERS = {
        DocumentType.CSV: [CSVParser, CSVParserAdvanced],
        DocumentType.DOCX: [DOCXParser],
        DocumentType.HTML: [HTMLParser],
        DocumentType.JSON: [JSONParser],
        DocumentType.MD: [MDParser],
        DocumentType.PDF: [PDFParser, PDFParserUnstructured],
        DocumentType.PPTX: [PPTParser],
        DocumentType.TXT: [TextParser],
        DocumentType.XLSX: [XLSXParser, XLSXParserAdvanced],
        DocumentType.GIF: [ImageParser],
        DocumentType.JPEG: [ImageParser],
        DocumentType.JPG: [ImageParser],
        DocumentType.PNG: [ImageParser],
        DocumentType.SVG: [ImageParser],
        DocumentType.MP3: [AudioParser],
        DocumentType.MP4: [MovieParser],
    }

    IMAGE_TYPES = {
        DocumentType.GIF,
        DocumentType.JPG,
        DocumentType.JPEG,
        DocumentType.PNG,
        DocumentType.SVG,
    }

    def __init__(
        self,
        excluded_parsers: list[DocumentType],
        override_parsers: Optional[list[dict]] = None,
        pipe_logger: Optional[KVLoggingSingleton] = None,
        type: PipeType = PipeType.INGESTOR,
        config: Optional[AsyncPipe.PipeConfig] = None,
        *args,
        **kwargs,
    ):
        super().__init__(
            pipe_logger=pipe_logger,
            type=type,
            config=config
            or AsyncPipe.PipeConfig(name="default_document_parsing_pipe"),
            *args,
            **kwargs,
        )

        self.parsers = {}

        if not override_parsers:
            override_parsers = []

        # Apply overrides if specified
        for parser_override in override_parsers:
            parser_name = getattr(parsers, parser_override["parser"])

            if parser_name:
                self.parsers[parser_override["document_type"]] = parser_name()

        for doc_type, parser_infos in self.AVAILABLE_PARSERS.items():
            for parser_info in parser_infos:
                if (
                    doc_type not in excluded_parsers
                    and doc_type not in self.parsers
                ):
                    # will chose the first parser in the list
                    self.parsers[doc_type] = parser_info()

    @property
    def supported_types(self) -> list[str]:
        """
        Lists the data types supported by the pipe.
        """
        return [entry_type for entry_type in DocumentType]

    async def _parse(
        self,
        document: Document,
        run_id: uuid.UUID,
        version: str,
    ) -> AsyncGenerator[Union[R2RDocumentProcessingError, Extraction], None]:
        if document.type not in self.parsers:
            yield R2RDocumentProcessingError(
                document_id=document.id,
                error_message=f"Parser for {document.type} not found in `ParsingPipe`.",
            )
            return
        parser = self.parsers[document.type]
        texts = parser.ingest(document.data)
        extraction_type = ExtractionType.TXT
        t0 = time.time()
        if document.type in self.IMAGE_TYPES:
            extraction_type = ExtractionType.IMG
            document.metadata["image_type"] = document.type.value
            # SAVE IMAGE DATA
            # try:
            #     import base64
            #     sanitized_data = base64.b64encode(document.data).decode('utf-8')
            # except Exception as e:
            #     sanitized_data = document.data

            # document.metadata["image_data"] = sanitized_data
        elif document.type == DocumentType.MP4:
            extraction_type = ExtractionType.MOV
            document.metadata["audio_type"] = document.type.value

        iteration = 0
        async for text in texts:
            extraction_id = generate_id_from_label(
                f"{document.id}-{iteration}-{version}"
            )
            document.metadata["version"] = version
            extraction = Extraction(
                id=extraction_id,
                data=text,
                metadata=document.metadata,
                document_id=document.id,
                type=extraction_type,
            )
            yield extraction
            # TODO - Add settings to enable extraction logging
            # extraction_dict = extraction.dict()
            # await self.enqueue_log(
            #     run_id=run_id,
            #     key="extraction",
            #     value=json.dumps(
            #         {
            #             "data": extraction_dict["data"],
            #             "document_id": str(extraction_dict["document_id"]),
            #             "extraction_id": str(extraction_dict["id"]),
            #         }
            #     ),
            # )
            iteration += 1
        logger.debug(
            f"Parsed document with id={document.id}, title={document.metadata.get('title', None)}, user_id={document.metadata.get('user_id', None)}, metadata={document.metadata} into {iteration} extractions in t={time.time() - t0:.2f} seconds."
        )

    async def _run_logic(
        self,
        input: Input,
        state: AsyncState,
        run_id: uuid.UUID,
        versions: Optional[list[str]] = None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[Extraction, None]:
        parse_tasks = []

        iteration = 0
        async for document in input.message:
            version = versions[iteration] if versions else "v0"
            iteration += 1
            parse_tasks.append(
                self._handle_parse_task(document, version, run_id)
            )

        # Await all tasks and yield results concurrently
        for parse_task in asyncio.as_completed(parse_tasks):
            for extraction in await parse_task:
                yield extraction

    async def _handle_parse_task(
        self, document: Document, version: str, run_id: uuid.UUID
    ) -> AsyncGenerator[Extraction, None]:
        extractions = []
        async for extraction in self._parse(document, run_id, version):
            extractions.append(extraction)
        return extractions
