from loguru import logger
import sys
from indoxArcg.core import Document
from indoxArcg.vector_stores.utils import filter_complex_metadata

# Set up logging
logger.remove()  # Remove the default logger
logger.add(
    sys.stdout, format="<green>{level}</green>: <level>{message}</level>", level="INFO"
)

logger.add(
    sys.stdout, format="<red>{level}</red>: <level>{message}</level>", level="ERROR"
)


class DoclingReader:
    def __init__(self, file_path: str):
        from docling.document_converter import DocumentConverter

        self.result = None
        try:
            self.converter = DocumentConverter()
            self.file_path = file_path
        except Exception as e:
            logger.error(f"Error initializing DocumentConverter: {e}")
            raise

    def load(self, max_num_pages=None, max_file_size=None):
        kwargs = {}

        if max_num_pages is not None:
            kwargs["max_num_pages"] = max_num_pages

        if max_file_size is not None:
            kwargs["max_file_size"] = max_file_size

        self.result = self.converter.convert(source=self.file_path, **kwargs)

        # if output_format == "txt":
        #     return self.result.document.export_to_text
        # elif output_format == "markdown":
        #     return self.result.document.export_to_markdown
        # elif output_format == "html":
        #     return self.result.document.export_to_html
        return self.result

    def load_and_split(self, max_tokens=512, **kwargs):
        from docling.chunking import HybridChunker

        result = self.load()
        chunker = HybridChunker()
        chunk_iter = chunker.chunk(dl_doc=result.document)

        documents = []  # Initialize a list to store Document objects
        for i, chunk in enumerate(chunk_iter):
            # Serialize the chunk to retain enriched text representation
            enriched_text = chunker.serialize(chunk=chunk)

            # Create a Document object with page_content and metadata
            doc = Document(
                page_content=chunk.text,
                metadata={
                    "chunk_index": i,
                    "enriched_text": enriched_text,
                },
            )
            documents.append(doc)
        documents = filter_complex_metadata(documents=documents)
        return documents
