import time, json
from typing import Union, List
from pathlib import Path

try:
    from azure.core.credentials import AzureKeyCredential
    from azure.ai.documentintelligence.aio import DocumentIntelligenceClient
    from azure.ai.documentintelligence.models import AnalyzeDocumentRequest, AnalyzeResult
    try:
        from azure.ai.documentintelligence.models import ContentFormat
    except ImportError:
        ContentFormat = None
except ImportError:
    AzureKeyCredential = None
    DocumentIntelligenceClient = None
    AnalyzeDocumentRequest = None
    ContentFormat = None
    AnalyzeResult = None

from langchain_text_splitters import MarkdownHeaderTextSplitter

from animuz_core.utils import init_logger

LOGGER = init_logger(__name__)

class AzureDocAiClient():
    def __init__(self, endpoint: str, key: str) -> None:
        self.client = DocumentIntelligenceClient(endpoint=endpoint, credential=AzureKeyCredential(key))

    async def analyze_document(self, file_dir: str, do_split: bool = True) -> Union[List[dict], str]:
        """
        Analyzes a document from the given file path.

        Args:
            file_dir (str): The path to the document file.
            do_split (bool; default=True): True if the document should be split with langchain markdown splitter.

        Returns:
            Union[List[dict], str]: If do_split is True: a list of dictionaries containing the text and metadata of each section of the document.
                                    If do_split is False: a string of the content of the analyzed document in markdown format.
        """
        start = time.time()
        with open(file_dir, "rb") as f:
            poller = await self.client.begin_analyze_document("prebuilt-layout", f, content_type="application/octet-stream", output_content_format=ContentFormat.MARKDOWN)
            
        result: AnalyzeResult = await poller.result()
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Azure Doc AI latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))

        if not do_split:
            return result.content

        headers_to_split_on = [
            ("#", "Header 1"),
            ("##", "Header 2"),
        ]
        splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on,  strip_headers=False)
        splitted = splitter.split_text(result.content)
        res = []
        for s in splitted:
            res.append({"text": s.page_content, 
                        "source": "azure_doc_ai",
                        "file_directory": str(Path(file_dir).parent),
                        "metadata": {
                            "filename": file_dir,
                            }
                        })
            
        return res
    
if __name__=="__main__":
    import asyncio
    import dotenv, os
    dotenv.load_dotenv()

    client = AzureDocAiClient(endpoint=os.getenv("AZURE_DOC_AI_ENDPOINT"), key=os.getenv("AZURE_DOC_AI_KEY"))
    asyncio.run(client.analyze_document("test_docs/test.pdf"))