"""
文本分块辅助工具

这里收纳了与分块元数据、层级处理等相关的工具类的精简版，
用于 PDF 处理流程，参考原 doc_worker.chunking 的实现，但保持独立。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List

import re


@dataclass
class ChunkMetaBuilder:
    """生成 chunk 元数据的简单工具。"""

    def build_chunk_metadata(
        self,
        sub_chunk: Dict[str, Any],
        title_info: Dict[str, Any],
        start_pos: int,
        pages_content: List[Dict[str, Any]],
        i: int,
        j: int,
    ) -> Dict[str, Any]:
        """
        根据子块与标题信息构造一个标准 chunk。

        为简化，只保留与现有输出 JSON 兼容的关键字段：
        - file_name
        - chunk_id
        - section_label
        - project_plan_type (由分类 category_code 决定)
        - element_tag.page / chunk_id / serial_number
        - review_chunk_content
        """
        content = sub_chunk["content"]
        page = self._get_page_from_pos(start_pos + sub_chunk["relative_start"], pages_content)

        level1_title = title_info.get("hierarchy_path", [title_info["title"]])[0]
        section_label = level1_title

        chunk_id = f"doc_chunk_{level1_title}_{j}"
        project_plan_type = title_info.get("category_code", "other")

        return {
            "file_name": "",  # 由上层填充
            "chunk_id": chunk_id,
            "section_label": section_label,
            "project_plan_type": project_plan_type,
            "element_tag": {
                "chunk_id": chunk_id,
                "page": page,
                "serial_number": level1_title,
            },
            "review_chunk_content": content,
            "_sort_key": start_pos + sub_chunk["relative_start"],
        }

    def _get_page_from_pos(self, pos: int, pages_content: List[Dict[str, Any]]) -> int:
        for page in pages_content:
            if page["start_pos"] <= pos < page["end_pos"]:
                return int(page["page_num"])
        return 1


class SimpleChunkSplitter:
    """
    非完整实现的分块器：
    - 暂不做子标题进一步拆分；
    - 仅基于 max_chunk_size / min_chunk_size 对正文文本做简单切分。
    """

    def split_by_text_length(
        self,
        content_block: str,
        max_chunk_size: int,
        min_chunk_size: int,
    ) -> List[Dict[str, Any]]:
        chunks: List[Dict[str, Any]] = []
        start = 0
        length = len(content_block)
        while start < length:
            end = min(start + max_chunk_size, length)
            # 简单按句号/换行回退，尽量不截断句子
            slice_text = content_block[start:end]
            split_pos = self._find_natural_split(slice_text, min_size=min_chunk_size)
            end = start + split_pos
            chunks.append(
                {
                    "relative_start": start,
                    "relative_end": end,
                    "content": content_block[start:end],
                }
            )
            start = end
        return chunks

    def _find_natural_split(self, text: str, min_size: int) -> int:
        if len(text) <= min_size:
            return len(text)
        # 从后往前找较自然的断点
        for ch in ["\n", "。", "！", "？", ".", "!", "?"]:
            idx = text.rfind(ch, min_size)
            if idx != -1:
                return idx + 1
        return len(text)


class HierarchicalChunkMixin:
    """
    分级目录切分的通用工具 Mixin。

    把原先 `PdfTextSplitter` 中完全相同的
    chunk 元数据构造、层级路径、编号提取等方法抽到这里，
    便于多种 worker 复用。
    """

    def _split_large_chunk(
        self,
        content: str,
        max_chunk_size: int,
        title: str,
        hierarchy_path: List[str] | None = None,
    ) -> List[Dict[str, Any]]:
        """
        将超大块按句子级分割（保持语义完整）。
        """
        sentences = re.split(r"([。！？\n])", content)

        combined_sentences = []
        for i in range(0, len(sentences) - 1, 2):
            if i + 1 < len(sentences):
                combined_sentences.append(sentences[i] + sentences[i + 1])
            else:
                combined_sentences.append(sentences[i])

        if not combined_sentences:
            combined_sentences = [content]

        chunks: List[Dict[str, Any]] = []
        current_chunk = ""
        current_start = 0

        for sentence in combined_sentences:
            if len(current_chunk) + len(sentence) <= max_chunk_size:
                current_chunk += sentence
            else:
                if current_chunk:
                    chunk_data: Dict[str, Any] = {
                        "content": current_chunk,
                        "relative_start": current_start,
                        "is_split": True,
                    }
                    if hierarchy_path is not None:
                        chunk_data["hierarchy_path"] = hierarchy_path
                    chunks.append(chunk_data)
                    current_start += len(current_chunk)
                current_chunk = sentence

        if current_chunk:
            chunk_data = {
                "content": current_chunk,
                "relative_start": current_start,
                "is_split": True,
            }
            if hierarchy_path is not None:
                chunk_data["hierarchy_path"] = hierarchy_path
            chunks.append(chunk_data)

        return chunks

    def _build_hierarchy_path_for_subtitle(
        self,
        sub_title_item: Dict[str, Any],
        all_toc_items: List[Dict[str, Any]],
        parent_title_info: Dict[str, Any],
    ) -> List[str]:
        """为子标题构建完整的层级路径。"""
        hierarchy_path: List[str] = []

        sub_title = sub_title_item.get("title", "")
        sub_title_idx = -1
        for idx, item in enumerate(all_toc_items):
            if item.get("title", "") == sub_title:
                sub_title_idx = idx
                break

        if sub_title_idx < 0:
            return [parent_title_info["title"], sub_title]

        level_paths: Dict[int, str] = {}
        current_level = sub_title_item.get("level", 2)

        for i in range(sub_title_idx, -1, -1):
            item = all_toc_items[i]
            item_level = item.get("level", 1)

            if item_level <= current_level and item_level not in level_paths:
                level_paths[item_level] = item["title"]
                if item_level == 1:
                    break

        for level in range(1, current_level + 1):
            if level in level_paths:
                hierarchy_path.append(level_paths[level])

        if not hierarchy_path:
            hierarchy_path = [parent_title_info["title"], sub_title]

        return hierarchy_path

    def _build_hierarchy_path(
        self, title: str, all_toc_items: List[Dict[str, Any]], target_level: int
    ) -> List[str]:
        """构建从1级到当前标题的完整层级路径。"""
        hierarchy_path: List[str] = []

        current_item: Dict[str, Any] | None = None
        current_idx = -1
        for idx, item in enumerate(all_toc_items):
            if item["title"] == title:
                current_item = item
                current_idx = idx
                break

        if not current_item:
            return [title]

        current_level = current_item.get("level", target_level)
        level_paths: Dict[int, str] = {}

        for i in range(current_idx, -1, -1):
            item = all_toc_items[i]
            item_level = item.get("level", 1)

            if item_level <= current_level and item_level not in level_paths:
                level_paths[item_level] = item["title"]
                if item_level == 1:
                    break

        for level in range(1, current_level + 1):
            if level in level_paths:
                hierarchy_path.append(level_paths[level])
            elif level == current_level:
                hierarchy_path.append(title)

        if not hierarchy_path:
            hierarchy_path = [title]

        return hierarchy_path

    def _build_chunk_metadata(
        self,
        sub_chunk: Dict[str, Any],
        title_info: Dict[str, Any],
        start_pos: int,
        pages_content: List[Dict[str, Any]],
        i: int,
        j: int,
        chapter_classification_map: Dict[str, Dict[str, Any]] | None = None,
    ) -> Dict[str, Any]:
        """构建文本块的元数据。"""
        content = sub_chunk["content"]
        chunk_start_pos = start_pos + sub_chunk["relative_start"]
        page_num = self._get_page_from_pos(chunk_start_pos, pages_content)

        hierarchy_path = sub_chunk.get("hierarchy_path", [])
        sub_title = sub_chunk.get("sub_title", "")

        if hierarchy_path:
            section_label = "->".join(hierarchy_path)
        elif sub_title:
            section_label = f"{title_info['title']}->{sub_title}"
        else:
            section_label = title_info["title"]

        if hierarchy_path:
            lowest_title = hierarchy_path[-1]
            title_number = self._extract_title_number(lowest_title)
        elif sub_title:
            title_number = self._extract_title_number(sub_title)
        else:
            title_number = self._extract_title_number(title_info["title"])

        chunk_id_str = f"doc_chunk_{title_number}_{j}" if title_number else f"doc_chunk_{j}"

        chapter_classification = None
        if chapter_classification_map:
            if hierarchy_path and len(hierarchy_path) > 0:
                chapter_title = hierarchy_path[0]
                chapter_classification = chapter_classification_map.get(chapter_title)
            elif not hierarchy_path:
                chapter_title = title_info.get("title", "")
                chapter_classification = chapter_classification_map.get(chapter_title)

        chunk_data: Dict[str, Any] = {
            "file_name": "",
            "chunk_id": chunk_id_str,
            "section_label": section_label,
            "project_plan_type": title_info.get("category_code", "other"),
            "chapter_classification": title_info.get("category_code", "other"),
            "element_tag": {
                "chunk_id": chunk_id_str,
                "page": page_num,
                "serial_number": title_number if title_number else str(i + 1),
            },
            "review_chunk_content": content,
            "_title_number": title_number,
            "_local_index": j,
            "_sort_key": chunk_start_pos,
        }

        # if chapter_classification:
        #     chunk_data["chapter_classification"] = chapter_classification

        return chunk_data

    def _finalize_chunk_ids(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """生成最终的chunk_id和serial_number。"""
        final_chunks: List[Dict[str, Any]] = []
        section_groups: Dict[str, int] = {}

        for chunk in chunks:
            section_label = chunk.get("section_label", "")

            if section_label not in section_groups:
                section_groups[section_label] = 1
            else:
                section_groups[section_label] += 1

            local_index = section_groups[section_label]
            title_number_path = self._extract_title_number_path(section_label)

            if title_number_path:
                chunk_id_str = f"doc_chunk_{title_number_path}_{local_index}"
            else:
                chunk_id_str = f"doc_chunk_{local_index}"

            serial_number = self._extract_number_from_section_label(section_label)

            final_chunk = {
                "file_name": chunk["file_name"],
                "chunk_id": chunk_id_str,
                "section_label": chunk["section_label"],
                "project_plan_type": chunk["project_plan_type"],
                "chapter_classification": chunk["chapter_classification"],
                "element_tag": {
                    "chunk_id": chunk_id_str,
                    "page": chunk["element_tag"]["page"],
                    "serial_number": serial_number,
                },
                "review_chunk_content": chunk["review_chunk_content"],
            }

            final_chunks.append(final_chunk)

        return final_chunks

    def _get_page_from_pos(self, pos: int, pages_content: List[Dict[str, Any]]) -> int:
        """根据位置获取页码。"""
        for page in pages_content:
            if page["start_pos"] <= pos < page["end_pos"]:
                return int(page["page_num"])
        return 1

    def _extract_title_number(self, title: str) -> str:
        """从标题中提取编号部分。"""
        if not title:
            return ""

        if re.match(r"^(第[一二三四五六七八九十\d]+[章节条款部分])", title):
            return re.match(r"^(第[一二三四五六七八九十\d]+[章节条款部分])", title).group(1)

        if re.match(r"^(【\d+】)", title):
            return re.match(r"^(【\d+】)", title).group(1)

        if re.match(r"^(〖\d+(?:\.\d+)*〗)", title):
            return re.match(r"^(〖\d+(?:\.\d+)*〗)", title).group(1)

        if re.match(r"^(\d+(?:\.\d+)*)", title):
            return re.match(r"^(\d+(?:\.\d+)*)", title).group(1)

        if re.match(r"^([一二三四五六七八九十]+)[、．.）\)]", title):
            return re.match(r"^([一二三四五六七八九十]+)[、．.）\)]", title).group(1)

        if re.match(r"^([\(（][一二三四五六七八九十\d]+[\)）])", title):
            return re.match(r"^([\(（][一二三四五六七八九十\d]+[\)）])", title).group(1)

        return ""

    def _extract_title_number_path(self, section_label: str) -> str:
        """从section_label中提取标题路径的编号路径。"""
        if not section_label:
            return ""

        parts = section_label.split("->")
        number_paths: List[str] = []

        for part in parts:
            part = part.strip()
            if part:
                number = self._extract_title_number(part)
                if number:
                    number_paths.append(number)

        if number_paths:
            return "->".join(number_paths)

        return ""

    def _extract_number_from_section_label(self, section_label: str) -> str:
        """从section_label中提取最底层级的编号。"""
        if not section_label:
            return ""

        if "->" in section_label:
            last_level_part = section_label.split("->")[-1].strip()
        else:
            last_level_part = section_label.strip()

        if " + " in last_level_part:
            merged_parts = last_level_part.split(" + ")
            numbers: List[str] = []
            for part in merged_parts:
                part = part.strip()
                number = self._extract_title_number(part)
                if number:
                    numbers.append(number)

            if numbers:
                return "+".join(numbers)

        return self._extract_title_number(last_level_part)
