"""Entity resolution — 3-phase: exact → fuzzy (string) → LLM.

Ported from advisor_api entity_resolution_v2.py.
Key difference: fuzzy phase uses difflib.SequenceMatcher (no pgvector dependency).
LLM phase uses injected LLMProvider.
"""

import json
import logging
import re
import unicodedata
import uuid as uuid_mod
from dataclasses import dataclass
from difflib import SequenceMatcher

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from arandu._helpers import cosine_similarity, strip_markdown_fences
from arandu.constants import ExtractedEntity, ExtractedFact
from arandu.models import MemoryEntity, MemoryEntityAlias
from arandu.protocols import EmbeddingProvider, LLMProvider

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
FUZZY_DIRECT_THRESHOLD = 0.85  # default; overridden by config.fuzzy_threshold
FUZZY_LLM_THRESHOLD = 0.5

# Multilingual self-references (EN, PT, ES, FR, DE).
# Best-effort optimization — the LLM extraction handles the rest.
_SELF_REFERENCES = frozenset(
    {
        # English
        "user",
        "me",
        "i",
        "myself",
        # Portuguese
        "eu",
        "mim",
        # Spanish
        "yo",
        "mi",
        # French
        "je",
        "moi",
        # German
        "ich",
    }
)

# Multilingual relationship hints (EN, PT, ES).
# Used to infer entity_type=person when a bare relationship term appears.
_RELATIONSHIP_HINTS = frozenset(
    {
        # English
        "girlfriend",
        "boyfriend",
        "brother",
        "sister",
        "friend",
        "mother",
        "father",
        "son",
        "daughter",
        "wife",
        "husband",
        "boss",
        "colleague",
        "manager",
        "ex",
        "cousin",
        "uncle",
        "aunt",
        "grandmother",
        "grandfather",
        "mother-in-law",
        "father-in-law",
        # Portuguese
        "namorada",
        "namorado",
        "irmao",
        "irma",
        "irmã",
        "irmão",
        "amigo",
        "amiga",
        "mae",
        "mãe",
        "pai",
        "filho",
        "filha",
        "esposa",
        "marido",
        "chefe",
        "colega",
        "gerente",
        "prima",
        "primo",
        "tio",
        "tia",
        "avo",
        "avó",
        "avô",
        "sogra",
        "sogro",
        # Spanish
        "novia",
        "novio",
        "hermano",
        "hermana",
        "madre",
        "padre",
        "hijo",
        "hija",
        "jefe",
        "abuela",
        "abuelo",
        "suegra",
        "suegro",
    }
)

_HINT_RE = re.compile(r"^(.+?)\s*\((.+?)\)\s*$")


# ---------------------------------------------------------------------------
# Result type
# ---------------------------------------------------------------------------
@dataclass
class ResolvedEntity:
    """Result of resolving a raw entity name to a canonical entity.

    Attributes:
        entity_type: Category of entity (e.g. "person", "place").
        entity_key: Canonical key in the format "type::normalized_name".
        display_name: Human-readable entity name.
        is_new: Whether a new entity was created during resolution.
        resolution_method: How the entity was resolved (e.g. "exact", "fuzzy", "llm", "new").
        fuzzy_score: Similarity score when resolved via fuzzy matching, None otherwise.
    """

    entity_type: str
    entity_key: str
    display_name: str
    is_new: bool
    resolution_method: str
    fuzzy_score: float | None = None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _slugify(name: str) -> str:
    """Convert display name to entity key slug."""
    name = name.lower().strip()
    name = unicodedata.normalize("NFKD", name)
    name = re.sub(r"[\u0300-\u036f]", "", name)
    name = re.sub(r"[^a-z0-9\s]", "", name)
    name = re.sub(r"\s+", "_", name).strip("_")
    return name or "unknown"


def _normalize_name(name: str) -> str:
    """Normalize entity name for exact matching."""
    name = name.lower().strip()
    name = unicodedata.normalize("NFKD", name)
    name = re.sub(r"[\u0300-\u036f]", "", name)
    return name


def _parse_subject_hint(raw_name: str) -> tuple[str, str | None]:
    """Extract (clean_name, hint) from a subject with optional parenthesised hint."""
    if not raw_name:
        return (raw_name, None)
    m = _HINT_RE.match(raw_name)
    if m:
        return (m.group(1).strip(), m.group(2).strip())
    return (raw_name, None)


