import json
from dataflow.core import OperatorABC
from dataflow.utils.registry import OPERATOR_REGISTRY
from dataflow.utils.storage import DataFlowStorage

from pathlib import Path

@OPERATOR_REGISTRY.register()
class MinerU2LLMInputOperator(OperatorABC):
    def __init__(self):
        pass
    
    @staticmethod
    def get_desc(lang: str = "zh") -> str:
        if lang == 'zh':
            return (
                "MinerU格式转换为LLM输入格式算子。"
                "将MinerU生成的内容列表JSON文件转换为适合LLM处理的格式，"
                "包括展平列表项并重新编号。"
            )
        else:
            return (
                "Convert MinerU format to LLM input format operator."
                "Transforms the content list JSON file generated by MinerU into a format suitable for LLM processing,"
                "including flattening list items and re-indexing."
            )

    def _convert_json(self, input_file, output_file):
        with open(input_file, 'r', encoding="utf-8") as infile:
            data = list(json.load(infile))
        
        new_data = []
        id = 0
        for item in data:
            item['id'] = id
            item.pop('bbox', None)
            item.pop('page_idx', None)
            if item.get('type','') == 'list':
                if item['sub_type'] == 'text':
                    for idx, list_item in enumerate(item.get('list_items', [])):
                        new_item = {
                            'type': 'text',
                            'text': list_item,
                            'id': id + idx,
                        }
                        new_data.append(new_item)
                    id += len(item.get('list_items', []))
            else:
                new_data.append(item)
                id += 1
        
        with open(output_file, 'w', encoding='utf-8') as outfile:
            json.dump(new_data, outfile, ensure_ascii=False)
    
    def run(self, storage: DataFlowStorage,
            input_markdown_path_key,
            output_converted_layout_key,
            ):
        dataframe = storage.read("dataframe")
    
        for index, row in dataframe.iterrows():
            md_path = Path(row[input_markdown_path_key])
            try:
                input_json_path = list(md_path.parent.glob("*_content_list.json"))[0]
            except Exception:
                raise ValueError("No _content_list.json file found in the api result. There might be an error with the Mineru api.")
            
            converted_path = str(input_json_path).replace('.json', '_converted.json')
            self._convert_json(input_json_path, converted_path)
            dataframe.at[index, output_converted_layout_key] = converted_path
            
            with open(converted_path, 'r', encoding='utf-8') as infile:
                data = json.load(infile)
                assert isinstance(data, list), f"Expected list, got {type(data)} for {input_json_path}"
                
        storage.write(dataframe)