from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import NamedTuple


@dataclass
class Document:
    page_content: str
    metadata: dict = field(default_factory=dict)


class MarkdownHeaderSplitter:
    """
    Разбивает Markdown-документ на чанки по иерархии заголовков.

    Каждый чанк — это содержимое одного раздела (от заголовка до
    следующего заголовка того же или более высокого уровня).
    В начало каждого чанка добавляется breadcrumb-путь заголовков.

    Args:
        headers_to_split_on: список пар (маркер, имя_поля), например
            [("#", "h1"), ("##", "h2"), ("###", "h3")].
            По умолчанию — все уровни H1–H6.
        max_chunk_size: максимальный размер чанка в символах.
            Секции больше этого значения дробятся fallback-сплиттером.
        strip_headers: не включать строку заголовка в текст чанка.
        inject_context: добавлять breadcrumb в начало page_content.
    """

    DEFAULT_HEADERS: list[tuple[str, str]] = [
        ("#",      "h1"),
        ("##",     "h2"),
        ("###",    "h3"),
        ("####",   "h4"),
        ("#####",  "h5"),
        ("######", "h6"),
    ]

    def __init__(
        self,
        headers_to_split_on: list[tuple[str, str]] | None = None,
        max_chunk_size: int = 2000,
        strip_headers: bool = True,
        inject_context: bool = True,
    ) -> None:
        raw = headers_to_split_on or self.DEFAULT_HEADERS
        # Сортируем от длинного к короткому: "###" должен проверяться раньше "##"
        self.headers = sorted(raw, key=lambda x: -len(x[0]))
        self.max_chunk_size = max_chunk_size
        self.strip_headers = strip_headers
        self.inject_context = inject_context
        self._fallback: "RecursiveCharacterSplitter | None" = None

    # ──────────────────────────────────────────────
    # Публичный API
    # ──────────────────────────────────────────────

    def split_text(self, text: str) -> list[Document]:
        """Разбить markdown-строку на список Document."""
        lines = text.split("\n")
        sections = self._parse_sections(lines)
        return self._sections_to_documents(sections)

    def split_documents(self, docs: list[Document]) -> list[Document]:
        """Обработать список документов (сохраняет исходные metadata)."""
        result = []
        for doc in docs:
            chunks = self.split_text(doc.page_content)
            for chunk in chunks:
                chunk.metadata = {**doc.metadata, **chunk.metadata}
            result.extend(chunks)
        return result

    # ──────────────────────────────────────────────
    # Внутренние методы
    # ──────────────────────────────────────────────

    def _parse_sections(self, lines: list[str]) -> list[dict]:
        """
        Один проход по строкам документа.
        Возвращает список словарей:
          {"path": [(level, field_name, title), ...], "content": "..."}
        """
        current_path: list[tuple[int, str, str]] = []  # (level, name, title)
        current_lines: list[str] = []
        sections: list[dict] = []

        for line in lines:
            matched_header = self._detect_header(line)
            if matched_header:
                level, field_name, title = matched_header
                # Сохраняем накопленный контент как секцию
                if current_lines or current_path:
                    content = "\n".join(current_lines).strip()
                    if content:
                        sections.append({
                            "path": list(current_path),
                            "content": content,
                        })
                # Обновляем иерархию: убираем заголовки того же и глубже уровня
                current_path = [
                    (l, n, t) for l, n, t in current_path if l < level
                ]
                current_path.append((level, field_name, title))
                # Если не вырезаем заголовки — добавляем строку в контент
                current_lines = [] if self.strip_headers else [line]
            else:
                current_lines.append(line)

        # Финальная секция
        if current_lines:
            content = "\n".join(current_lines).strip()
            if content:
                sections.append({
                    "path": list(current_path),
                    "content": content,
                })

        return sections

    def _detect_header(
        self, line: str
    ) -> tuple[int, str, str] | None:
        """
        Возвращает (level, field_name, title) если строка — заголовок,
        иначе None.
        Условие: строка начинается с N символов '#', затем ровно один пробел.
        """
        for marker, field_name in self.headers:
            prefix = marker + " "
            if line.startswith(prefix):
                # Убедимся, что это не "### нет" вместо "## нет"
                # (уже гарантировано сортировкой по длине маркера)
                title = line[len(prefix):].strip()
                level = len(marker)
                return (level, field_name, title)
        return None

    def _build_path_string(self, path: list[tuple]) -> str:
        """'FastAPI Guide > Установка > Конфигурация'"""
        return " > ".join(title for _, _, title in path)

    def _sections_to_documents(self, sections: list[dict]) -> list[Document]:
        docs: list[Document] = []

        for section in sections:
            path_str = self._build_path_string(section["path"])
            content = section["content"]

            # Препендим breadcrumb в текст (влияет на embedding)
            if self.inject_context and path_str:
                full_text = f"[{path_str}]\n\n{content}"
            else:
                full_text = content

            metadata: dict = {"heading_path": path_str}
            # Также сохраняем отдельные уровни для фильтрации
            for _, field_name, title in section["path"]:
                metadata[field_name] = title

            if len(full_text) <= self.max_chunk_size:
                docs.append(Document(page_content=full_text, metadata=metadata))
            else:
                # Fallback: дробим большую секцию recursive splitter'ом
                sub_chunks = self._fallback_split(content)
                for sub in sub_chunks:
                    sub_text = f"[{path_str}]\n\n{sub}" if (self.inject_context and path_str) else sub
                    docs.append(Document(
                        page_content=sub_text,
                        metadata={**metadata, "sub_split": True},
                    ))

        return docs

    def _fallback_split(self, text: str) -> list[str]:
        """Рекурсивный сплиттер как fallback для больших секций."""
        if self._fallback is None:
            self._fallback = RecursiveCharacterSplitter(
                chunk_size=self.max_chunk_size,
                chunk_overlap=200,
            )
        return self._fallback.split_text(text)


