# Copyright (C) 2026 NPO NT
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

import hashlib
from pathlib import Path

import pymupdf as fitz
from PIL import Image, ImageOps, ImageSequence, UnidentifiedImageError

from .models import Box, DocumentData, PageData, WordToken
from .progress import CancelSignal, ProgressCallback, ensure_not_cancelled, report_progress


MAX_SOURCE_BYTES = 100 * 1024 * 1024
MAX_PAGES = 200
MAX_IMAGE_PIXELS = 100_000_000
DEFAULT_DPI = 220


class DocumentLoadError(RuntimeError):
    pass


def _sha256_file(
    path: Path,
    *,
    source_size: int,
    progress_callback: ProgressCallback | None = None,
    cancel_signal: CancelSignal | None = None,
) -> str:
    digest = hashlib.sha256()
    processed = 0
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            ensure_not_cancelled(cancel_signal)
            digest.update(chunk)
            processed += len(chunk)
            report_progress(progress_callback, "hashing", processed, source_size)
    return digest.hexdigest()


def _register_heif() -> bool:
    try:
        from pillow_heif import register_heif_opener  # type: ignore[import-not-found]
    except ImportError:
        return False
    register_heif_opener(thumbnails=False)
    return True


def heif_available() -> bool:
    return _register_heif()


def _load_pdf(
    path: Path,
    dpi: int,
    max_pages: int,
    *,
    progress_callback: ProgressCallback | None = None,
    cancel_signal: CancelSignal | None = None,
) -> tuple[list[PageData], list[str]]:
    pages: list[PageData] = []
    warnings: list[str] = []
    try:
        document = fitz.open(path)
    except Exception as exc:
        raise DocumentLoadError("Не удалось безопасно открыть PDF") from exc
    try:
        if document.needs_pass:
            raise DocumentLoadError("Зашифрованные PDF в этой версии не поддерживаются")
        if document.page_count <= 0:
            raise DocumentLoadError("PDF не содержит страниц")
        if document.page_count > max_pages:
            raise DocumentLoadError(f"PDF содержит больше {max_pages} страниц")
        scale = dpi / 72.0
        matrix = fitz.Matrix(scale, scale)
        for page_index in range(document.page_count):
            ensure_not_cancelled(cancel_signal)
            report_progress(progress_callback, "rendering", page_index + 1, document.page_count)
            page = document.load_page(page_index)
            rect = page.rect
            try:
                pixmap = page.get_pixmap(matrix=matrix, alpha=False, colorspace=fitz.csRGB)
            except Exception as exc:
                raise DocumentLoadError(f"Не удалось отрисовать страницу {page_index + 1}") from exc
            image = Image.frombytes("RGB", (pixmap.width, pixmap.height), pixmap.samples).copy()
            native_tokens: list[WordToken] = []
            try:
                words = page.get_text("words", sort=True)
            except Exception:
                words = []
                warnings.append(f"Страница {page_index + 1}: не удалось извлечь нативный текст; требуется OCR")
            for item in words:
                if len(item) < 8:
                    continue
                x0, y0, x1, y1, text, block_no, line_no, _word_no = item[:8]
                if not str(text).strip():
                    continue
                native_tokens.append(
                    WordToken(
                        text=str(text),
                        box=Box(float(x0) * scale, float(y0) * scale, float(x1) * scale, float(y1) * scale),
                        line_key=(page_index, int(block_no), int(line_no)),
                        confidence=100.0,
                        source="native",
                    )
                )
            pages.append(
                PageData(
                    index=page_index,
                    image=image,
                    width_points=float(rect.width),
                    height_points=float(rect.height),
                    native_tokens=native_tokens,
                )
            )
            report_progress(progress_callback, "render_complete", page_index + 1, document.page_count)
    finally:
        document.close()
    return pages, warnings


def _load_image(
    path: Path,
    max_pages: int,
    *,
    progress_callback: ProgressCallback | None = None,
    cancel_signal: CancelSignal | None = None,
) -> tuple[list[PageData], list[str]]:
    _register_heif()
    Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
    pages: list[PageData] = []
    warnings: list[str] = []
    try:
        source = Image.open(path)
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise DocumentLoadError("Формат изображения не поддерживается или файл повреждён") from exc
    try:
        frame_count = min(int(getattr(source, "n_frames", 1)), max_pages)
        for page_index, frame in enumerate(ImageSequence.Iterator(source)):
            ensure_not_cancelled(cancel_signal)
            if page_index >= max_pages:
                raise DocumentLoadError(f"Изображение содержит больше {max_pages} кадров")
            report_progress(progress_callback, "rendering", page_index + 1, frame_count)
            oriented = ImageOps.exif_transpose(frame.copy()).convert("RGB")
            # Image output uses a neutral 150 DPI physical size; pixels remain unchanged.
            width_points = oriented.width * 72.0 / 150.0
            height_points = oriented.height * 72.0 / 150.0
            pages.append(
                PageData(
                    index=page_index,
                    image=oriented,
                    width_points=width_points,
                    height_points=height_points,
                )
            )
            report_progress(progress_callback, "render_complete", page_index + 1, frame_count)
        if not pages:
            raise DocumentLoadError("Изображение не содержит доступных кадров")
        if source.info:
            warnings.append("Метаданные исходного изображения будут удалены из новой копии")
    finally:
        source.close()
    return pages, warnings


def load_document(
    source_path: str | Path,
    *,
    dpi: int = DEFAULT_DPI,
    max_source_bytes: int = MAX_SOURCE_BYTES,
    max_pages: int = MAX_PAGES,
    progress_callback: ProgressCallback | None = None,
    cancel_signal: CancelSignal | None = None,
) -> DocumentData:
    ensure_not_cancelled(cancel_signal)
    report_progress(progress_callback, "opening")
    path = Path(source_path).expanduser().resolve()
    if not path.is_file():
        raise DocumentLoadError("Файл не найден")
    source_size = path.stat().st_size
    if source_size <= 0:
        raise DocumentLoadError("Файл пуст")
    if source_size > max_source_bytes:
        raise DocumentLoadError(f"Файл превышает лимит {max_source_bytes // (1024 * 1024)} МиБ")

    with path.open("rb") as stream:
        header = stream.read(16)
    ensure_not_cancelled(cancel_signal)
    if header.startswith(b"%PDF-"):
        source_type = "pdf"
        pages, warnings = _load_pdf(
            path,
            dpi,
            max_pages,
            progress_callback=progress_callback,
            cancel_signal=cancel_signal,
        )
    else:
        source_type = "image"
        pages, warnings = _load_image(
            path,
            max_pages,
            progress_callback=progress_callback,
            cancel_signal=cancel_signal,
        )

    report_progress(progress_callback, "hashing", 0, source_size)
    source_sha256 = _sha256_file(
        path,
        source_size=source_size,
        progress_callback=progress_callback,
        cancel_signal=cancel_signal,
    )
    ensure_not_cancelled(cancel_signal)

    return DocumentData(
        source_path=path,
        source_type=source_type,
        source_size=source_size,
        source_sha256=source_sha256,
        pages=pages,
        warnings=warnings,
    )
