"""
extractor.analytics.entity_index — In-memory entity index backed by JSON.

Reads extracted JSONL files, canonicalizes entity surface forms, and exposes
look-up / top-N / mention-list APIs. Read-only after build.
"""

from __future__ import annotations

import json
import os
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from .canonicalizer import Canonicalizer, derive_canonical_id

if TYPE_CHECKING:
    pass

_MAX_MENTIONS = 500  # cap per canonical entity to bound memory


@dataclass
class EntityMention:
    entity_type: str
    text: str  # surface form
    score: float
    source_filename: str
    source_sha256: str
    element_id: str
    chunk_id: str | None = None
    page: int | None = None
    section_path: list[str] = field(default_factory=list)
    canonical_id: str | None = None
    canonical_text: str | None = None

    def to_dict(self) -> dict:
        return {
            "entity_type": self.entity_type,
            "text": self.text,
            "score": self.score,
            "source_filename": self.source_filename,
            "source_sha256": self.source_sha256,
            "element_id": self.element_id,
            "chunk_id": self.chunk_id,
            "page": self.page,
            "section_path": self.section_path,
            "canonical_id": self.canonical_id,
            "canonical_text": self.canonical_text,
        }

    @classmethod
    def from_dict(cls, d: dict) -> EntityMention:
        return cls(
            entity_type=d["entity_type"],
            text=d["text"],
            score=d.get("score", 1.0),
            source_filename=d["source_filename"],
            source_sha256=d.get("source_sha256", ""),
            element_id=d.get("element_id", ""),
            chunk_id=d.get("chunk_id"),
            page=d.get("page"),
            section_path=d.get("section_path", []),
            canonical_id=d.get("canonical_id"),
            canonical_text=d.get("canonical_text"),
        )


@dataclass
class CanonicalEntity:
    canonical_id: str
    canonical_text: str
    entity_type: str
    aliases: list[str] = field(default_factory=list)
    mention_count: int = 0
    doc_count: int = 0
    mentions: list[EntityMention] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "canonical_id": self.canonical_id,
            "canonical_text": self.canonical_text,
            "entity_type": self.entity_type,
            "aliases": self.aliases,
            "mention_count": self.mention_count,
            "doc_count": self.doc_count,
            "mentions": [m.to_dict() for m in self.mentions],
        }

    @classmethod
    def from_dict(cls, d: dict) -> CanonicalEntity:
        obj = cls(
            canonical_id=d["canonical_id"],
            canonical_text=d["canonical_text"],
            entity_type=d["entity_type"],
            aliases=d.get("aliases", []),
            mention_count=d.get("mention_count", 0),
            doc_count=d.get("doc_count", 0),
        )
        obj.mentions = [EntityMention.from_dict(m) for m in d.get("mentions", [])]
        return obj


