"""
chapter_segmenter.py
对外接口：segment_chapters(data) -> List[Chapter]

纯规则章节切分，不依赖大模型，不依赖 extract_util / entity_registry。
输入：content_list JSON（list of dict，每条含 type/text/text_level/page_idx/bbox 等字段）
输出：Chapter 列表，每个 Chapter 包含标题信息、所属 records、以及二级子章节。

多信号融合标题识别（v8）：
  - 两遍扫描：先扫描所有候选标题 → 全局校验编号连续性 → 再切分
  - 置信度融合：text_level + 编号格式 + 文本长度 + 标点密度
  - 全局校验：编号跳跃过大的假标题自动剔除
"""

import re
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Set


# ==============================
# 配置（可在调用前修改模块级变量）
# ==============================

# 一级标题含以下关键词时整章跳过
FILTER_KEYWORDS: List[str] = [
    "维修工作卡", "维修项目索引",
    "常见故障", "维修项目表"
]

# 二级标题含以下关键词时跳过该二级节
FILTER_L2_KEYWORDS: List[str] = [
     "常见故障表", "维修项目表", "维修工作卡","操作程序"
]

# 特殊章节关键词：一级标题含以下任一时，不作为整体实体，二级标题各自独立
SPECIAL_CHAPTER_KEYWORDS: List[str] = [
    "组成、技术特征及接口",
    "组成和技术特征",
    "组成与技术特征",
    "组成、技术特征",
]

# 目录关键词（去空格后匹配，兼容"目 次"/"目 录"等 OCR 带空格的写法）
TOC_KEYWORDS: List[str] = ["目次", "目录", "contents"]

# 置信度阈值：低于此值的不认为是标题
HEADING_CONFIDENCE_THRESHOLD: float = 0.5

# 一级标题编号最大允许跳跃数（连续性校验）
MAX_L1_GAP: int = 3


# ==============================
# 数据结构
# ==============================

@dataclass
class L2Section:
    """二级子章节"""
    title: str          # 去编号后的标题名
    numbering: str      # 编号，如 "3.2"
    records: List[dict] = field(default_factory=list)  # 该二级节的所有 records


@dataclass
class Chapter:
    """一级章节（含所有二级子节）"""
    title: str              # 去编号后的标题名（已应用 L1_TITLE_MAPPING 等由调用方处理）
    raw_title: str          # 原始文本（含编号）
    numbering: str          # 编号，如 "3"
    is_special: bool        # 是否特殊章节（二级各自独立）
    is_skipped: bool        # 是否被 FILTER_KEYWORDS 过滤
    records: List[dict] = field(default_factory=list)      # 本章所有 records（含标题行）
    loose_records: List[dict] = field(default_factory=list) # 一级标题到第一个二级标题之间的 records
    l2_sections: List[L2Section] = field(default_factory=list)  # 二级子章节列表


@dataclass
class _HeadingNode:
    """内部用：统一的标题节点"""
    raw_text: str
    clean_title: str
    numbering: str
    level: int
    index_in_data: int
    confidence: float = 1.0


# ==============================
# 正则
# ==============================

_CHAPTER_NUM_RE = re.compile(r'^(\d{1,2}(?:\.\d{1,2})*\.?)')
_LIST_ITEM_RE   = re.compile(r'^[\d]+[）\)]\s*')
_ALPHA_LIST_RE  = re.compile(r'^[a-zA-Z][）\)\.]\s*')
_TOC_ENTRY_SUFFIX_RE = re.compile(r'([\s.]+\d{1,4}|\.{2,})\s*$')


# ==============================
# 标题工具函数
# ==============================

def _is_real_chapter_title(text: str) -> bool:
    text = text.strip()
    if _LIST_ITEM_RE.match(text):
        return False
    if _ALPHA_LIST_RE.match(text):
        return False
    if re.search(r'[。！？：]\s*$', text):
        return False
    chapter_num = r'^\d{1,2}(?:\.\d{1,2})*'
    if re.match(chapter_num + r'\s+\S', text):
        return True
    if re.match(chapter_num + r'\.\s+\S', text):
        return True
    if re.match(chapter_num + r'\.?[\u4e00-\u9fff]', text):
        return True
    return False


def _is_toc_entry_line(text: str) -> bool:
    if '\n' in text:
        return False
    return bool(_TOC_ENTRY_SUFFIX_RE.search(text))