# ─────────────────────────────────────────────────────────────
# Минимальный RecursiveCharacterSplitter для fallback
# ─────────────────────────────────────────────────────────────

class RecursiveCharacterSplitter:
    """Упрощённый recursive splitter для использования как fallback."""

    DEFAULT_SEPS = ["\n\n", "\n", ". ", " ", ""]

    def __init__(
        self,
        chunk_size: int = 2000,
        chunk_overlap: int = 200,
        separators: list[str] | None = None,
    ) -> None:
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.seps = separators or self.DEFAULT_SEPS

    def split_text(self, text: str) -> list[str]:
        return self._split(text, self.seps)

    def _split(self, text: str, seps: list[str]) -> list[str]:
        if len(text) <= self.chunk_size:
            return [text] if text.strip() else []
        sep = next((s for s in seps if s in text), seps[-1])
        parts = text.split(sep)
        chunks: list[str] = []
        buf = ""
        for part in parts:
            candidate = buf + (sep if buf else "") + part
            if len(candidate) <= self.chunk_size:
                buf = candidate
            else:
                if buf:
                    chunks.append(buf)
                    # Overlap: сохраняем конец предыдущего чанка
                    overlap_start = max(0, len(buf) - self.chunk_overlap)
                    buf = buf[overlap_start:] + (sep if buf[overlap_start:] else "") + part
                else:
                    # Часть сама по себе слишком велика — рекурсия
                    sub_seps = seps[seps.index(sep) + 1:] if sep in seps[1:] else [""]
                    chunks.extend(self._split(part, sub_seps))
                    buf = ""
        if buf.strip():
            chunks.append(buf)
        return chunks

SAMPLE = """
# FastAPI: полное руководство

## Установка

Установите FastAPI и uvicorn через pip:

```
pip install fastapi uvicorn[standard]
```

## Первый маршрут

### Декоратор @app.get

Декоратор регистрирует функцию как обработчик GET-запросов.
Путь передаётся первым аргументом.

### Параметры пути

FastAPI автоматически извлекает параметры из URL:

```python
@app.get("/items/{item_id}")
def get_item(item_id: int):
    return {"id": item_id}
```

## Развёртывание

Запустите сервер командой:

    uvicorn main:app --host 0.0.0.0 --port 8000
"""

splitter = MarkdownHeaderSplitter(
    headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")],
    max_chunk_size=2000,
    inject_context=True,
)

chunks = splitter.split_text(SAMPLE)

for i, chunk in enumerate(chunks, 1):
    print(f"=== Chunk {i} ===")
    print(f"Metadata: {chunk.metadata}")
    print(f"Content:\n{chunk.page_content[:200]}")
    print()

# Вывод:
# === Chunk 1 ===
# Metadata: {'heading_path': 'FastAPI: полное руководство > Установка', 'h1': 'FastAPI: полное руководство', 'h2': 'Установка'}
# Content:
# [FastAPI: полное руководство > Установка]
#
# Установите FastAPI и uvicorn через pip:
# ...
#
# === Chunk 2 ===
# Metadata: {'heading_path': '... > Первый маршрут > Декоратор @app.get', 'h1': ..., 'h2': ..., 'h3': ...}
# Content:
# [FastAPI: полное руководство > Первый маршрут > Декоратор @app.get]
# ...

from bs4 import BeautifulSoup, NavigableString, Tag


