"""Indexador: orquestra extração, parsing, embeddings e persistência vetorial.

O manifesto (`manifest.json`) registra, por arquivo, o SHA-256 do conteúdo, o
estado (`complete`, `partial`, `error`) e as páginas com falha. Arquivos cujo
hash não mudou não são reprocessados.
"""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from .bible_parser import BibleParser, parse_book_index
from .config import Settings
from .embedding import Embedder
from .extraction import PageText, file_sha256
from .vector_store import VectorStore

logger = logging.getLogger(__name__)

STATE_COMPLETE = "complete"
STATE_PARTIAL = "partial"
STATE_ERROR = "error"


@dataclass
class IndexStatus:
    """Estado público da indexação."""

    state: str = "ready"
    pdfs_total: int = 0
    pdfs_indexed: int = 0
    chunks: int = 0
    model: str = ""
    last_indexed_at: str | None = None
    errors: list[str] = field(default_factory=list)
    pages_processed: int = 0
    pages_total: int = 0
    current_file: str | None = None


@dataclass
class _ManifestEntry:
    sha256: str
    state: str
    chunks: int = 0
    indexed_at: str = ""
    failed_pages: list[int] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "sha256": self.sha256,
            "state": self.state,
            "chunks": self.chunks,
            "indexed_at": self.indexed_at,
            "failed_pages": self.failed_pages,
        }


