# src/extractor/canon/canon_step.py

from pathlib import Path
from typing import Any

from .registry import EntityRegistry


class CanonStep:
    def __init__(
        self,
        registry_path: Path,
        enabled: bool = True,
        use_embeddings: bool = False,
    ) -> None:
        self.enabled = enabled
        self.registry_path = registry_path
        self.use_embeddings = use_embeddings
        self._registry: EntityRegistry | None = None

    def _get_registry(self) -> EntityRegistry:
        if self._registry is None:
            self._registry = EntityRegistry.load(self.registry_path)
        return self._registry

    def canonicalize(self, elements: list[Any], doc_id: str, doc_sha256: str) -> list[Any]:
        if not self.enabled:
            return elements
        registry = self._get_registry()
        for el in elements:
            for ent in getattr(el, "entities", None) or []:
                if ent.is_pronoun and not ent.canonical_text:
                    continue
                surface = ent.canonical_text or ent.text
                cid = registry.resolve(
                    surface=surface,
                    entity_type=ent.entity_type,
                    doc_id=doc_id,
                    chunk_id=getattr(el, "id", None),
                    use_embeddings=self.use_embeddings,
                )
                ent.canonical_id = cid
                if not ent.canonical_text:
                    ent.canonical_text = registry._entries[cid].canonical_text
        return elements

    def flush(self) -> None:
        if self._registry is not None:
            self._registry.save(self.registry_path)