def _hint_implies_person(hint: str | None) -> bool:
    """Return True if the hint contains a recognised relationship pattern."""
    if not hint:
        return False
    normalized = unicodedata.normalize("NFKD", hint.lower())
    normalized = re.sub(r"[\u0300-\u036f]", "", normalized)
    tokens = normalized.split()
    return any(t in _RELATIONSHIP_HINTS for t in tokens)


# ---------------------------------------------------------------------------
# Cache loading
# ---------------------------------------------------------------------------
async def _load_entity_cache(
    session: AsyncSession,
    user_id: str,
) -> tuple[list[MemoryEntity], dict[str, str]]:
    """Load existing entities + alias mapping for this user."""
    stmt = select(MemoryEntity).where(MemoryEntity.user_id == user_id, MemoryEntity.is_active.is_(True))
    result = await session.execute(stmt)
    entities = list(result.scalars().all())

    alias_stmt = select(MemoryEntityAlias).where(MemoryEntityAlias.user_id == user_id)
    alias_result = await session.execute(alias_stmt)
    aliases = alias_result.scalars().all()
    alias_map = {_normalize_name(a.alias): a.canonical_entity_key for a in aliases}

    return entities, alias_map


# ---------------------------------------------------------------------------
# Phase 1: Exact match
# ---------------------------------------------------------------------------
def _exact_match(
    name_normalized: str,
    existing_entities: list[MemoryEntity],
    alias_cache: dict[str, str],
) -> MemoryEntity | None:
    """Phase 1: exact match against alias cache, entity slugs, display names."""
    # (a) Check alias cache
    if name_normalized in alias_cache:
        target_key = alias_cache[name_normalized]
        for ent in existing_entities:
            if ent.canonical_key == target_key:
                return ent

    # (b) Check entity canonical_key slugs
    for ent in existing_entities:
        slug = ent.canonical_key.split(":", 1)[1] if ":" in ent.canonical_key else ent.canonical_key
        if _normalize_name(slug) == name_normalized:
            return ent

    # (c) Check display names
    for ent in existing_entities:
        if ent.display_name and _normalize_name(ent.display_name) == name_normalized:
            return ent

    return None


# ---------------------------------------------------------------------------
# Phase 2: Fuzzy match (embedding-based when available, difflib fallback)
# ---------------------------------------------------------------------------


async def _fuzzy_match(
    name: str,
    existing_entities: list[MemoryEntity],
    embeddings: EmbeddingProvider | None = None,
) -> list[tuple[MemoryEntity, float]]:
    """Phase 2: fuzzy match using embeddings (when available) or difflib fallback.

    When embeddings provider is available and entity has embedding_vec,
    uses cosine similarity for semantic matching (handles misspellings,
    abbreviations, transliterations). Falls back to difflib SequenceMatcher
    when embeddings are unavailable.

    Returns list of (entity, score) sorted by score descending.
    """
    name_norm = _normalize_name(name)
    candidates: list[tuple[MemoryEntity, float]] = []

    # Try embedding-based matching first
    name_embedding: list[float] | None = None
    if embeddings is not None:
        try:
            name_embedding = await embeddings.embed_one(name)
        except Exception as e:
            logger.debug("entity_resolution: embedding for fuzzy match failed: %s", e)

    for ent in existing_entities:
        best_score = 0.0

        # Embedding-based similarity (when both have embeddings)
        if name_embedding and getattr(ent, "embedding_vec", None) is not None:
            sim = cosine_similarity(name_embedding, list(ent.embedding_vec))
            best_score = max(best_score, sim)
        else:
            # Difflib fallback
            if ent.display_name:
                ent_name = _normalize_name(ent.display_name)
                score = SequenceMatcher(None, name_norm, ent_name).ratio()
                best_score = max(best_score, score)

            slug = ent.canonical_key.split(":", 1)[1] if ":" in ent.canonical_key else ent.canonical_key
            slug_norm = _normalize_name(slug).replace("_", " ")
            score = SequenceMatcher(None, name_norm, slug_norm).ratio()
            best_score = max(best_score, score)

        if best_score >= FUZZY_LLM_THRESHOLD:
            candidates.append((ent, best_score))

    candidates.sort(key=lambda x: x[1], reverse=True)
    return candidates[:3]


