import copy
import json as json_parse
import os
import re

import click
from loguru import logger

import magic_pdf.model as model_config
from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_line_sort_bbox,
                                      draw_model_bbox, draw_span_bbox)
from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
from magic_pdf.pipe.OCRPipe import OCRPipe
from magic_pdf.pipe.TXTPipe import TXTPipe
from magic_pdf.pipe.UNIPipe import UNIPipe
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
from magic_pdf.post_proc.remove_spaces_html import remove_extra_spaces_html_txt

def sanitize_filename(filename: str) -> str:
    return re.sub(r'[^a-zA-Z0-9_\u4e00-\u9fff-]', '_', filename)

def prepare_env(output_dir, pdf_file_name, method):
    # logger.info(f'pdf_file_name:{pdf_file_name}')
    pdf_file_name = sanitize_filename(pdf_file_name)
    local_parent_dir = os.path.join(output_dir, pdf_file_name)

    local_image_dir = os.path.join(str(local_parent_dir), 'images')
    local_md_dir = local_parent_dir
    os.makedirs(local_image_dir, exist_ok=True)
    os.makedirs(local_md_dir, exist_ok=True)
    return local_image_dir, local_md_dir

def remove_empty_lines_from_file(file_path):
    # 读取文件内容
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()

    # 删除空行
    non_empty_lines = [line for line in lines if line.strip()]

    # 将非空行写回原文件
    with open(file_path, 'w', encoding='utf-8') as file:
        file.writelines(non_empty_lines)


def do_parse(
    ocr_status,
    config_path,
    output_dir,
    pdf_file_name,
    pdf_bytes,
    model_list,
    parse_method,
    debug_able,
    model,
    f_draw_span_bbox=True,
    f_draw_layout_bbox=True,
    f_dump_md=True,
    f_dump_middle_json=True,
    f_dump_model_json=True,
    f_dump_orig_pdf=True,
    f_dump_content_list=True,
    f_make_md_mode=MakeMode.MM_MD,
    f_draw_model_bbox=False,
    f_draw_line_sort_bbox=False,
    start_page_id=0,
    end_page_id=None,
    lang=None,
    layout_model=None,
    formula_enable=None,
    table_enable=None,
):
    if debug_able:
        logger.warning('debug mode is on')
        f_draw_model_bbox = True
        f_draw_line_sort_bbox = True

    orig_model_list = copy.deepcopy(model_list)
    local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name,
                                                parse_method)

    image_writer, md_writer = DiskReaderWriter(
        local_image_dir), DiskReaderWriter(local_md_dir)
    image_dir = str(os.path.basename(local_image_dir))

    if parse_method == 'auto':
        jso_useful_key = {'_pdf_type': '', 'model_list': model_list}
        pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True,
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang,
                       layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable)
    elif parse_method == 'txt':
        pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True,
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang,
                       layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable)
    elif parse_method == 'ocr':
        pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True,
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang,
                       layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable)
    else:
        logger.error('unknown parse method')
        exit(1)

    pipe.pipe_classify()

    if len(model_list) == 0:
        if model_config.__use_inside_model__:
            pipe.pipe_analyze(model)
            orig_model_list = copy.deepcopy(pipe.model_list)
        else:
            logger.error('need model list input')
            exit(2)

    pipe.pipe_parse(ocr_status,config_path,local_image_dir)
    pdf_info = pipe.pdf_mid_data['pdf_info']
    if f_draw_layout_bbox:
        draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)

    md_content = pipe.pipe_mk_markdown(image_dir,
                                       drop_mode=DropMode.NONE,
                                       md_make_mode=f_make_md_mode)

    try:
        pdf_file_name = sanitize_filename(pdf_file_name)
        txt_file = f'{pdf_file_name}.txt'
        md_writer.write(
            content=md_content,
            path=txt_file,
            mode=AbsReaderWriter.MODE_TXT,
        )
        md_writer.write(
            f'{pdf_file_name}.md',
            md_content
        )
        filepath = os.path.join(str(local_md_dir), f'{pdf_file_name}.txt')
        # logger.info(f'txt文件保存在filepath:{filepath}')
        remove_empty_lines_from_file(filepath)

        if f_dump_orig_pdf:
            md_writer.write(
                content=pdf_bytes,
                path=f'{pdf_file_name}_origin.pdf',
                mode=AbsReaderWriter.MODE_BIN,
            )
        remove_extra_spaces_html_txt(filepath)


        return filepath
    except Exception as e:
        logger.error(f'{pdf_file_name}导出txt文件失败，具体原因为：\n{e}')
        return None

parse_pdf_methods = click.Choice(['ocr', 'txt', 'auto'])
