# src/extractor/canon/registry.py

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


@dataclass
class RegistryEntry:
    canonical_id: str
    canonical_text: str
    entity_type: str
    aliases: list[str] = field(default_factory=list)
    first_seen_doc: str = ""
    doc_count: int = 1
    mention_count: int = 1
    chunk_ids: list[str] = field(default_factory=list)
    pending_review: list[dict] = field(default_factory=list)
    merge_log: list[dict] = field(default_factory=list)


class EntityRegistry:
    """
    JSON-backed canonical entity store.
    No SQLite. Thread-safe reads; write lock for updates.
    """

    def __init__(self) -> None:
        self._entries: dict[str, RegistryEntry] = {}
        self._alias_index: dict[str, str] = {}  # "entity_type:normalized" -> canonical_id
        self._lock = threading.Lock()

    @classmethod
    def load(cls, path: Path) -> "EntityRegistry":
        reg = cls()
        if path.exists():
            raw = json.loads(path.read_text(encoding="utf-8"))
            for k, v in raw.items():
                if k.startswith("_"):
                    continue
                entry = RegistryEntry(
                    **{fk: v[fk] for fk in RegistryEntry.__dataclass_fields__ if fk in v}
                )
                reg._entries[k] = entry
                reg._index_entry(entry)
        return reg

    def save(self, path: Path) -> None:
        with self._lock:
            out = {
                "_meta": {
                    "schema_version": "1.0.0",
                    "updated_at": datetime.now(timezone.utc).isoformat(),
                    "total_entities": len(self._entries),
                },
                **{k: asdict(v) for k, v in self._entries.items()},
            }
            tmp = path.with_suffix(".tmp")
            tmp.write_text(json.dumps(out, indent=2, ensure_ascii=False), encoding="utf-8")
            tmp.replace(path)

    def resolve(
        self,
        surface: str,
        entity_type: str,
        doc_id: str = "",
        chunk_id: str | None = None,
        use_embeddings: bool = False,
    ) -> str:
        """
        Look up or create a canonical_id for the given surface form.
        Updates alias lists and mention counts. Returns canonical_id.
        """
        from .clusterer import AUTO_MERGE_THRESHOLD, REVIEW_THRESHOLD, similarity_score
        from .normalizer import normalize_surface

        with self._lock:
            normalized = normalize_surface(surface, entity_type)
            alias_key = f"{entity_type}:{normalized.lower()}"

            # Exact alias match
            if alias_key in self._alias_index:
                cid = self._alias_index[alias_key]
                self._record_mention(cid, surface, doc_id, chunk_id)
                return cid

            # String similarity search
            best_score = 0.0
            best_cid = None
            for cid, entry in self._entries.items():
                if entry.entity_type != entity_type:
                    continue
                score = max(
                    similarity_score(normalized, normalize_surface(a, entity_type))
                    for a in [entry.canonical_text] + entry.aliases
                )
                if score > best_score:
                    best_score = score
                    best_cid = cid

            if best_score >= AUTO_MERGE_THRESHOLD and best_cid:
                self._add_alias(
                    best_cid, surface, normalized, entity_type, "similarity", best_score / 100
                )
                self._record_mention(best_cid, surface, doc_id, chunk_id)
                return best_cid

            if best_score >= REVIEW_THRESHOLD and best_cid:
                if use_embeddings:
                    from .embedder import EMBEDDING_MERGE_THRESHOLD, embedding_similarity

                    sim = embedding_similarity(normalized, self._entries[best_cid].canonical_text)
                    if sim >= EMBEDDING_MERGE_THRESHOLD:
                        self._add_alias(
                            best_cid, surface, normalized, entity_type, "embedding", sim
                        )
                        self._record_mention(best_cid, surface, doc_id, chunk_id)
                        return best_cid

            # New entity
            return self._create_entry(surface, normalized, entity_type, doc_id, chunk_id)

    def _create_entry(
        self,
        surface: str,
        normalized: str,
        entity_type: str,
        doc_id: str,
        chunk_id: str | None,
    ) -> str:
        from .normalizer import derive_canonical_id, derive_canonical_id_with_disambiguation

        base_id = derive_canonical_id(entity_type, normalized)
        if base_id in self._entries:
            existing = self._entries[base_id].canonical_text
            if existing.lower() != normalized.lower():
                base_id = derive_canonical_id_with_disambiguation(entity_type, normalized, existing)

        entry = RegistryEntry(
            canonical_id=base_id,
            canonical_text=normalized,
            entity_type=entity_type,
            aliases=[surface] if surface != normalized else [],
            first_seen_doc=doc_id,
            chunk_ids=[chunk_id] if chunk_id else [],
        )
        self._entries[base_id] = entry
        self._index_entry(entry)
        return base_id

    def _add_alias(
        self,
        cid: str,
        surface: str,
        normalized: str,
        entity_type: str,
        decision: str,
        confidence: float,
    ) -> None:
        entry = self._entries[cid]
        if surface not in entry.aliases:
            entry.aliases.append(surface)
            entry.merge_log.append(
                {"merged_from": surface, "decision": decision, "confidence": round(confidence, 4)}
            )
        self._alias_index[f"{entity_type}:{normalized.lower()}"] = cid
        self._alias_index[f"{entity_type}:{surface.lower()}"] = cid

    def _record_mention(self, cid: str, surface: str, doc_id: str, chunk_id: str | None) -> None:
        entry = self._entries[cid]
        entry.mention_count += 1
        if chunk_id and chunk_id not in entry.chunk_ids:
            entry.chunk_ids = (entry.chunk_ids + [chunk_id])[:1000]

    def export_graph(self, path: Path) -> None:
        """Export entity graph as JSON with nodes and edges."""
        nodes = []
        for entry in self._entries.values():
            nodes.append(
                {
                    "canonical_id": entry.canonical_id,
                    "text": entry.canonical_text,
                    "type": entry.entity_type,
                    "aliases": list(entry.aliases),
                    "doc_count": entry.doc_count,
                    "mention_count": entry.mention_count,
                }
            )
        # Edges: co-occurrence edges require ChunkRecord access; emit empty list for now
        edges: list[dict] = []
        graph = {"nodes": nodes, "edges": edges}
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(graph, ensure_ascii=False, indent=2), encoding="utf-8")
        tmp.replace(path)

    def _index_entry(self, entry: RegistryEntry) -> None:
        key = f"{entry.entity_type}:{entry.canonical_text.lower()}"
        self._alias_index[key] = entry.canonical_id
        for alias in entry.aliases:
            self._alias_index[f"{entry.entity_type}:{alias.lower()}"] = entry.canonical_id
