"""
cortex/core/entity_resolution.py — Entity Resolution Engine (Module 3 / M3).

Resolves raw text mentions (e.g., "Andrew", "Andy", "@andrew") to stable
canonical Entity IDs. This is critical for the claim deduplication — without
entity resolution, "Andrew prefers dark mode" and "Andy prefers dark mode"
would be stored as two separate claims about two different entities.

Resolution pipeline (cascading, stops at first confident match):
    1. Exact alias match (case-insensitive, zero LLM cost)
    2. Fuzzy alias match via rapidfuzz (token_sort_ratio ≥ 85%)
    3. LLM disambiguation for ambiguous matches (60% ≤ score < 85%)
    4. New entity creation when no match is found

Design decisions:
    - Uses rapidfuzz for O(n) fuzzy matching against all stored aliases.
    - LLM disambiguation is the expensive fallback — only triggered for
      ambiguous cases (multiple candidates within 5 points of each other).
    - Entity merging (suggest_merges) is a dream-cycle operation, not hot-path.

Dependencies:
    - cortex.types (Entity, EntityAlias, ResolvedEntity, etc.)
    - cortex.storage.adapter.StorageAdapter (M1) for alias/entity CRUD
    - cortex.llm.provider.LLMProvider for disambiguation fallback
    - rapidfuzz for fuzzy string matching

Data flow:
    CortexPlugin.remember()
      → ExtractionPipeline.extract() → ClaimExtraction with EntityMention list
      → EntityResolver.resolve(mention) → ResolvedEntity
          → entity_id used as FK in SemanticClaim.subject_entity_id
      → ReconciliationEngine.reconcile() (M4)
"""
from __future__ import annotations

import json
import logging
import math
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np
from rapidfuzz import fuzz, process as rf_process