class HtmlStructureSplitter:
    """
    Разбивает HTML-документ на чанки по заголовкам h1–h3.

    Args:
        split_on_headings: теги, которые считаются разделителями секций.
        max_chunk_size: максимальная длина чанка (символы).
        content_selector: CSS-селектор корневого элемента (e.g. "main", "article").
    """

    BLOCK_TAGS = {"p", "li", "blockquote", "pre", "td", "th", "dd", "dt"}

    def __init__(
        self,
        split_on_headings: list[str] | None = None,
        max_chunk_size: int = 2000,
        content_selector: str | None = None,
    ) -> None:
        self.split_tags = split_on_headings or ["h1", "h2", "h3"]
        self.max_chunk_size = max_chunk_size
        self.content_selector = content_selector
        self._fallback = RecursiveCharacterSplitter(chunk_size=max_chunk_size)

    def split_html(self, html: str) -> list[Document]:
        soup = BeautifulSoup(html, "html.parser")

        # Находим корневой элемент с контентом
        if self.content_selector:
            root = soup.select_one(self.content_selector) or soup
        else:
            root = soup.find("main") or soup.find("article") or soup.body or soup

        return self._extract_sections(root)

    def _extract_sections(self, root) -> list[Document]:
        docs: list[Document] = []
        current_headings: dict[str, str] = {}  # {"h1": "...", "h2": "..."}
        current_parts: list[str] = []

        def flush():
            if not current_parts:
                return
            path = self._build_path(current_headings)
            content = "\n\n".join(p for p in current_parts if p)
            if not content.strip():
                return
            full = f"[{path}]\n\n{content}" if path else content
            if len(full) <= self.max_chunk_size:
                docs.append(Document(page_content=full, metadata={"heading_path": path, **current_headings}))
            else:
                for sub in self._fallback.split_text(content):
                    sub_text = f"[{path}]\n\n{sub}" if path else sub
                    docs.append(Document(page_content=sub_text, metadata={"heading_path": path, **current_headings, "sub_split": True}))

        # Обходим только прямые дочерние элементы верхнего уровня
        # (не рекурсируем вглубь — текст берём через get_text)
        for el in root.children:
            if not isinstance(el, Tag):
                continue
            tag = el.name.lower() if el.name else ""

            if tag in self.split_tags:
                flush()
                current_parts = []
                level = int(tag[1])
                title = el.get_text(" ", strip=True)
                current_headings[tag] = title
                # Удаляем более глубокие заголовки
                for i in range(level + 1, 7):
                    current_headings.pop(f"h{i}", None)

            elif tag in self.BLOCK_TAGS or tag in {"ul", "ol", "table", "figure"}:
                text = el.get_text(" ", strip=True)
                if text:
                    current_parts.append(text)

        flush()
        return docs

    def _build_path(self, headings: dict[str, str]) -> str:
        return " > ".join(
            headings[f"h{i}"]
            for i in range(1, 7)
            if f"h{i}" in headings
        )


# Использование
html = """


  
Python asyncio
  
Зачем нужен async
  
Синхронный код блокирует поток...
  
Coroutines
  
Корутина объявляется через async def...
  
Ключевое слово await
  
await приостанавливает корутину до завершения...


"""

splitter = HtmlStructureSplitter(split_on_headings=["h1", "h2", "h3"])
chunks = splitter.split_html(html)
for chunk in chunks:
    print(chunk.metadata["heading_path"])
    print(chunk.page_content[:100])
    print()

from docx import Document as DocxFile


class DocxStructureSplitter:
    """
    Разбивает DOCX-файл на чанки по стилям абзацев Heading 1–3
    (и их русскоязычным эквивалентам).
    """

    # Маппинг стилей → уровень (учитываем оба языка)
    HEADING_STYLES: dict[str, int] = {
        "Heading 1": 1, "Heading 2": 2, "Heading 3": 3,
        "Heading 4": 4, "Heading 5": 5, "Heading 6": 6,
        "Заголовок 1": 1, "Заголовок 2": 2, "Заголовок 3": 3,
        "Заголовок 4": 4, "Заголовок 5": 5, "Заголовок 6": 6,
    }

    def __init__(
        self,
        max_chunk_size: int = 2000,
        inject_context: bool = True,
        extra_heading_styles: dict[str, int] | None = None,
    ) -> None:
        self.max_chunk_size = max_chunk_size
        self.inject_context = inject_context
        self._styles = {**self.HEADING_STYLES, **(extra_heading_styles or {})}
        self._fallback = RecursiveCharacterSplitter(chunk_size=max_chunk_size)

    def split_file(self, path: str) -> list[Document]:
        doc = DocxFile(path)
        return self._split_paragraphs(doc.paragraphs)

    def split_paragraphs(self, paragraphs) -> list[Document]:
        """Принимает список Paragraph-объектов python-docx."""
        return self._split_paragraphs(paragraphs)

    def _split_paragraphs(self, paragraphs) -> list[Document]:
        docs: list[Document] = []
        current_path: list[tuple[int, str]] = []  # (level, title)
        current_parts: list[str] = []

        def flush():
            if not current_parts:
                return
            content = "\n\n".join(p for p in current_parts if p)
            if not content.strip():
                return
            path_str = " > ".join(t for _, t in current_path)
            full = f"[{path_str}]\n\n{content}" if (self.inject_context and path_str) else content
            metadata = {
                "heading_path": path_str,
                **{f"h{l}": t for l, t in current_path},
            }
            if len(full) <= self.max_chunk_size:
                docs.append(Document(page_content=full, metadata=metadata))
            else:
                for sub in self._fallback.split_text(content):
                    sub_text = f"[{path_str}]\n\n{sub}" if (self.inject_context and path_str) else sub
                    docs.append(Document(page_content=sub_text, metadata={**metadata, "sub_split": True}))

        for para in paragraphs:
            style_name = para.style.name if para.style else ""
            level = self._styles.get(style_name)
            text = para.text.strip()

            if level and text:
                flush()
                current_parts = []
                # Обновляем иерархию: убираем уровни ≥ текущего
                current_path = [(l, t) for l, t in current_path if l < level]
                current_path.append((level, text))
            elif text:
                current_parts.append(text)

        flush()
        return docs