# ---------------------------------------------------------------------------
# Phase 3: LLM fallback
# ---------------------------------------------------------------------------
async def _llm_resolve(
    name: str,
    entity_type: str,
    candidates: list[tuple[MemoryEntity, float]],
    llm: LLMProvider,
) -> MemoryEntity | None:
    """Phase 3: ask LLM if entity matches any candidate."""
    candidate_list = "\n".join(
        f'{i + 1}. "{c.display_name or c.canonical_key}" (type: {c.entity_type}, key: {c.canonical_key})'
        for i, (c, score) in enumerate(candidates)
    )

    prompt = (
        f'Is the entity "{name}" (type: {entity_type}) the same as any of these existing entities?\n\n'
        f"{candidate_list}\n\n"
        f'Respond with JSON: {{"match_index": N}} where N is the 1-based index, '
        f'or {{"match_index": 0}} if none match.'
    )

    try:
        _llm_result = await llm.complete(
            messages=[
                {"role": "system", "content": "You resolve entity identity. Respond only with JSON."},
                {"role": "user", "content": prompt},
            ],
            response_format={"type": "json_object"},
            temperature=0,
        )
        cleaned = strip_markdown_fences(_llm_result.text)
        data = json.loads(cleaned)
        idx = data.get("match_index", 0)
        if isinstance(idx, int) and 1 <= idx <= len(candidates):
            return candidates[idx - 1][0]
    except Exception as e:
        logger.warning("entity_resolution: LLM fallback failed: %s", e)

    return None


# ---------------------------------------------------------------------------
# Entity creation
# ---------------------------------------------------------------------------
async def _create_entity(
    session: AsyncSession,
    user_id: str,
    entity: ExtractedEntity,
    embeddings: EmbeddingProvider | None = None,
    aliases: list[str] | None = None,
) -> ResolvedEntity:
    """Create a new MemoryEntity and register alias. Uses savepoint pattern.

    Args:
        session: Database session.
        user_id: User identifier.
        entity: Extracted entity to create.
        embeddings: Optional embedding provider for entity embedding.
        aliases: Additional alias names to register for this entity.
    """
    slug = _slugify(entity.name)
    entity_type_norm = entity.type.lower().strip() if entity.type else "other"
    key = f"{entity_type_norm}:{slug}"

    # Generate embedding for entity name (matches original behavior)
    embedding_vec = None
    if embeddings is not None:
        try:
            embedding_vec = await embeddings.embed_one(entity.name)
        except Exception as e:
            logger.debug("entity_resolution: embedding failed for '%s': %s", entity.name, e)

    new_entity = MemoryEntity()
    new_entity.id = uuid_mod.uuid4()
    new_entity.user_id = user_id
    new_entity.canonical_key = key
    new_entity.display_name = entity.name
    new_entity.entity_type = entity_type_norm
    new_entity.embedding_vec = embedding_vec
    new_entity.fact_count = 0
    new_entity.is_active = True

    alias_record = MemoryEntityAlias()
    alias_record.id = uuid_mod.uuid4()
    alias_record.user_id = user_id
    alias_record.alias = _normalize_name(entity.name)
    alias_record.canonical_entity_key = key
    alias_record.canonical_entity_type = entity_type_norm

    try:
        async with session.begin_nested():
            session.add(new_entity)
            session.add(alias_record)
            await session.flush()
    except Exception as e:
        logger.warning(
            "entity_resolution: _create_entity failed for '%s' (key=%s): %s",
            entity.name,
            key,
            e,
        )
        return ResolvedEntity(
            entity_type=entity_type_norm,
            entity_key=key,
            display_name=entity.name,
            is_new=False,
            resolution_method="conflict_fallback",
        )

    # Register additional aliases from extraction (e.g., nicknames)
    for alias_name in aliases or []:
        await _register_alias(session, user_id, alias_name, key, entity_type_norm)

    return ResolvedEntity(
        entity_type=entity_type_norm,
        entity_key=key,
        display_name=entity.name,
        is_new=True,
        resolution_method="new",
    )


async def _register_alias(
    session: AsyncSession,
    user_id: str,
    alias_name: str,
    entity_key: str,
    entity_type: str,
) -> None:
    """Register a new alias mapping (first-write-wins). Uses savepoint."""
    normalized = _normalize_name(alias_name)
    try:
        async with session.begin_nested():
            alias_record = MemoryEntityAlias()
            alias_record.id = uuid_mod.uuid4()
            alias_record.user_id = user_id
            alias_record.alias = normalized
            alias_record.canonical_entity_key = entity_key
            alias_record.canonical_entity_type = entity_type
            session.add(alias_record)
            await session.flush()
    except Exception:
        logger.debug("entity_resolution: alias '%s' already exists, skipping", normalized)