from cortex.types import (
    CandidateEntity,
    Entity,
    EntityAlias,
    LLMProvider,
    MergeSuggestion,
    ResolvedEntity,
    StorageAdapter,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration defaults
# ---------------------------------------------------------------------------

DEFAULT_FUZZY_THRESHOLD = 85.0       # ≥ this → fuzzy match accepted
DEFAULT_AMBIGUOUS_LOWER = 60.0       # confidence range that triggers LLM disambiguation
DEFAULT_AMBIGUOUS_UPPER = 85.0       # same upper bound as fuzzy threshold
DEFAULT_MERGE_SIMILARITY = 80.0      # threshold for suggest_merges()
DEFAULT_CONFIDENCE_THRESHOLD = 0.5   # min normalised confidence (0-1) for merge suggestions
DEFAULT_EMBEDDING_PREFILTER = 0.60   # cosine-similarity lower bound for embedding pre-filter

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_DISAMBIGUATION_SYSTEM = """\
You are an entity disambiguation assistant for a knowledge-base system.
Given a text mention and a list of candidate entities (with their known aliases),
decide which candidate is the best match — or indicate that none match.

Respond with JSON only:
{
  "best_match_index": <int or null>,   // index into candidates array, null if no match
  "confidence": <float 0.0-1.0>,
  "reasoning": "<brief explanation>"
}
"""


def _build_disambiguation_user(
    mention: str,
    context: str,
    candidates: List[Tuple[Entity, str, float]],  # (entity, matched_alias, score)
) -> str:
    cand_lines = []
    for i, (entity, alias, score) in enumerate(candidates):
        cand_lines.append(
            f"[{i}] id={entity.id!r} canonical={entity.canonical_name!r} "
            f"type={entity.entity_type!r} alias_matched={alias!r} score={score:.1f}"
        )
    return (
        f"Mention: {mention!r}\n"
        f"Context: {context!r}\n\n"
        f"Candidates:\n" + "\n".join(cand_lines)
    )


# ---------------------------------------------------------------------------
# EntityResolver
# ---------------------------------------------------------------------------


class EntityResolver:
    """
    Resolve entity mentions to stable canonical entities.

    Usage::

        resolver = EntityResolver(storage=adapter, llm=provider,
                                  fuzzy_threshold=85.0)
        result = await resolver.resolve("Andrew", context="...", owner_id="eva")
    """

    def __init__(
        self,
        storage: StorageAdapter,
        llm: LLMProvider,
        fuzzy_threshold: float = DEFAULT_FUZZY_THRESHOLD,
        ambiguous_lower: float = DEFAULT_AMBIGUOUS_LOWER,
        ambiguous_upper: float = DEFAULT_AMBIGUOUS_UPPER,
    ) -> None:
        self.storage = storage
        self.llm = llm
        self.fuzzy_threshold = fuzzy_threshold
        self.ambiguous_lower = ambiguous_lower
        self.ambiguous_upper = ambiguous_upper

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def resolve(
        self,
        mention: str,
        context: str = "",
        owner_id: str = "default",
        entity_type: str = "unknown",
    ) -> ResolvedEntity:
        """Resolve a text mention to a stable canonical entity.

        Cascading resolution pipeline — stops at first confident match:
            1. Exact alias match (case-insensitive, free)
            2. Fuzzy alias match (score ≥ fuzzy_threshold, free)
            3. LLM disambiguation (ambiguous range, costs tokens)
            4. Create new entity (free)

        Args:
            mention: Raw text mention to resolve (e.g., "Andrew", "Andy").
            context: Optional surrounding text for disambiguation context.
            owner_id: Multi-tenancy scope key.
            entity_type: Classification hint from extraction (e.g. 'person',
                'technology', 'project', 'organization'). Used when creating
                new entities; existing entities may be upgraded from 'unknown'.

        Returns:
            ResolvedEntity with entity_id, canonical_name, confidence, and
            is_new flag indicating whether a new entity was created.

        Raises:
            ValueError: If mention is empty or whitespace-only.

        Note:
            The resolved entity_id is used as subject_entity_id or
            object_entity_id in SemanticClaim records (M4).
        """
        mention = mention.strip()
        if not mention:
            raise ValueError("mention must be a non-empty string")

        # --- Step 1: Exact match ---
        entity = await self.storage.find_entity_by_alias(mention, owner_id=owner_id)
        if entity is not None:
            logger.debug("resolve: exact match %r → %s", mention, entity.id)
            # Upgrade entity type from 'unknown' if extraction provides a real type
            effective_type = entity.entity_type
            if (
                entity.entity_type in ("unknown", "", None)
                and entity_type not in ("unknown", "", None)
            ):
                try:
                    await self.storage.update_entity(
                        entity.id, entity_type=entity_type, owner_id=owner_id,
                    )
                    effective_type = entity_type
                    logger.info(
                        "resolve: upgraded entity %r type %r → %r",
                        entity.canonical_name, entity.entity_type, entity_type,
                    )
                except Exception as exc:
                    logger.debug("resolve: entity type upgrade failed: %s", exc)
            return ResolvedEntity(
                entity_id=entity.id,
                canonical_name=entity.canonical_name,
                entity_type=effective_type,
                confidence=1.0,
                is_new=False,
                matched_alias=mention,
            )

        # --- Step 2 + 3: Load all aliases and fuzzy-score them ---
        # We load ALL aliases upfront (typically <1000 for most deployments)
        # and score them in-memory with rapidfuzz. This avoids N+1 queries.
        all_aliases = await self.storage.get_all_entity_aliases(owner_id=owner_id)

        if not all_aliases:
            # No entities at all — create one
            return await self._create_new(mention, owner_id, entity_type)

        alias_texts = [a.alias for a in all_aliases]
        # rapidfuzz.process.extract returns ALL matches above score_cutoff.
        # token_sort_ratio normalizes word order: "John Smith" ≈ "Smith John"
        scored = rf_process.extract(
            mention,
            alias_texts,
            scorer=fuzz.token_sort_ratio,
            limit=None,
            score_cutoff=self.ambiguous_lower,  # 60% — below this, don't even consider
        )
        # scored = [(alias_text, score, index), ...]

        if not scored:
            return await self._create_new(mention, owner_id, entity_type)

        # Collapse alias-level scores to entity-level scores.
        # Multiple aliases can point to the same entity — keep only the best.
        entity_best: dict[str, Tuple[float, str]] = {}  # entity_id → (score, alias)
        alias_map = {a.alias: a for a in all_aliases}

        for alias_text, score, _ in scored:
            alias_obj = alias_map.get(alias_text)
            if alias_obj is None:
                continue
            eid = alias_obj.entity_id
            if eid not in entity_best or score > entity_best[eid][0]:
                entity_best[eid] = (score, alias_text)

        # Sort by score descending
        ranked = sorted(entity_best.items(), key=lambda x: x[1][0], reverse=True)
        top_score = ranked[0][1][0] if ranked else 0.0

        if top_score >= self.fuzzy_threshold:
            # Clear winner — accept if it's unambiguous (no other entity within 5 pts).
            # The 5-point gap prevents auto-accepting when two entities score similarly
            # (e.g., "Alex" matching both "Alexander" at 90 and "Alexa" at 87).
            top_eid, (top_sc, top_alias) = ranked[0]
            ambiguous = len(ranked) > 1 and (ranked[0][1][0] - ranked[1][1][0]) < 5.0

            if not ambiguous:
                entity = await self.storage.get_entity_by_id(top_eid)
                if entity:
                    logger.debug(
                        "resolve: fuzzy match %r → %s (score=%.1f)",
                        mention, entity.id, top_sc,
                    )
                    return ResolvedEntity(
                        entity_id=entity.id,
                        canonical_name=entity.canonical_name,
                        entity_type=entity.entity_type,
                        confidence=top_sc / 100.0,
                        is_new=False,
                        matched_alias=top_alias,
                    )

        # --- Step 3: LLM disambiguation for ambiguous / mid-range matches ---
        candidates: List[Tuple[Entity, str, float]] = []
        for eid, (sc, alias_text) in ranked[:5]:  # limit to top-5
            e = await self.storage.get_entity_by_id(eid)
            if e:
                candidates.append((e, alias_text, sc))

        if candidates:
            resolved = await self._llm_disambiguate(mention, context, candidates)
            if resolved is not None:
                return resolved

        # --- Step 4: No match — create new entity ---
        return await self._create_new(mention, owner_id, entity_type)

    async def resolve_batch(
        self,
        mentions: List[CandidateEntity],
        context: str = "",
        owner_id: str = "default",
    ) -> List[ResolvedEntity]:
        """Resolve multiple entity mentions in sequence.

        Currently resolves sequentially (not parallel) because:
        - Alias lookups are local SQLite queries (fast)
        - LLM disambiguation is rare (only for ambiguous cases)
        - Sequential resolution lets later mentions benefit from entities
          created by earlier mentions in the same batch

        Args:
            mentions: List of CandidateEntity objects to resolve.
            context: Shared context string for disambiguation.
            owner_id: Multi-tenancy scope key.

        Returns:
            List of ResolvedEntity in the same order as input mentions.
        """
        results: List[ResolvedEntity] = []
        for candidate in mentions:
            resolved = await self.resolve(
                mention=candidate.name,
                context=context,
                owner_id=owner_id,
            )
            results.append(resolved)
        return results

    async def create_entity(
        self,
        canonical_name: str,
        entity_type: str = "unknown",
        owner_id: str = "default",
    ) -> Entity:
        """Create a new entity and register its canonical name as the first alias.

        Args:
            canonical_name: The primary name for this entity.
            entity_type: Classification — 'person', 'project', 'concept', etc.
            owner_id: Multi-tenancy scope key.

        Returns:
            Newly created Entity with its first alias registered.
        """
        entity = await self.storage.create_entity(
            canonical_name=canonical_name,
            entity_type=entity_type,
            owner_id=owner_id,
        )
        await self.storage.add_entity_alias(
            entity_id=entity.id,
            alias=canonical_name,
            alias_type="name",
            source="operator",
            confidence=1.0,
        )
        logger.info("create_entity: created %r (%s) → %s", canonical_name, entity_type, entity.id)
        return entity

    async def add_alias(
        self,
        entity_id: str,
        alias: str,
        alias_type: str = "name",
        source: str = "extracted",
        confidence: float = 1.0,
    ) -> EntityAlias:
        """
        Add a new alias to an existing entity.
        """
        result = await self.storage.add_entity_alias(
            entity_id=entity_id,
            alias=alias,
            alias_type=alias_type,
            source=source,
            confidence=confidence,
        )
        logger.info("add_alias: %r → entity %s", alias, entity_id)
        return result

    async def merge_entities(
        self,
        keep_id: str,
        merge_id: str,
    ) -> Entity:
        """Merge two entities that refer to the same real-world thing.

        All aliases and semantic claim FK references (subject_entity_id,
        object_entity_id) are atomically reassigned from merge_id → keep_id.
        The merge_id entity is then deleted. The StorageAdapter (M1) handles
        this in a single DB transaction to ensure consistency.

        Args:
            keep_id: Entity to keep (the canonical one).
            merge_id: Entity to absorb and delete.

        Returns:
            The surviving Entity (keep_id) with updated alias count.

        Raises:
            ValueError: If keep_id == merge_id (can't merge entity with itself).

        Note:
            This is a dream-cycle operation — called from suggest_merges()
            review, not on the hot path.
        """
        if keep_id == merge_id:
            raise ValueError("keep_id and merge_id must be different entities")

        entity = await self.storage.merge_entities(keep_id=keep_id, merge_id=merge_id)
        logger.info("merge_entities: %s ← %s (kept: %r)", keep_id, merge_id, entity.canonical_name)
        return entity

    async def suggest_merges(
        self,
        owner_id: str = "default",
        similarity_threshold: float = DEFAULT_MERGE_SIMILARITY,
        confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
        cross_session_owner_ids: Optional[List[str]] = None,
        use_embedding_prefilter: bool = True,
        embedding_prefilter_threshold: float = DEFAULT_EMBEDDING_PREFILTER,
    ) -> List[MergeSuggestion]:
        """Find entity pairs that might be duplicates, for operator/dream-cycle review.

        Algorithm:
            1. Load aliases for ``owner_id`` (and optionally additional owner IDs
               for cross-session co-reference detection).
            2. *Embedding pre-filter* (optional, enabled by default): embed all
               entity canonical names and compute cosine-similarity. Only entity
               pairs whose embedding cosine-similarity is ≥ ``embedding_prefilter_threshold``
               are passed to the O(n²) fuzzy-string step.  This cuts the number
               of expensive string comparisons dramatically for large entity sets.
            3. Fuzzy-string comparison (token_sort_ratio) on the pre-filtered pairs.
               Only the best alias-pair score between each entity pair is kept.
            4. Results are filtered by both ``similarity_threshold`` (raw 0-100
               fuzzy score) and ``confidence_threshold`` (normalised 0-1 score).
            5. Cross-session results are tagged in the ``reason`` field so operators
               can distinguish intra-session dedup from cross-session co-reference
               linking.

        This is intentionally a dream-cycle operation (batch, offline) because:
        - O(n²) is too expensive for the hot-path
        - Merges need human review before executing

        Args:
            owner_id: Primary multi-tenancy scope key.
            similarity_threshold: Minimum fuzzy score (0-100) to suggest a merge.
            confidence_threshold: Minimum normalised confidence (0.0-1.0) to
                include in results.  Use 0.0 to disable.
            cross_session_owner_ids: Optional list of additional owner IDs whose
                entities should be compared against ``owner_id`` entities.  This
                enables cross-session co-reference linking where the same real-world
                entity (e.g., "Andrew") appears in multiple sessions under slightly
                different surface forms.
            use_embedding_prefilter: When True (default), embed entity canonical
                names and skip fuzzy comparison for pairs with low cosine similarity.
                Falls back gracefully to full O(n²) if the LLM embed call fails.
            embedding_prefilter_threshold: Cosine-similarity lower bound (0-1) for
                the embedding pre-filter.  Pairs below this threshold are skipped.

        Returns:
            List of MergeSuggestion sorted by similarity (highest first).
        """
        # ------------------------------------------------------------------
        # 1. Load aliases — primary owner + optional cross-session owners
        # ------------------------------------------------------------------
        all_aliases = await self.storage.get_all_entity_aliases(owner_id=owner_id)

        # Map each entity_id → owner label for cross-session tagging
        entity_owner: Dict[str, str] = {a.entity_id: owner_id for a in all_aliases}

        if cross_session_owner_ids:
            for xid in cross_session_owner_ids:
                if xid == owner_id:
                    continue
                extra = await self.storage.get_all_entity_aliases(owner_id=xid)
                for a in extra:
                    entity_owner.setdefault(a.entity_id, xid)
                all_aliases = all_aliases + extra

        # Group aliases by entity_id
        entity_aliases: Dict[str, List[str]] = {}
        for a in all_aliases:
            entity_aliases.setdefault(a.entity_id, []).append(a.alias)

        entity_ids = list(entity_aliases.keys())
        n = len(entity_ids)

        if n < 2:
            return []

        # ------------------------------------------------------------------
        # 2. Embedding pre-filter (optional)
        # ------------------------------------------------------------------
        # Build a set of entity-pair indices that we will run fuzzy matching on.
        # By default, ALL pairs — override below if embeddings succeed.
        pairs_to_check: set[Tuple[int, int]] = {
            (i, j) for i in range(n) for j in range(i + 1, n)
        }

        if use_embedding_prefilter:
            try:
                # Use canonical name of first alias as representative text
                canonical_names = [entity_aliases[eid][0] for eid in entity_ids]
                embeddings_raw = await self.llm.embed(canonical_names)
                # Normalise to unit vectors for cosine similarity via dot product
                emb = np.array(embeddings_raw, dtype=np.float32)
                norms = np.linalg.norm(emb, axis=1, keepdims=True)
                norms = np.where(norms == 0, 1.0, norms)
                emb_unit = emb / norms
                # Compute full cosine-similarity matrix in one BLAS call
                cos_sim = emb_unit @ emb_unit.T  # shape (n, n)
                pairs_to_check = {
                    (i, j)
                    for i in range(n)
                    for j in range(i + 1, n)
                    if cos_sim[i, j] >= embedding_prefilter_threshold
                }
                logger.debug(
                    "suggest_merges: embedding pre-filter reduced pairs from %d → %d "
                    "(threshold=%.2f)",
                    n * (n - 1) // 2,
                    len(pairs_to_check),
                    embedding_prefilter_threshold,
                )
            except Exception as exc:
                logger.warning(
                    "suggest_merges: embedding pre-filter failed (%s) — "
                    "falling back to full O(n²) comparison",
                    exc,
                )
                # pairs_to_check already set to full set above

        # ------------------------------------------------------------------
        # 3. Fuzzy string comparison on pre-filtered pairs
        # ------------------------------------------------------------------
        suggestions: List[MergeSuggestion] = []
        seen: set[frozenset] = set()

        for i, j in pairs_to_check:
            eid_a = entity_ids[i]
            eid_b = entity_ids[j]
            pair = frozenset([eid_a, eid_b])
            if pair in seen:
                continue
            seen.add(pair)

            aliases_a = entity_aliases[eid_a]
            aliases_b = entity_aliases[eid_b]

            # Best alias-pair fuzzy score between the two entities
            best_score = 0.0
            best_pair: Tuple[str, str] = (aliases_a[0], aliases_b[0])
            for a in aliases_a:
                for b in aliases_b:
                    score = fuzz.token_sort_ratio(a, b)
                    if score > best_score:
                        best_score = score
                        best_pair = (a, b)

            normalised = best_score / 100.0
            if best_score < similarity_threshold:
                continue
            if normalised < confidence_threshold:
                continue

            ea = await self.storage.get_entity_by_id(eid_a)
            eb = await self.storage.get_entity_by_id(eid_b)
            if not (ea and eb):
                continue

            # Determine whether this is a cross-session match
            owner_a = entity_owner.get(eid_a, owner_id)
            owner_b = entity_owner.get(eid_b, owner_id)
            is_cross = owner_a != owner_b
            if is_cross:
                reason = (
                    f"Cross-session co-reference: alias {best_pair[0]!r} (session {owner_a!r}) "
                    f"≈ {best_pair[1]!r} (session {owner_b!r}) — "
                    f"fuzzy {best_score:.1f}/100"
                )
            else:
                reason = (
                    f"Alias fuzzy similarity {best_score:.1f}/100 "
                    f"({best_pair[0]!r} ≈ {best_pair[1]!r})"
                )

            suggestions.append(
                MergeSuggestion(
                    keep_id=eid_a,
                    merge_id=eid_b,
                    keep_name=ea.canonical_name,
                    merge_name=eb.canonical_name,
                    similarity_score=normalised,
                    reason=reason,
                )
            )

        suggestions.sort(key=lambda s: s.similarity_score, reverse=True)
        logger.info(
            "suggest_merges: %d suggestions for owner=%r (threshold=%.1f, "
            "cross_session=%s, embedding_prefilter=%s)",
            len(suggestions),
            owner_id,
            similarity_threshold,
            bool(cross_session_owner_ids),
            use_embedding_prefilter,
        )
        return suggestions

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    async def _create_new(
        self, mention: str, owner_id: str, entity_type: str = "unknown",
    ) -> ResolvedEntity:
        """Create a new entity from a mention and return a ResolvedEntity."""
        entity = await self.create_entity(
            canonical_name=mention,
            entity_type=entity_type or "unknown",
            owner_id=owner_id,
        )
        return ResolvedEntity(
            entity_id=entity.id,
            canonical_name=entity.canonical_name,
            entity_type=entity.entity_type,
            confidence=1.0,
            is_new=True,
            matched_alias=None,
        )

    async def _llm_disambiguate(
        self,
        mention: str,
        context: str,
        candidates: List[Tuple[Entity, str, float]],
    ) -> Optional[ResolvedEntity]:
        """
        Ask the LLM to pick the best candidate entity for a mention.
        Returns None if the LLM cannot confidently match any candidate.
        """
        user_msg = _build_disambiguation_user(mention, context, candidates)
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "entity_resolution")
            _resp = await self.llm.complete(
                system=_DISAMBIGUATION_SYSTEM,
                user=user_msg,
                response_format="json",
                **({"model": _routed_model} if _routed_model else {}),
            )
            raw = _resp.text if hasattr(_resp, "text") else str(_resp)
            data = json.loads(raw)
        except Exception as exc:
            logger.warning("_llm_disambiguate: LLM call failed: %s", exc)
            return None

        idx = data.get("best_match_index")
        confidence = float(data.get("confidence", 0.0))

        if idx is None or not (0 <= idx < len(candidates)):
            logger.debug("_llm_disambiguate: LLM chose no match for %r", mention)
            return None

        entity, matched_alias, _ = candidates[idx]
        logger.debug(
            "_llm_disambiguate: %r → %s (confidence=%.2f)",
            mention, entity.id, confidence,
        )
        return ResolvedEntity(
            entity_id=entity.id,
            canonical_name=entity.canonical_name,
            entity_type=entity.entity_type,
            confidence=confidence,
            is_new=False,
            matched_alias=matched_alias,
        )
