from marker.convert import convert_single_pdf
from marker.models import load_all_models
import tempfile
from indexify_extractor_sdk import Content, Extractor, Feature

from pydantic import BaseModel
from typing import Optional, Literal, List, Union

class MarkdownExtractorConfig(BaseModel):
    max_pages: Optional[int] = None
    langs: Optional[str] = None
    batch_multiplier: Optional[int] = 2

class MarkdownExtractor(Extractor):
    name = "tensorlake/marker"
    description = "Markdown Extractor for PDFs"
    system_dependencies = []
    input_mime_types = ["application/pdf"]

    def __init__(self):
        super(MarkdownExtractor, self).__init__()
        self.model_lst = load_all_models()

    def extract(self, content: Content, params: MarkdownExtractorConfig) -> List[Union[Feature, Content]]:
        contents = []
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as inputtmpfile:
            inputtmpfile.write(content.data)
            inputtmpfile.flush()

            full_text, images, out_meta = convert_single_pdf(inputtmpfile.name, self.model_lst, max_pages=params.max_pages, langs=params.langs, batch_multiplier=params.batch_multiplier)
            
            feature = Feature.metadata(value=out_meta, name="text")
            contents.append(Content.from_text(full_text, features=[feature]))

        return contents