"""Document chunking strategies used before Milvus ingestion."""
from __future__ import annotations

import re
from dataclasses import dataclass

from rag.chunk_config import ChunkConfig, resolve_chunk_config


SUPPORTED_STRATEGIES = {"default", "qa_pair", "chapter_semantic"}


@dataclass(frozen=True)
class Chunk:
    text: str
    section_title: str | None = None


@dataclass(frozen=True)
class ChunkingResult:
    chunks: list[Chunk]
    requested_strategy: str
    actual_strategy: str
    degraded: bool = False
    warning: str | None = None


def chunk_document(
    text: str,
    strategy: str,
    *,
    config: ChunkConfig | None = None,
) -> ChunkingResult:
    if strategy not in SUPPORTED_STRATEGIES:
        raise ValueError(f"Unsupported chunk strategy: {strategy}")
    config = config or resolve_chunk_config()
    if strategy == "default":
        return ChunkingResult(_default_chunks(text, config), strategy, strategy)
    if strategy == "qa_pair":
        return ChunkingResult(_qa_pair_chunks(text, config), strategy, strategy)
    return _chapter_chunks(text, config)


def _default_chunks(text: str, config: ChunkConfig) -> list[Chunk]:
    paragraphs = [part.strip() for part in re.split(r"\n\s*\n", text) if part.strip()]
    chunks: list[Chunk] = []
    for paragraph in paragraphs:
        chunks.extend(Chunk(part) for part in _split_long_text(paragraph, config))
    return chunks


def _split_long_text(text: str, config: ChunkConfig) -> list[str]:
    if len(text) <= config.size:
        return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + config.size, len(text))
        chunks.append(text[start:end])
        if end == len(text):
            break
        start = end - config.overlap
    return chunks


def _qa_pair_chunks(text: str, config: ChunkConfig) -> list[Chunk]:
    matches = list(
        re.finditer(
            r"^\s*Q:\s*(.*?)\r?\n\s*A:\s*(.*?)(?=^\s*Q:\s*|\Z)",
            text,
            flags=re.MULTILINE | re.DOTALL,
        )
    )
    if not matches or any(match.group(0).strip() == "" for match in matches):
        raise ValueError("未识别FAQ问答格式，请确认文档包含Q:/A:标记")
    if any(not match.group(1).strip() or not match.group(2).strip() for match in matches):
        raise ValueError("FAQ问答对必须同时包含问题和答案")
    chunks = [
        Chunk(f"question: {match.group(1).strip()}\nanswer: {match.group(2).strip()}")
        for match in matches
    ]
    if any(len(chunk.text) > config.size for chunk in chunks):
        raise ValueError("FAQ问答对超过chunk_size，整份文档拒绝上传")
    return chunks


def _chapter_chunks(text: str, config: ChunkConfig) -> ChunkingResult:
    heading_pattern = re.compile(r"^(#{1,3})\s+(.+?)\s*$", re.MULTILINE)
    headings = list(heading_pattern.finditer(text))
    if not headings:
        return ChunkingResult(
            _default_chunks(text, config),
            "chapter_semantic",
            "default",
            degraded=True,
            warning="未发现Markdown标题，已降级为default策略",
        )

    chunks: list[Chunk] = []
    title_stack: list[str] = []
    for index, heading in enumerate(headings):
        level = len(heading.group(1))
        title = heading.group(2).strip()
        title_stack = title_stack[: level - 1] + [title]
        body_start = heading.end()
        body_end = headings[index + 1].start() if index + 1 < len(headings) else len(text)
        body = text[body_start:body_end].strip()
        if not body:
            continue
        section_title = " > ".join(title_stack)
        prefix = f"【章节：{section_title}】\n"
        for part in _split_long_text(body, config):
            chunks.append(Chunk(prefix + part, section_title=section_title))
    return ChunkingResult(chunks, "chapter_semantic", "chapter_semantic")