import os
import pandas as pd
import re
from tqdm import tqdm
from collections import defaultdict
from langchain_text_splitters import RecursiveCharacterTextSplitter


def merge_small_chunks(chunks, min_size):
    """
    合并小于指定大小的文本块,反复执行直到不再出现小于min_size的块

    参数:
        chunks: 文本块列表
        min_size: 最小块大小阈值

    返回:
        合并后的文本块列表

    逻辑:
        - 遍历所有块,如果某块小于min_size,则与下一个块合并
        - 如果是最后一个块且小于min_size,则与上一个块合并
        - 反复执行上述过程,直到不再出现任何小于min_size的块
    """
    if not chunks or min_size <= 0:
        return chunks

    # 反复合并,直到不再出现小块
    has_small_chunks = True
    merged_chunks = [chunk.page_content for chunk in chunks]
    max_iterations = 100  # 防止无限循环的最大迭代次数
    iteration_count = 0

    while has_small_chunks:
        iteration_count += 1

        # 防止无限循环
        if iteration_count > max_iterations:
            print(f"警告: 合并迭代超过{max_iterations}次,停止合并。可能存在无法满足min_size的块。")
            break

        has_small_chunks = False
        new_merged_chunks = []
        i = 0
        previous_chunks_count = len(merged_chunks)

        while i < len(merged_chunks):
            current_chunk = merged_chunks[i]
            current_length = len(current_chunk)

            # 如果当前块大小足够,直接添加
            if current_length >= min_size:
                new_merged_chunks.append(current_chunk)
                i += 1
            # 如果当前块太小
            else:
                # 如果不是最后一个块,与下一个块合并
                if i < len(merged_chunks) - 1:
                    next_chunk = merged_chunks[i + 1]
                    merged_chunk = current_chunk + next_chunk
                    new_merged_chunks.append(merged_chunk)
                    has_small_chunks = True  # 标记发现小块并进行了合并
                    i += 2  # 跳过下一个块,因为已经合并了
                # 如果是最后一个块,与上一个块合并
                else:
                    if new_merged_chunks:
                        # 将当前小块合并到上一个块
                        new_merged_chunks[-1] = new_merged_chunks[-1] + current_chunk
                        # 注意:这里不设置has_small_chunks=True,因为合并后不会产生新的独立小块
                    else:
                        # 如果没有上一个块(即只有一个块且很小),直接添加并停止循环
                        new_merged_chunks.append(current_chunk)
                        print(f"警告: 文件整体内容长度({current_length})小于最小块大小({min_size}),保留原样。")
                    i += 1

        merged_chunks = new_merged_chunks

        # 额外的终止条件:如果块数量不再减少,说明无法继续合并
        if len(merged_chunks) == previous_chunks_count:
            # 检查是否还有小块
            small_chunks_exist = any(len(chunk) < min_size for chunk in merged_chunks)
            if small_chunks_exist:
                small_count = sum(1 for chunk in merged_chunks if len(chunk) < min_size)
                print(f"警告: 仍有{small_count}个块小于最小块大小({min_size}),但无法继续合并。")
            break

    return merged_chunks