class EntityIndex:
    """Read-only after build. JSON-backed."""

    def __init__(self) -> None:
        self._entities: dict[str, CanonicalEntity] = {}  # canonical_id → CanonicalEntity
        self._alias_index: dict[str, str] = {}  # "entity_type:surface.lower()" → canonical_id

    # ------------------------------------------------------------------
    # Build
    # ------------------------------------------------------------------

    @classmethod
    def build_from_jsonl(
        cls,
        paths: list[str | Path],
        canonicalizer: Canonicalizer | None = None,
    ) -> EntityIndex:
        """Read extracted JSONL files and build entity index."""
        if canonicalizer is None:
            canonicalizer = Canonicalizer()

        index = cls()

        for path in paths:
            path = Path(path)
            source_filename = path.name
            try:
                lines = path.read_text(encoding="utf-8").splitlines()
            except OSError:
                continue

            for line in lines:
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue

                record_type = record.get("type", "")
                # Skip envelope / control lines
                if record_type in ("stream_start", "stream_end", "manifest"):
                    continue

                entities = record.get("entities") or []
                if not entities:
                    continue

                source_sha256 = record.get("sha256") or record.get("source_sha256") or ""
                element_id = record.get("element_id") or record.get("id") or ""
                chunk_id = record.get("chunk_id")
                page = record.get("page")
                section_path = record.get("section_path") or []

                for ent in entities:
                    if not isinstance(ent, dict):
                        continue
                    # Skip pronouns
                    if ent.get("is_pronoun"):
                        continue

                    entity_type = ent.get("entity_type") or ent.get("type") or ""
                    surface = ent.get("text") or ent.get("surface") or ""
                    if not surface or not entity_type:
                        continue

                    score = float(ent.get("score", 1.0))

                    # Use pre-assigned canonical_id if present, else derive
                    canonical_id = ent.get("canonical_id") or derive_canonical_id(
                        entity_type, surface
                    )

                    # Check if we can merge with an existing entity via alias lookup
                    alias_key = f"{entity_type.lower()}:{surface.lower()}"
                    existing_id = index._alias_index.get(alias_key)

                    if existing_id is None:
                        # Try fuzzy merge against existing entities of same type
                        merged_id = index._try_fuzzy_merge(surface, entity_type, canonicalizer)
                        if merged_id:
                            existing_id = merged_id

                    if existing_id:
                        canonical_id = existing_id
                        canonical_entity = index._entities[canonical_id]
                        canonical_text = canonical_entity.canonical_text
                        # Add alias if new surface form
                        if surface not in canonical_entity.aliases:
                            canonical_entity.aliases.append(surface)
                    else:
                        # New entity
                        from .canonicalizer import normalize_surface

                        canonical_text = normalize_surface(surface, entity_type)
                        canonical_entity = CanonicalEntity(
                            canonical_id=canonical_id,
                            canonical_text=canonical_text,
                            entity_type=entity_type,
                            aliases=[surface],
                        )
                        index._entities[canonical_id] = canonical_entity

                    # Register alias
                    index._alias_index[alias_key] = canonical_id

                    # Build mention
                    mention = EntityMention(
                        entity_type=entity_type,
                        text=surface,
                        score=score,
                        source_filename=source_filename,
                        source_sha256=source_sha256,
                        element_id=element_id,
                        chunk_id=chunk_id,
                        page=page,
                        section_path=list(section_path),
                        canonical_id=canonical_id,
                        canonical_text=canonical_entity.canonical_text,
                    )

                    canonical_entity.mention_count += 1

                    # Track unique docs
                    existing_docs = {m.source_filename for m in canonical_entity.mentions}
                    if source_filename not in existing_docs:
                        canonical_entity.doc_count += 1

                    # Cap mentions
                    if len(canonical_entity.mentions) < _MAX_MENTIONS:
                        canonical_entity.mentions.append(mention)

        return index

    def _try_fuzzy_merge(
        self,
        surface: str,
        entity_type: str,
        canonicalizer: Canonicalizer,
    ) -> str | None:
        """Check existing entities for a merge candidate. Returns canonical_id or None."""
        for cid, entity in self._entities.items():
            if entity.entity_type.lower() != entity_type.lower():
                continue
            merge, _decision, _conf = canonicalizer.should_merge(
                surface, entity.canonical_text, entity_type
            )
            if merge:
                return cid
        return None

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str | Path) -> None:
        """Atomic JSON write."""
        path = Path(path)
        data = {
            "entities": {cid: e.to_dict() for cid, e in self._entities.items()},
            "alias_index": self._alias_index,
        }
        # Atomic write via temp file
        dir_ = path.parent
        dir_.mkdir(parents=True, exist_ok=True)
        fd, tmp = tempfile.mkstemp(dir=dir_, suffix=".tmp")
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            os.replace(tmp, path)
        except Exception:
            import contextlib

            with contextlib.suppress(OSError):
                os.unlink(tmp)
            raise

    @classmethod
    def load(cls, path: str | Path) -> EntityIndex:
        path = Path(path)
        data = json.loads(path.read_text(encoding="utf-8"))
        index = cls()
        for cid, edict in data.get("entities", {}).items():
            index._entities[cid] = CanonicalEntity.from_dict(edict)
        index._alias_index = data.get("alias_index", {})
        return index

    # ------------------------------------------------------------------
    # Query API
    # ------------------------------------------------------------------

    def get(self, canonical_id: str) -> CanonicalEntity | None:
        return self._entities.get(canonical_id)

    def resolve(self, surface: str, entity_type: str) -> str | None:
        """Look up canonical_id for a surface form. Returns None if not found."""
        key = f"{entity_type}:{surface.lower()}"
        return self._alias_index.get(key)

    def top_entities(
        self,
        entity_type: str | None = None,
        n: int = 20,
    ) -> list[CanonicalEntity]:
        """Return top N entities by mention_count, optionally filtered by type."""
        entities = list(self._entities.values())
        if entity_type is not None:
            entities = [e for e in entities if e.entity_type.lower() == entity_type.lower()]
        entities.sort(key=lambda e: e.mention_count, reverse=True)
        return entities[:n]

    def mentions_for(self, canonical_id: str) -> list[EntityMention]:
        entity = self._entities.get(canonical_id)
        if entity is None:
            return []
        return list(entity.mentions)

    def all_entity_types(self) -> list[str]:
        seen: set[str] = set()
        result: list[str] = []
        for e in self._entities.values():
            if e.entity_type not in seen:
                seen.add(e.entity_type)
                result.append(e.entity_type)
        return sorted(result)
