from __future__ import annotations

import mimetypes
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path


@dataclass(frozen=True)
class ExtractedBlock:
    text: str
    page_start: int | None = None
    page_end: int | None = None
    section_title: str = ""
    section_path: str = ""
    paragraph_index: int | None = None


@dataclass(frozen=True)
class ExtractedDocument:
    text: str
    mime_type: str
    blocks: list[ExtractedBlock]


TEXT_EXTENSIONS = {".txt", ".md", ".markdown", ".json", ".csv"}
DOCX_EXTENSIONS = {".docx"}
PDF_EXTENSIONS = {".pdf"}
SUPPORTED_EXTENSIONS = TEXT_EXTENSIONS | DOCX_EXTENSIONS | PDF_EXTENSIONS


def supported_extensions_label() -> str:
    return ", ".join(sorted(SUPPORTED_EXTENSIONS))


def _guess_mime_type(filename: str) -> str:
    guessed, _ = mimetypes.guess_type(filename)
    if guessed:
        return guessed
    ext = Path(filename).suffix.lower()
    if ext == ".docx":
        return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    if ext == ".pdf":
        return "application/pdf"
    return "text/plain"


def extract_document(raw_bytes: bytes, filename: str) -> ExtractedDocument:
    ext = Path(filename or "").suffix.lower()
    if ext in TEXT_EXTENSIONS:
        return _extract_text(raw_bytes, filename)
    if ext in DOCX_EXTENSIONS:
        return _extract_docx(raw_bytes, filename)
    if ext in PDF_EXTENSIONS:
        return _extract_pdf(raw_bytes, filename)
    raise ValueError(f"Неподдерживаемый формат документа. Поддерживаются: {supported_extensions_label()}")


def _extract_text(raw_bytes: bytes, filename: str) -> ExtractedDocument:
    text = raw_bytes.decode("utf-8-sig", errors="replace").strip()
    blocks = [
        ExtractedBlock(text=line.strip(), paragraph_index=idx)
        for idx, line in enumerate(text.splitlines())
        if line.strip()
    ]
    if not blocks and text:
        blocks = [ExtractedBlock(text=text, paragraph_index=0)]
    return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)


def _extract_docx(raw_bytes: bytes, filename: str) -> ExtractedDocument:
    try:
        from docx import Document as DocxDocument
    except ImportError as exc:
        raise ValueError("Для загрузки DOCX установите зависимость python-docx") from exc

    document = DocxDocument(BytesIO(raw_bytes))
    blocks: list[ExtractedBlock] = []
    heading_stack: list[str] = []

    for idx, paragraph in enumerate(document.paragraphs):
        text = paragraph.text.strip()
        if not text:
            continue

        style_name = (paragraph.style.name if paragraph.style else "").strip()
        if style_name.lower().startswith("heading"):
            level = _heading_level(style_name)
            heading_stack = heading_stack[: max(level - 1, 0)]
            heading_stack.append(text)
            section_title = text
        else:
            section_title = heading_stack[-1] if heading_stack else ""

        blocks.append(
            ExtractedBlock(
                text=text,
                section_title=section_title,
                section_path=" > ".join(heading_stack),
                paragraph_index=idx,
            )
        )

    text = "\n\n".join(block.text for block in blocks).strip()
    return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)


def _heading_level(style_name: str) -> int:
    parts = style_name.rsplit(" ", 1)
    if len(parts) == 2 and parts[1].isdigit():
        return max(1, int(parts[1]))
    return 1


def _extract_pdf(raw_bytes: bytes, filename: str) -> ExtractedDocument:
    try:
        from pypdf import PdfReader
    except ImportError as exc:
        raise ValueError("Для загрузки PDF установите зависимость pypdf") from exc

    reader = PdfReader(BytesIO(raw_bytes))
    blocks: list[ExtractedBlock] = []
    for page_index, page in enumerate(reader.pages, start=1):
        page_text = (page.extract_text() or "").strip()
        if not page_text:
            continue
        paragraphs = [part.strip() for part in page_text.split("\n\n") if part.strip()]
        if not paragraphs:
            paragraphs = [page_text]
        for paragraph_index, paragraph in enumerate(paragraphs):
            blocks.append(
                ExtractedBlock(
                    text=paragraph,
                    page_start=page_index,
                    page_end=page_index,
                    paragraph_index=paragraph_index,
                )
            )

    text = "\n\n".join(block.text for block in blocks).strip()
    return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)