# Использование
splitter = DocxStructureSplitter(max_chunk_size=2000)
chunks = splitter.split_file("technical_spec.docx")

print(f"Всего чанков: {len(chunks)}")
for chunk in chunks[:3]:
    print(f"\n[{chunk.metadata['heading_path']}]")
    print(chunk.page_content[:150] + "...")

from pathlib import Path


class DocumentAwareChunker:
    """
    Единая точка входа для document-aware chunking.
    Автоматически выбирает сплиттер по расширению файла или
    по явно указанному формату.

    Поддерживаемые форматы: markdown, html, docx.
    """

    def __init__(
        self,
        max_chunk_size: int = 2000,
        inject_context: bool = True,
        md_headers: list[tuple[str, str]] | None = None,
        html_split_on: list[str] | None = None,
        html_selector: str | None = None,
    ) -> None:
        self._md = MarkdownHeaderSplitter(
            headers_to_split_on=md_headers,
            max_chunk_size=max_chunk_size,
            inject_context=inject_context,
        )
        self._html = HtmlStructureSplitter(
            split_on_headings=html_split_on,
            max_chunk_size=max_chunk_size,
            content_selector=html_selector,
        )
        self._docx = DocxStructureSplitter(
            max_chunk_size=max_chunk_size,
            inject_context=inject_context,
        )

    def split_file(self, path: str | Path) -> list[Document]:
        """Автоопределение формата по расширению файла."""
        p = Path(path)
        suffix = p.suffix.lower()
        text = None

        if suffix in (".md", ".markdown"):
            text = p.read_text(encoding="utf-8")
            return self._md.split_text(text)

        elif suffix in (".html", ".htm"):
            text = p.read_text(encoding="utf-8")
            return self._html.split_html(text)

        elif suffix == ".docx":
            return self._docx.split_file(str(p))

        else:
            raise ValueError(
                f"Неподдерживаемый формат: {suffix}. "
                f"Используйте split_text() или split_html() напрямую."
            )

    def split_text(self, text: str, format: str = "markdown") -> list[Document]:
        """Разбить текст явно указанного формата."""
        if format == "markdown":
            return self._md.split_text(text)
        elif format == "html":
            return self._html.split_html(text)
        else:
            raise ValueError(f"Неподдерживаемый формат: {format}")


# ─────────────────────────────────────────────────────────────
# Пример использования в RAG-пайплайне
# ─────────────────────────────────────────────────────────────

from pathlib import Path

chunker = DocumentAwareChunker(
    max_chunk_size=2000,
    inject_context=True,
)

# Обработка директории с документацией
docs_dir = Path("docs/")
all_chunks: list[Document] = []

for file_path in docs_dir.rglob("*.md"):
    try:
        chunks = chunker.split_file(file_path)
        # Добавляем имя файла в метаданные
        for chunk in chunks:
            chunk.metadata["source"] = str(file_path)
        all_chunks.extend(chunks)
    except Exception as e:
        print(f"Ошибка при обработке {file_path}: {e}")

print(f"Всего чанков: {len(all_chunks)}")

# Статистика по размерам
sizes = [len(c.page_content) for c in all_chunks]
print(f"Средний размер: {sum(sizes) / len(sizes):.0f} символов")
print(f"Максимальный:   {max(sizes)} символов")
print(f"Минимальный:    {min(sizes)} символов")

# Фильтрация при поиске: только чанки из раздела "API Reference"
api_chunks = [
    c for c in all_chunks
    if "API Reference" in c.metadata.get("heading_path", "")
]
print(f"\nЧанков в разделе API Reference: {len(api_chunks)}")