def _get_chapter_level(text: str) -> int:
    match = _CHAPTER_NUM_RE.match(text.strip())
    if match:
        num = match.group(1).rstrip('.')
        return num.count('.') + 1
    return 0


def _get_chapter_numbering(text: str) -> str:
    match = _CHAPTER_NUM_RE.match(text.strip())
    return match.group(1).rstrip('.') if match else ""


def _clean_title(text: str) -> str:
    return re.sub(r'^\d{1,2}(?:\.\d{1,2})*\.?\s*', '', text.strip()).strip()


def _should_skip_l2(title_clean: str) -> bool:
    return any(kw in title_clean for kw in FILTER_L2_KEYWORDS)


def _is_special_chapter(title: str) -> bool:
    return any(kw in title for kw in SPECIAL_CHAPTER_KEYWORDS)


# ==============================
# 目录页检测
# ==============================

def _is_toc_content(ins: dict) -> bool:
    """
    判断是否是目录页内容（三层判断）。

    层1：text_level>=1 且含目录关键词 → 直接判定
    层2：严格目录条目格式（"文字..页码" ≥2次）
    层3：宽松兜底：换行多 + '..'密集 + 排除正文列表/技术参数
    """
    text = ins.get("text", "")
    text_level = ins.get("text_level", 0)

    text_nospace = text.replace(" ", "").replace("\u3000", "")
    if text_level and any(kw in text_nospace.lower() for kw in TOC_KEYWORDS):
        return True

    toc_entry_re = re.compile(
        r'[\u4e00-\u9fffa-zA-Z）\)]\s*\.{2,}\s*\.?\s*(\d{1,4})\s*$',
        re.MULTILINE
    )
    if len(toc_entry_re.findall(text)) >= 2:
        return True

    if text.count('\n') > 3 and '..' in text:
        has_list_markers = bool(re.search(
            r'(?:^|\n)\s*(?:\d+[）\)]|[a-zA-Z][）\).])', text
        ))
        if has_list_markers:
            return False
        range_pattern_count = len(re.findall(r'[\d℃%]\.\.[+\-\d]', text))
        dot_count = len(re.findall(r'\.{2,}', text))
        if range_pattern_count > 0 and range_pattern_count >= dot_count * 0.5:
            return False
        line_count = text.count('\n') + 1
        if dot_count >= max(3, line_count * 0.3):
            return True

    return False


def _get_toc_pages(data: list) -> Set[int]:
    """找出所有目录页的 page_idx，支持多页目录，并向后延伸。"""
    toc_pages: Set[int] = set()
    for ins in data:
        if _is_toc_content(ins):
            toc_pages.add(ins.get("page_idx", -1))
        text = ins.get("text", "")
        text_level = ins.get("text_level", 0)
        text_nospace = text.replace(" ", "").replace("\u3000", "")
        if text_level and any(kw in text_nospace.lower() for kw in TOC_KEYWORDS):
            toc_pages.add(ins.get("page_idx", -1))
    toc_pages.discard(-1)

    if not toc_pages:
        return toc_pages

    pages_records: Dict[int, list] = {}
    for ins in data:
        p = ins.get("page_idx", -1)
        if p < 0:
            continue
        pages_records.setdefault(p, []).append(ins)

    max_toc = max(toc_pages)
    for next_page in range(max_toc + 1, max_toc + 10):
        if next_page not in pages_records:
            break
        records = pages_records[next_page]
        text_records = [
            r for r in records
            if r.get("type") == "text" and (r.get("text") or "").strip()
        ]
        if not text_records:
            toc_pages.add(next_page)
            continue
        toc_like = sum(
            1 for r in text_records
            if _is_toc_content(r) or _is_toc_entry_line((r.get("text") or ""))
        )
        if toc_like == len(text_records):
            toc_pages.add(next_page)
        else:
            break

    return toc_pages


# ==============================
# 多信号融合标题识别
# ==============================

def _compute_heading_confidence(ins: dict, text: str, level: int) -> float:
    """
    4个信号融合（0~1）：
      text_level>=1（0.35）+ 编号格式规范（0.30）+ 长度合理（0.20）+ 标点密度低（0.15）
    """
    score = 0.0
    text_level = ins.get("text_level", 0)

    if text_level >= 1:
        score += 0.35

    normalized = re.sub(r'\s+', ' ', text.replace('\u3000', ' '))
    if re.match(r'^\d{1,2}(?:\.\d{1,2})*\.?\s+[\u4e00-\u9fffa-zA-Z]', normalized):
        score += 0.30
    elif re.match(r'^\d{1,2}(?:\.\d{1,2})*[\u4e00-\u9fff]', normalized):
        score += 0.20

    clean = _clean_title(text)
    if 1 <= len(clean) <= 20:
        score += 0.20
    elif len(clean) <= 35:
        score += 0.10

    punct_count = len(re.findall(r'[，。；：、！？,.;:!?]', clean))
    if punct_count <= 1:
        score += 0.15
    elif punct_count <= 3:
        score += 0.05

    return score