def process_md_to_csv(input_folder, output_folder, chunk_size=350, min_chunk_size=100):
    """
    将 Markdown/文本文件处理为 CSV 文件

    参数:
        input_folder: 输入文件夹路径
        output_folder: 输出文件夹路径
        chunk_size: 文本切分大小
        min_chunk_size: 最小块大小,小于此值的块将被合并

    逻辑:
        - 读取所有 .md 和 .txt 文件
        - 使用 RecursiveCharacterTextSplitter 切分文本
        - 合并过小的块
        - 提取图片URL、处理Markdown替换
        - 生成 Position 层级结构
        - 输出包含 Text, Text_pure, Img_url, Position 列的 CSV
    """
    # 创建输出文件夹(如果不存在)
    os.makedirs(output_folder, exist_ok=True)

    # 定义文本切分器
    text_splitter = RecursiveCharacterTextSplitter(
        separators=[
            "\n\n",
            "\n",
            " ",
            ".", # 半角句读 大头使用
            ",",
            "،", # 阿拉伯文逗号
            "،", # 全角句读
            "。",
            "、", # 日文逗号
            "।", # 印地语句号
            "\u200b",  # Zero-width space
            "\uff0c",  # Fullwidth comma
            "\u3001",  # Ideographic comma
            "\uff0e",  # Fullwidth full stop
            "\u3002",  # Ideographic full stop
            "",
        ],
        chunk_size=chunk_size,
        chunk_overlap=0,
        length_function=len,
        is_separator_regex=False,
    )

    # 获取输入文件夹中的所有.md和.txt文件
    files = [f for f in os.listdir(input_folder) if f.endswith(".md") or f.endswith(".txt")]

    # 字典用于存储替换信息
    replacement_dict = defaultdict(list)

    def replace_img(match, img_counter, replacement_dict):
        key = f"<img{img_counter}>"
        replacement_dict[key] = match.group(0)
        return key, img_counter + 1, match.group(1)

    def replace_text(match, text_counter, replacement_dict):
        text = match.group(1)
        text_counter[text] += 1
        key = f"<{text}{text_counter[text] if text_counter[text] > 1 else ''}>"
        replacement_dict[key] = match.group(0)
        return key

    # 定义Markdown语法及其优先级
    markdown_priority = {
        r'^# ': 0,
        r'^## ': 1,
        r'^### ': 2,
        r'^#### ': 3,
        r'^##### ': 4,
        r'^###### ': 5,
        r'^\d+\. ': 6,
        r'^- ': 6,
        r'^\* ': 6,
        r'^\*\*.*\*\*$': 6
    }

    # 特殊Markdown语法检测
    special_syntax = {
        r'^\s*=+\s*$': 0,  # 一级标题
        r'^\s*-+\s*$': 1   # 二级标题
    }

    # 遍历文件并显示进度条
    for filename in tqdm(files, desc="Processing files"):
        input_filepath = os.path.join(input_folder, filename)

        # 读取文件内容
        with open(input_filepath, encoding='utf-8-sig') as f:
            file_content = f.read()

        # 预处理步骤
        # 替换连续超过两个的*
        file_content = re.sub(r'\*{3,}', '**', file_content)
        # 替换连续超过三个的- 或 =
        file_content = re.sub(r'\-{4,}', '---', file_content)
        file_content = re.sub(r'\={4,}', '===', file_content)

        img_counter = 1
        text_counter = defaultdict(int)
        img_urls_per_line = []
        replacement_dict = {}

        # 替换 ![img](url)
        img_pattern = re.compile(r'!\[.*?\]\((.*?)\)')
        img_matches = list(img_pattern.finditer(file_content))
        for match in img_matches:
            replacement, img_counter, img_url = replace_img(match, img_counter, replacement_dict)
            file_content = file_content.replace(match.group(0), replacement)
            img_urls_per_line.append(img_url)

        # 替换 [text](url)
        text_pattern = re.compile(r'\[(.*?)\]\(.*?\)')
        text_matches = list(text_pattern.finditer(file_content))
        for match in text_matches:
            replacement = replace_text(match, text_counter, replacement_dict)
            file_content = file_content.replace(match.group(0), replacement)

        # 执行文本切分
        texts = text_splitter.create_documents([file_content])

        # 合并过小的块
        merged_texts = merge_small_chunks(texts, min_chunk_size)

        # 创建DataFrame
        df = pd.DataFrame(merged_texts, columns=['Text'])

        # 还原替换的 <img%d> 和 <text%d>
        def restore_replacement(text):
            for key, value in replacement_dict.items():
                text = text.replace(key, value)
            return text

        # 创建 Text_pure 列
        df['Text_pure'] = df['Text']

        def extract_img_urls(text):
            img_keys = re.findall(r'<img\d+>', text)
            img_urls = []
            for key in img_keys:
                if key in replacement_dict:
                    url_match = re.search(r'\((.*?)\)', replacement_dict[key])
                    if url_match:
                        img_urls.append(url_match.group(1))
            return ';'.join(img_urls)

        df['Img_url'] = df['Text_pure'].apply(extract_img_urls)
        # 还原 Text 列中的内容
        df['Text'] = df['Text'].apply(restore_replacement)

        # 保存为CSV文件
        output_filename = f"{os.path.splitext(filename)[0]}.csv"
        output_filepath = os.path.join(output_folder, output_filename)
        df.to_csv(output_filepath, index=False, encoding='utf-8-sig')

        print(f"Processed and saved: {output_filepath}")

        # 读取CSV文件内容
        df = pd.read_csv(output_filepath, encoding='utf-8-sig')

        # 初始化栈并将文件名作为初始位置
        initial_position = os.path.splitext(filename)[0]
        position_stack = [initial_position]
        priority_stack = [-1]  # 初始优先级为-1,以确保文件名始终在最顶部

        # 定义位置列
        df['Position'] = ""

        # 用于记录连续相同位置的次数
        last_position = ""
        position_counter = 0

        # 逐行处理文本
        for idx, row in df.iterrows():
            text = row['Text']
            sentences = re.split(r'\n+', text)  # 以换行符进行分句

            # 首先对每一行赋初始位置
            current_position = ' > '.join(position_stack)

            # 检查是否与上一行位置相同
            if current_position == last_position:
                position_counter += 1
                if position_counter == 2:  # 在第二次遇到时给上一个添加序号
                    df.at[idx - 1, 'Position'] += f" (part 1)"
                current_position += f" (part {position_counter})"
            else:
                last_position = current_position
                position_counter = 1

            df.at[idx, 'Position'] = current_position

            # 逐句处理文本
            previous_sentence = ""
            for sentence in sentences:
                current_priority = None

                # 判断当前句子的标题信息及其优先级
                for grammar, priority in markdown_priority.items():
                    if re.match(grammar, sentence.strip()):
                        current_priority = priority
                        break

                # 特殊处理 `===` 和 `---` 的情况
                for special_grammar, special_priority in special_syntax.items():
                    if re.match(special_grammar, sentence.strip()):
                        current_priority = special_priority
                        sentence = previous_sentence
                        break

                # 更新位置栈和优先级栈
                if current_priority is not None:
                    while priority_stack and priority_stack[-1] >= current_priority:
                        position_stack.pop()
                        priority_stack.pop()

                    position_stack.append(sentence.strip())
                    priority_stack.append(current_priority)

                previous_sentence = sentence

        # 保存更新后的CSV文件
        df.to_csv(output_filepath, index=False, encoding='utf-8-sig')

        print(f"Refined and saved: {output_filepath}")


if __name__ == "__main__":
    # 定义输入和输出文件夹
    input_folder = "./input/GT_ML_TEST/"
    output_folder = "./output/GT_ML_TEST/"

    # 执行处理
    process_md_to_csv(input_folder, output_folder)