# ---------------------------------------------------------------------------
# Single entity resolution
# ---------------------------------------------------------------------------
async def _resolve_single(
    session: AsyncSession,
    user_id: str,
    entity: ExtractedEntity,
    existing_entities: list[MemoryEntity],
    alias_cache: dict[str, str],
    llm: LLMProvider,
    embeddings: EmbeddingProvider | None = None,
    *,
    fuzzy_threshold: float = FUZZY_DIRECT_THRESHOLD,
    enable_llm_resolution: bool = True,
) -> ResolvedEntity:
    """Resolve a single extracted entity to a canonical entity key.

    3 phases: exact match → fuzzy (string) → LLM fallback.
    When enable_llm_resolution is False, Phase 3 (LLM) is skipped and
    ambiguous candidates (score 0.50–0.85) create a new entity instead.
    """
    clean_name, hint = _parse_subject_hint(entity.name)
    entity_type = entity.type.lower().strip() if entity.type else "other"

    if _hint_implies_person(hint):
        entity_type = "person"

    clean_entity = ExtractedEntity(name=clean_name, type=entity_type, aliases=entity.aliases)

    # Phase 0: Self-reference
    name_lower = clean_name.lower().strip()
    if name_lower in _SELF_REFERENCES:
        return ResolvedEntity(
            entity_type="person",
            entity_key="user:self",
            display_name="user",
            is_new=False,
            resolution_method="self",
        )

    name_normalized = _normalize_name(clean_name)

    # Phase 0b: Bare relationship term
    if name_normalized in _RELATIONSHIP_HINTS:
        return ResolvedEntity(
            entity_type="person",
            entity_key="user:self",
            display_name="user",
            is_new=False,
            resolution_method="relationship_term",
        )

    # Phase 1: Exact match
    matched = _exact_match(name_normalized, existing_entities, alias_cache)
    if matched is not None:
        return ResolvedEntity(
            entity_type=matched.entity_type,
            entity_key=matched.canonical_key,
            display_name=matched.display_name or clean_name,
            is_new=False,
            resolution_method="exact",
        )

    # Phase 1.5: Prefix/diminutive match (person entities only)
    if entity_type == "person" or not entity_type:
        name_norm_stripped = _normalize_name(clean_name)
        for ent in existing_entities:
            if ent.entity_type != "person":
                continue
            ent_name_norm = _normalize_name(ent.display_name or "")
            if not ent_name_norm or not name_norm_stripped:
                continue
            if (
                len(name_norm_stripped) >= 3
                and len(ent_name_norm) >= 3
                and (ent_name_norm.startswith(name_norm_stripped) or name_norm_stripped.startswith(ent_name_norm))
            ):
                await _register_alias(
                    session,
                    user_id,
                    clean_name,
                    ent.canonical_key,
                    ent.entity_type,
                )
                return ResolvedEntity(
                    entity_type=ent.entity_type,
                    entity_key=ent.canonical_key,
                    display_name=ent.display_name or clean_name,
                    is_new=False,
                    resolution_method="nickname",
                )

    # Phase 2: Fuzzy match (embedding-based when available, difflib fallback)
    candidates = await _fuzzy_match(clean_name, existing_entities, embeddings=embeddings)

    if candidates:
        best_entity, best_score = candidates[0]

        if best_score >= fuzzy_threshold:
            await _register_alias(
                session,
                user_id,
                clean_name,
                best_entity.canonical_key,
                best_entity.entity_type,
            )
            return ResolvedEntity(
                entity_type=best_entity.entity_type,
                entity_key=best_entity.canonical_key,
                display_name=best_entity.display_name or clean_name,
                is_new=False,
                resolution_method="fuzzy",
                fuzzy_score=best_score,
            )

        # Ambiguous: score 0.5-0.85 → Phase 3 LLM (when enabled)
        ambiguous = [(ent, score) for ent, score in candidates if score >= FUZZY_LLM_THRESHOLD]
        if ambiguous and enable_llm_resolution:
            llm_match = await _llm_resolve(clean_name, entity_type, ambiguous, llm)
            if llm_match is not None:
                await _register_alias(
                    session,
                    user_id,
                    clean_name,
                    llm_match.canonical_key,
                    llm_match.entity_type,
                )
                return ResolvedEntity(
                    entity_type=llm_match.entity_type,
                    entity_key=llm_match.canonical_key,
                    display_name=llm_match.display_name or clean_name,
                    is_new=False,
                    resolution_method="llm",
                    fuzzy_score=ambiguous[0][1],
                )

    # No match found → create new entity
    return await _create_entity(session, user_id, clean_entity, embeddings=embeddings, aliases=clean_entity.aliases)


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
async def resolve_entities(
    session: AsyncSession,
    user_id: str,
    extracted_entities: list[ExtractedEntity],
    extracted_facts: list[ExtractedFact],
    llm: LLMProvider,
    embeddings: EmbeddingProvider | None = None,
    *,
    fuzzy_threshold: float | None = None,
    enable_llm_resolution: bool = True,
) -> dict[str, ResolvedEntity]:
    """Resolve extracted entity names to canonical entity keys.

    Also resolves fact subjects that don't have a corresponding entity.

    Args:
        session: Database session.
        user_id: User identifier.
        extracted_entities: Entities from extraction.
        extracted_facts: Facts from extraction (subjects may need resolution).
        llm: Injected LLM provider for Phase 3 fallback.
        embeddings: Optional embedding provider for entity embedding on creation.
        enable_llm_resolution: When False, skip Phase 3 LLM fallback.
            Ambiguous candidates create a new entity instead.

    Returns:
        Mapping from entity name (lowercase) → ResolvedEntity.
    """
    _ft = fuzzy_threshold if fuzzy_threshold is not None else FUZZY_DIRECT_THRESHOLD

    if not extracted_entities and not extracted_facts:
        return {}

    existing_entities, alias_cache = await _load_entity_cache(session, user_id)

    entity_map: dict[str, ResolvedEntity] = {}

    # Resolve explicit entities
    for entity in extracted_entities:
        name_key = entity.name.lower().strip()
        if name_key not in entity_map:
            resolved = await _resolve_single(
                session,
                user_id,
                entity,
                existing_entities,
                alias_cache,
                llm,
                embeddings=embeddings,
                fuzzy_threshold=_ft,
                enable_llm_resolution=enable_llm_resolution,
            )
            entity_map[name_key] = resolved
            if resolved.is_new:
                clean_name, _ = _parse_subject_hint(entity.name)
                alias_cache[_normalize_name(clean_name)] = resolved.entity_key
                # Register extraction-provided aliases in the in-memory cache
                # so subsequent entities in this batch can resolve via Phase 1
                for alias_name in entity.aliases:
                    alias_cache[_normalize_name(alias_name)] = resolved.entity_key

    # Resolve fact subjects not covered by entities
    for fact in extracted_facts:
        subj_key = fact.subject.lower().strip()
        if subj_key not in entity_map:
            subj_normalized = _normalize_name(fact.subject)
            if subj_normalized in _RELATIONSHIP_HINTS:
                entity_map[subj_key] = ResolvedEntity(
                    entity_type="person",
                    entity_key="user:self",
                    display_name="user",
                    is_new=False,
                    resolution_method="relationship_term",
                )
                continue
            synthetic = ExtractedEntity(name=fact.subject, type="person")
            resolved = await _resolve_single(
                session,
                user_id,
                synthetic,
                existing_entities,
                alias_cache,
                llm,
                embeddings=embeddings,
                fuzzy_threshold=_ft,
                enable_llm_resolution=enable_llm_resolution,
            )
            entity_map[subj_key] = resolved
            if resolved.is_new:
                clean_name, _ = _parse_subject_hint(fact.subject)
                alias_cache[_normalize_name(clean_name)] = resolved.entity_key

    # Ensure user:self exists in MemoryEntity table so that sleep-time
    # compute (importance scoring, summary refresh) can process it.
    if any(e.entity_key == "user:self" for e in entity_map.values()):
        from arandu.write.entity_helpers import create_or_update_entity

        try:
            await create_or_update_entity(
                session,
                user_id,
                "user:self",
                display_name="user",
                entity_type="person",
            )
        except Exception as e:
            logger.debug("entity_resolution: user:self upsert failed (non-critical): %s", e)

    logger.info(
        "entity_resolution: user_id=%s resolved=%d entities (new=%d)",
        user_id,
        len(entity_map),
        sum(1 for e in entity_map.values() if e.is_new),
    )

    return entity_map