def _scan_all_headings(data: list, toc_pages: Set[int] = None) -> List[_HeadingNode]:
    """第一遍：识别所有置信度足够的标题节点。"""
    headings = []
    debug_rejected = []

    for i, ins in enumerate(data):
        if ins.get("type") != "text":
            continue
        if toc_pages and ins.get("page_idx") in toc_pages:
            continue

        text = ins.get("text", "").strip()
        if not text:
            continue
        if _is_toc_content(ins):
            debug_rejected.append(f"  [i={i}] 目录页跳过: {text[:40]!r}")
            continue
        if _is_toc_entry_line(text):
            debug_rejected.append(f"  [i={i}] 目录条目行: {text[:60]!r}")
            continue
        if not _is_real_chapter_title(text):
            first_char = text[0] if text else ''
            if first_char.isdigit() and len(debug_rejected) < 50:
                debug_rejected.append(f"  [i={i}] 非章节标题格式: {text[:60]!r}")
            continue

        level = _get_chapter_level(text)
        if level == 0:
            continue

        numbering = _get_chapter_numbering(text)
        clean = _clean_title(text)
        confidence = _compute_heading_confidence(ins, text, level)

        if confidence >= HEADING_CONFIDENCE_THRESHOLD:
            headings.append(_HeadingNode(
                raw_text=text,
                clean_title=clean,
                numbering=numbering,
                level=level,
                index_in_data=i,
                confidence=confidence,
            ))
        else:
            debug_rejected.append(
                f"  [i={i}] 置信度不足({confidence:.2f}): {text[:60]!r} "
                f"(text_level={ins.get('text_level', 0)})"
            )

    print(f"\n  [DEBUG] data总条数={len(data)}, "
          f"text类型条数={sum(1 for x in data if x.get('type')=='text')}")
    print(f"  [DEBUG] 候选标题={len(headings)}, 被拒绝={len(debug_rejected)}")
    if debug_rejected:
        print("  [DEBUG] 被拒绝样本（最多20条）:")
        for msg in debug_rejected[:20]:
            print(msg)

    return headings


def _validate_heading_sequence(headings: List[_HeadingNode]) -> List[_HeadingNode]:
    """第二遍：全局校验编号连续性，去除可疑标题并去重。"""
    if not headings:
        return headings

    level1_nums = sorted(set(
        int(h.numbering) for h in headings if h.level == 1
    ))

    valid_level1: Set[int] = set()
    if level1_nums:
        for i, num in enumerate(level1_nums):
            if i == 0:
                valid_level1.add(num)
                continue
            gap = num - level1_nums[i - 1]
            if gap <= MAX_L1_GAP:
                valid_level1.add(num)
            else:
                print(f"  ⚠️ 一级标题编号跳跃过大: {level1_nums[i-1]} → {num}，丢弃 {num}")

    validated = []
    for h in headings:
        if h.level == 1:
            if int(h.numbering) in valid_level1:
                validated.append(h)
        elif h.level >= 2:
            parent_num = int(h.numbering.split('.')[0])
            if parent_num in valid_level1:
                validated.append(h)
        else:
            validated.append(h)

    def _has_lower_l2_after(h: _HeadingNode) -> bool:
        if h.level != 1:
            return False
        h_num = int(h.numbering)
        after_l2 = [x for x in validated if x.index_in_data > h.index_in_data and x.level == 2]
        for x in after_l2[:8]:
            parent = int(x.numbering.split('.')[0])
            if parent < h_num:
                return True
            if parent >= h_num:
                break
        return False

    from collections import defaultdict
    by_num: Dict[str, list] = defaultdict(list)
    for h in validated:
        by_num[h.numbering].append(h)

    best: Dict[str, _HeadingNode] = {}
    for numbering, candidates in by_num.items():
        if len(candidates) == 1:
            best[numbering] = candidates[0]
        else:
            ok = [c for c in candidates if not _has_lower_l2_after(c)]
            chosen = ok[0] if ok else candidates[-1]
            for c in candidates:
                if c is not chosen:
                    print(f"  ⚠️ 编号重复，丢弃误识别: {c.numbering} 【{c.clean_title}】")
            best[numbering] = chosen

    return [h for h in validated if best.get(h.numbering) is h]


