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

from __future__ import annotations

import csv
import io
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path

from PIL import Image

from .models import Box, WordToken
from .progress import CancelSignal, ensure_not_cancelled


class OcrError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class OcrResult:
    tokens: list[WordToken]
    mean_confidence: float | None
    languages: tuple[str, ...]


def project_tessdata_dir() -> Path:
    return Path(__file__).resolve().parents[2] / "models" / "tessdata"


class TesseractOcr:
    def __init__(
        self,
        executable: str | None = None,
        preferred_languages: tuple[str, ...] = ("rus", "eng"),
        tessdata_dir: Path | None = None,
    ) -> None:
        self.executable = executable or shutil.which("tesseract") or ""
        candidate = tessdata_dir or project_tessdata_dir()
        self.tessdata_dir = candidate if candidate.exists() and any(candidate.glob("*.traineddata")) else None
        self.preferred_languages = preferred_languages
        self._available_languages: tuple[str, ...] | None = None

    @property
    def available(self) -> bool:
        return bool(self.executable) and bool(self.available_languages)

    @property
    def available_languages(self) -> tuple[str, ...]:
        if self._available_languages is not None:
            return self._available_languages
        if not self.executable:
            self._available_languages = ()
            return self._available_languages
        command = [self.executable]
        if self.tessdata_dir is not None:
            command.extend(["--tessdata-dir", str(self.tessdata_dir)])
        command.append("--list-langs")
        try:
            completed = subprocess.run(command, capture_output=True, text=True, check=True, timeout=20)
        except (OSError, subprocess.SubprocessError):
            self._available_languages = ()
            return self._available_languages
        lines = [line.strip() for line in completed.stdout.splitlines() if line.strip()]
        self._available_languages = tuple(line for line in lines if not line.lower().startswith("list of available"))
        return self._available_languages

    @property
    def selected_languages(self) -> tuple[str, ...]:
        available = set(self.available_languages)
        selected = tuple(language for language in self.preferred_languages if language in available)
        if selected:
            return selected
        if "eng" in available:
            return ("eng",)
        return tuple(sorted(available))[:1]

    def recognize(
        self,
        image: Image.Image,
        page_index: int,
        *,
        cancel_signal: CancelSignal | None = None,
    ) -> OcrResult:
        ensure_not_cancelled(cancel_signal)
        if not self.executable:
            raise OcrError("Tesseract не найден в PATH")
        languages = self.selected_languages
        if not languages:
            raise OcrError("Не найдено ни одной локальной языковой модели Tesseract")
        buffer = io.BytesIO()
        image.convert("RGB").save(buffer, format="PNG", optimize=False)
        command = [self.executable]
        if self.tessdata_dir is not None:
            command.extend(["--tessdata-dir", str(self.tessdata_dir)])
        command.extend(
            [
                "stdin",
                "stdout",
                "-l",
                "+".join(languages),
                "--psm",
                "6",
                "-c",
                "preserve_interword_spaces=1",
                "-c",
                "tessedit_create_tsv=1",
            ]
        )
        try:
            completed = subprocess.run(
                command,
                input=buffer.getvalue(),
                capture_output=True,
                check=True,
                timeout=120,
            )
        except subprocess.TimeoutExpired as exc:
            raise OcrError("OCR превысил безопасный лимит времени") from exc
        except subprocess.CalledProcessError as exc:
            message = exc.stderr.decode("utf-8", errors="replace").strip()
            raise OcrError(f"Tesseract завершился с ошибкой: {message[:240]}") from exc
        ensure_not_cancelled(cancel_signal)

        text = completed.stdout.decode("utf-8", errors="replace")
        tokens, mean = _parse_tsv(text, page_index)
        return OcrResult(tokens=tokens, mean_confidence=mean, languages=languages)


def _parse_tsv(text: str, page_index: int) -> tuple[list[WordToken], float | None]:
    """Parse Tesseract TSV literally; OCR quote glyphs are data, not CSV quoting."""

    reader = csv.DictReader(io.StringIO(text), delimiter="\t", quoting=csv.QUOTE_NONE)
    tokens: list[WordToken] = []
    confidences: list[float] = []
    for row in reader:
        word = (row.get("text") or "").strip()
        if not word:
            continue
        try:
            confidence = float(row.get("conf") or -1)
            left = int(row.get("left") or 0)
            top = int(row.get("top") or 0)
            width = int(row.get("width") or 0)
            height = int(row.get("height") or 0)
            line_key = (
                page_index,
                int(row.get("block_num") or 0),
                int(row.get("par_num") or 0),
                int(row.get("line_num") or 0),
            )
        except (TypeError, ValueError):
            continue
        if width <= 0 or height <= 0:
            continue
        tokens.append(
            WordToken(
                text=word,
                box=Box(left, top, left + width, top + height),
                line_key=line_key,
                confidence=confidence,
                source="ocr",
            )
        )
        if confidence >= 0:
            confidences.append(confidence)
    mean = sum(confidences) / len(confidences) if confidences else None
    return tokens, mean
