import os
import sys
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod

def convert_pdf(pdf_file_path):
    # Get filename and prepare output paths
    pdf_file_name = os.path.basename(pdf_file_path)
    name_without_suff = os.path.splitext(pdf_file_name)[0]
    
    # prepare env
    local_image_dir, local_md_dir = "output/images", "output"
    image_dir = str(os.path.basename(local_image_dir))
    
    os.makedirs(local_image_dir, exist_ok=True)
    os.makedirs(local_md_dir, exist_ok=True)
    
    image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(
        local_md_dir
    )
    
    # read bytes
    reader1 = FileBasedDataReader(os.path.dirname(pdf_file_path))
    pdf_bytes = reader1.read(pdf_file_name)  # read the pdf content
    
    print(f"Processing PDF: {pdf_file_path}")
    
    # proc
    ## Create Dataset Instance
    ds = PymuDocDataset(pdf_bytes)
    
    ## inference
    if ds.classify() == SupportedPdfParseMethod.OCR:
        infer_result = ds.apply(doc_analyze, ocr=True)
    
        ## pipeline
        pipe_result = infer_result.pipe_ocr_mode(image_writer)
    
    else:
        infer_result = ds.apply(doc_analyze, ocr=False)
    
        ## pipeline
        pipe_result = infer_result.pipe_txt_mode(image_writer)
    
    ### draw model result on each page
    infer_result.draw_model(os.path.join(local_md_dir, f"{name_without_suff}_model.pdf"))
    
    ### get model inference result
    model_inference_result = infer_result.get_infer_res()
    
    ### draw layout result on each page
    pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf"))
    
    ### draw spans result on each page
    pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_suff}_spans.pdf"))
    
    ### get markdown content
    md_content = pipe_result.get_markdown(image_dir)
    
    ### dump markdown
    md_file_path = f"{name_without_suff}.md"
    pipe_result.dump_md(md_writer, md_file_path, image_dir)
    print(f"Markdown saved to: {os.path.join(local_md_dir, md_file_path)}")
    
    ### get content list content
    content_list_content = pipe_result.get_content_list(image_dir)
    
    ### dump content list
    pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir)
    
    ### get middle json
    middle_json_content = pipe_result.get_middle_json()
    
    ### dump middle json
    pipe_result.dump_middle_json(md_writer, f'{name_without_suff}_middle.json')
    
    return md_content

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python convert_pdf.py <pdf_path>")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    
    if not os.path.exists(pdf_path):
        print(f"Error: PDF file not found at {pdf_path}")
        sys.exit(1)
        
    convert_pdf(pdf_path)