def _build_heading_index(data: list) -> Dict[int, _HeadingNode]:
    """构建 data索引 → HeadingNode 的映射。"""
    toc_pages = _get_toc_pages(data)
    if toc_pages:
        print(f"  [TOC页] 检测到目录页 page_idx={sorted(toc_pages)}，跳过这些页的标题识别")

    headings = _scan_all_headings(data, toc_pages)
    headings = _validate_heading_sequence(headings)

    if headings:
        print(f"  [标题识别] 共识别 {len(headings)} 个有效标题：")
        for h in headings:
            indent = "    " * h.level
            conf_str = f" (conf={h.confidence:.2f})" if h.confidence < 1.0 else ""
            print(f"  {indent}{h.numbering} {h.clean_title}{conf_str}")

    return {h.index_in_data: h for h in headings}


# ==============================
# 对外接口
# ==============================

def segment_chapters(data: list) -> List[Chapter]:
    """
    对外唯一接口。

    将 content_list JSON 按章节结构切分，返回 Chapter 列表。
    每个 Chapter 包含：
      - title / raw_title / numbering：标题信息
      - is_special：是否特殊章节（二级节各自独立，不合并）
      - is_skipped：是否被 FILTER_KEYWORDS 过滤（内容为空）
      - records：本章全部 records（含标题行，不含目录页记录）
      - loose_records：一级标题到第一个二级标题之间的 records
      - l2_sections：List[L2Section]，每个二级节含 title/numbering/records

    :param data: content_list JSON，list of dict
    :return: List[Chapter]，按文档顺序排列
    """
    toc_pages = _get_toc_pages(data)
    heading_index = _build_heading_index(data)

    chapters: List[Chapter] = []

    # 当前章节状态
    cur: Optional[Chapter] = None
    cur_l2: Optional[L2Section] = None
    has_met_l2 = False

    def _flush_l2():
        nonlocal cur_l2
        if cur_l2 is not None and cur is not None:
            if not _should_skip_l2(cur_l2.title):
                cur.l2_sections.append(cur_l2)
        cur_l2 = None

    def _flush_chapter():
        nonlocal cur, cur_l2, has_met_l2
        if cur is None:
            return
        _flush_l2()
        chapters.append(cur)
        cur = None
        cur_l2 = None
        has_met_l2 = False

    for i, ins in enumerate(data):
        # 跳过目录页
        if toc_pages and ins.get("page_idx") in toc_pages:
            continue
        if _is_toc_content(ins):
            continue
        # discarded 记录：仍然归入当前章节（保留位置信息供调用方使用）
        if ins.get("type") == "discarded":
            if cur is not None and not cur.is_skipped:
                cur.records.append(ins)
            continue

        heading = heading_index.get(i)

        if heading is not None and heading.level == 1:
            # 新的一级标题 → 刷出上一章
            _flush_chapter()
            is_skipped = any(kw in heading.raw_text for kw in FILTER_KEYWORDS)
            cur = Chapter(
                title=heading.clean_title,
                raw_title=heading.raw_text,
                numbering=heading.numbering,
                is_special=_is_special_chapter(heading.raw_text),
                is_skipped=is_skipped,
            )
            has_met_l2 = False
            if not is_skipped:
                cur.records.append(ins)
                tag = " [特殊]" if cur.is_special else ""
                print(f"\n[一级标题] {heading.raw_text}{tag} (conf={heading.confidence:.2f})")
            continue

        if cur is None or cur.is_skipped:
            continue

        if heading is not None and heading.level == 2:
            # 二级标题
            _flush_l2()
            l2_clean = heading.clean_title
            if not _should_skip_l2(l2_clean):
                cur_l2 = L2Section(
                    title=l2_clean,
                    numbering=heading.numbering,
                    records=[ins],
                )
            has_met_l2 = True
            cur.records.append(ins)
            continue

        # 普通内容记录
        cur.records.append(ins)
        if has_met_l2:
            if cur_l2 is not None:
                cur_l2.records.append(ins)
        else:
            cur.loose_records.append(ins)

    _flush_chapter()

    print(f"\n[chapter_segmenter] 共切分章节: {len(chapters)} 个")
    return chapters

