from io import BytesIO
import os, logging
import aspose.words as aw


ACCEPTED_FILE_EXTENSIONS = {
    ".doc" : aw.LoadFormat.DOC,
    ".docx" : aw.LoadFormat.DOCX,
    ".rtf" : aw.LoadFormat.RTF,
    ".html" : aw.LoadFormat.HTML,
    ".htm" : aw.LoadFormat.HTML,
    ".mhtml" : aw.LoadFormat.MHTML,
    ".mht" : aw.LoadFormat.MHTML,
    ".mobi" : aw.LoadFormat.MOBI,
    ".azw3" : aw.LoadFormat.AZW3,
    ".epub" : aw.LoadFormat.EPUB,
    ".odt" : aw.LoadFormat.ODT,
    ".txt" : aw.LoadFormat.TEXT,
    ".md" : aw.LoadFormat.MARKDOWN,
    ".pdf" : aw.LoadFormat.PDF,
    ".xml" : aw.LoadFormat.XML
}

ACCEPTED_MIME_TYPES = {
    "application/msword" : aw.LoadFormat.DOC,
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document" : aw.LoadFormat.DOCX,
    "application/rtf" : aw.LoadFormat.RTF,
    "text/rtf" : aw.LoadFormat.RTF,
    "text/richtext" : aw.LoadFormat.RTF,
    "text/html" : aw.LoadFormat.HTML,
    "application/xhtml" : aw.LoadFormat.HTML,
    "multipart/related" : aw.LoadFormat.MHTML,
    "application/x-mobipocket-ebook" : aw.LoadFormat.MOBI,
    "application/vnd.amazon.ebook" : aw.LoadFormat.AZW3,
    "application/epub+zip" : aw.LoadFormat.EPUB,
    "application/vnd.oasis.opendocument.text" : aw.LoadFormat.ODT,
    "text/plain" : aw.LoadFormat.TEXT,
    "text/markdown" : aw.LoadFormat.MARKDOWN,
    "application/pdf" : aw.LoadFormat.PDF,
    "application/xml" : aw.LoadFormat.XML
}

MARKDOWN_ENCODING = "utf-8"


class AsposeWordsConverter:

    def convert(self, document_data: BytesIO, document_format: aw.LoadFormat,
                filename: str = None) -> tuple[BytesIO, str]:
        """Converts the given document to a markdown and returns it along with the document title."""

        load_options = aw.loading.LoadOptions()
        load_options.load_format = document_format
        if filename is not None and filename != "": 
            load_options.base_uri = os.path.dirname(filename)
        load_options.resource_loading_callback = ResourceLoadingCallback()
        document = aw.Document(document_data, load_options)

        md_data = BytesIO()
        save_options = aw.saving.MarkdownSaveOptions()
        save_options.export_images_as_base64 = True
        save_options.office_math_export_mode = aw.saving.MarkdownOfficeMathExportMode.LATEX
        save_options.encoding = MARKDOWN_ENCODING
        document.save(stream=md_data, save_options=save_options)
        return (md_data, document.built_in_document_properties.title or "")


    def get_aspose_format(extension: str, mimetype: str) -> aw.LoadFormat:
        """Returns the Aspose.Words.LoadFormat for the given file extension and MIME type, or
         aw.LoadFormat.UNKNOWN if it can't be determined."""

        if extension in ACCEPTED_FILE_EXTENSIONS:
            return ACCEPTED_FILE_EXTENSIONS[extension]

        if mimetype is not None:
            for minetype_prefix, load_format in ACCEPTED_MIME_TYPES.items():
                if mimetype.startswith(minetype_prefix):
                    return load_format

        return aw.LoadFormat.UNKNOWN

class ResourceResolver:
    """ Helper class to resolve resources (e.g., images) during documents loading. 
        Note: This class is not thread-safe! Only one instance should be used at a time. """
    
    _callback = None

    def __init__(self, callback):
        """ Initializes the ResourceResolver with a callback function.
            The callback function should accept a single string argument (the resource URI)
            and return either: bytes (the resource data), str (a new URI to load the resource from),
            or None (to skip the resource). """

        if callback is None:
            raise ValueError("callback cannot be None")
        if not callable(callback):
            raise ValueError("callback must be callable")
        if ResourceResolver._callback is not None:
            raise RuntimeError("Only one ResourceResolver instance can be active at a time.")
        ResourceResolver._callback = callback


    def __enter__(self):
        return self


    def __exit__(self, exc_type, exc_value, traceback):
        """ Cleans up the ResourceResolver by removing the callback. """
        self.close()

    
    def close(self):
        ResourceResolver._callback = None


class ResourceLoadingCallback(aw.loading.IResourceLoadingCallback):
    def resource_loading(self, args: aw.loading.ResourceLoadingArgs) -> aw.loading.ResourceLoadingAction:
        if ResourceResolver._callback is not None:
            result = ResourceResolver._callback(args.uri)
            if result is None:
                return aw.loading.ResourceLoadingAction.SKIP
            if isinstance(result, bytes):
                args.set_data(result)
                return aw.loading.ResourceLoadingAction.USER_PROVIDED
            if isinstance(result, str):
                args.uri = result
                return aw.loading.ResourceLoadingAction.USER_PROVIDED
            
        return aw.loading.ResourceLoadingAction.DEFAULT
    