class Indexer:
    """Coordena a indexação dos PDFs presentes em `settings.pdf_dir`."""

    def __init__(
        self,
        settings: Settings,
        extractor,
        parser: BibleParser,
        embedder: Embedder,
        store: VectorStore,
    ):
        self.settings = settings
        self.extractor = extractor
        self.parser = parser
        self.embedder = embedder
        self.store = store
        self._manifest: dict[str, _ManifestEntry] = {}
        self._lock = threading.Lock()
        self._last_indexed_at: str | None = None
        self._errors: list[str] = []
        self._indexing = False
        self._pages_processed = 0
        self._pages_total = 0
        self._current_file: str | None = None
        self._load_manifest()

    # ------------------------------------------------------------------ #
    # Manifesto
    # ------------------------------------------------------------------ #

    def _load_manifest(self) -> None:
        path = self.settings.manifest_path
        if not path.exists():
            return
        try:
            raw = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            logger.warning("Manifesto ilegível em %s: %s", path, exc)
            return
        for name, entry in raw.items():
            self._manifest[name] = _ManifestEntry(
                sha256=entry.get("sha256", ""),
                state=entry.get("state", STATE_ERROR),
                chunks=int(entry.get("chunks", 0)),
                indexed_at=entry.get("indexed_at", ""),
                failed_pages=list(entry.get("failed_pages", [])),
            )

    def _save_manifest(self) -> None:
        path = self.settings.manifest_path
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = {name: entry.to_dict() for name, entry in self._manifest.items()}
        temporary = path.with_suffix(".json.tmp")
        temporary.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        temporary.replace(path)

    # ------------------------------------------------------------------ #
    # Índice de livros
    # ------------------------------------------------------------------ #

    def _book_starts(self, pages: list[PageText]) -> dict[str, int]:
        """Índice `{livro: página}` do PDF, memoizado em disco.

        As primeiras páginas do documento contêm o índice de livros. Se ele não
        puder ser lido, devolve um dicionário vazio: o parser cairá para o
        cabeçalho de cada página.
        """
        cached_path = self.settings.book_index_path
        if cached_path.exists():
            try:
                cached = json.loads(cached_path.read_text(encoding="utf-8"))
                if cached:
                    return {str(key): int(value) for key, value in cached.items()}
            except (OSError, json.JSONDecodeError, ValueError) as exc:
                logger.warning("Índice de livros ilegível em %s: %s", cached_path, exc)

        starts: dict[str, int] = {}
        for page in pages[:12]:
            if not page.text.strip():
                continue
            starts.update(parse_book_index(page.text))
            if len(starts) >= 20:
                break

        if starts:
            try:
                cached_path.parent.mkdir(parents=True, exist_ok=True)
                cached_path.write_text(
                    json.dumps(starts, ensure_ascii=False, indent=2), encoding="utf-8"
                )
            except OSError as exc:  # pragma: no cover
                logger.warning("Não foi possível gravar o índice de livros: %s", exc)
        return starts

    # ------------------------------------------------------------------ #
    # Indexação
    # ------------------------------------------------------------------ #

    def mark_indexing(self) -> None:
        """Sinaliza que uma indexação está em andamento."""
        self._indexing = True
        self._pages_processed = 0
        self._pages_total = 0
        self._current_file = None

    def report_progress(
        self, pdf_name: str, pages_done: int, pages_total: int
    ) -> None:
        """Registra o progresso da extração de um arquivo."""
        self._current_file = pdf_name
        self._pages_processed = pages_done
        self._pages_total = pages_total

    def discover_pdfs(self) -> list[Path]:
        """Lista os PDFs do diretório configurado, ordenados por nome."""
        directory = self.settings.pdf_dir
        if not directory.exists():
            return []
        return sorted(path for path in directory.glob("*.pdf") if path.is_file())

    def reindex(self, force: bool = False) -> IndexStatus:
        """Indexa os PDFs novos ou alterados."""
        with self._lock:
            self._errors = []
            self.mark_indexing()
            pdfs = self.discover_pdfs()

            try:
                for pdf in pdfs:
                    try:
                        self._index_file(pdf, force=force)
                    except Exception as exc:  # pragma: no cover - falha inesperada
                        logger.exception("Falha ao indexar %s", pdf.name)
                        self._manifest[pdf.name] = _ManifestEntry(
                            sha256="", state=STATE_ERROR, indexed_at=_now()
                        )
                        self._errors.append(f"{pdf.name}: {exc}")

                self._save_manifest()
                self._last_indexed_at = _now()
            finally:
                self._indexing = False

            return self.status()

    def _extract_pages(self, pdf: Path, file_hash: str) -> list:
        """Extrai as páginas do PDF, usando OCR paralelo quando possível."""
        extract_parallel = getattr(self.extractor, "extract_pages_parallel", None)
        if extract_parallel is None:
            return list(self.extractor.extract_document(pdf, file_hash=file_hash))

        page_count = _pdf_page_count(pdf)
        pages: list = []
        batch_size = max(self.settings.ocr_workers, 1) * 4

        for batch in self._page_batches(pdf, batch_size):
            raw_pages = [
                (page_number, text)
                for page_number, text in batch
            ]
            extracted = extract_parallel(
                raw_pages,
                file_hash=file_hash,
                path=pdf,
                workers=self.settings.ocr_workers,
            )
            pages.extend(extracted)
            self.report_progress(pdf.name, len(pages), page_count)

        return sorted(pages, key=lambda page: page.page)

    def _page_batches(self, pdf: Path, batch_size: int):
        """Gera lotes de (número_da_página, texto_da_camada)."""
        import pymupdf

        with pymupdf.open(pdf) as document:
            batch: list[tuple[int, str]] = []
            for index in range(len(document)):
                page = document[index]
                try:
                    text = page.get_text("text") or ""
                except Exception:  # pragma: no cover - PDF malformado
                    text = ""
                batch.append((index + 1, text))
                if len(batch) >= batch_size:
                    yield batch
                    batch = []
            if batch:
                yield batch

    def _index_file(self, pdf: Path, force: bool) -> None:
        file_hash = file_sha256(pdf)
        entry = self._manifest.get(pdf.name)

        if entry is not None and entry.sha256 == file_hash and not force:
            logger.info("%s já está indexado; pulando", pdf.name)
            return

        logger.info("Indexando %s", pdf.name)
        pages = self._extract_pages(pdf, file_hash)
        failed_pages = [page.page for page in pages if page.failed]

        book_starts = self._book_starts(pages)
        chunks = self.parser.parse(pages, pdf.name, book_starts)

        if not chunks:
            state = STATE_ERROR
            self._errors.append(f"{pdf.name}: nenhum versículo pôde ser extraído")
        elif failed_pages:
            state = STATE_PARTIAL
            self._errors.append(
                f"{pdf.name}: {len(failed_pages)} página(s) falharam na extração"
            )
        else:
            state = STATE_COMPLETE

        # Remove os chunks anteriores deste arquivo antes de gravar os novos.
        self.store.delete_file(pdf.name)

        if chunks:
            vectors = self.embedder.encode([chunk.text for chunk in chunks])
            self.store.upsert(chunks, vectors)

        self._manifest[pdf.name] = _ManifestEntry(
            sha256=file_hash,
            state=state,
            chunks=len(chunks),
            indexed_at=_now(),
            failed_pages=failed_pages,
        )

    def status(self) -> IndexStatus:
        """Estado atual da indexação."""
        indexed = sum(
            1
            for entry in self._manifest.values()
            if entry.state in (STATE_COMPLETE, STATE_PARTIAL)
        )
        if self._indexing:
            state = "indexing"
        elif self._errors and indexed == 0:
            state = "error"
        else:
            state = "ready"
        return IndexStatus(
            state=state,
            pdfs_total=len(self.discover_pdfs()),
            pdfs_indexed=indexed,
            chunks=self.store.count(),
            model=self.settings.embedding_model,
            last_indexed_at=self._last_indexed_at,
            errors=list(self._errors),
            pages_processed=self._pages_processed,
            pages_total=self._pages_total,
            current_file=self._current_file,
        )


def _pdf_page_count(pdf: Path) -> int:
    """Número de páginas do PDF, ou 0 se não for possível determinar."""
    try:
        import pymupdf

        with pymupdf.open(pdf) as document:
            return len(document)
    except Exception:  # pragma: no cover - PDF malformado
        return 0


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()