"""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.constants import ExtractedEntity, ExtractedFact
from arandu.models import MemoryEntity, MemoryEntityAlias
from arandu.protocols import EmbeddingProvider, LLMProvider
from arandu.write.extract import _strip_markdown_fences

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
FUZZY_DIRECT_THRESHOLD = 0.85
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)
# ---------------------------------------------------------------------------


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine similarity between two vectors."""
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = sum(x * x for x in a) ** 0.5
    norm_b = sum(x * x for x in b) ** 0.5
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(dot / (norm_a * norm_b))


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:
        raw = 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(raw)
        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,
) -> ResolvedEntity:
    """Create a new MemoryEntity and register alias. Uses savepoint pattern."""
    slug = _slugify(entity.name)
    key = f"{entity.type}:{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  # type: ignore[assignment]
    new_entity.canonical_key = key
    new_entity.display_name = entity.name
    new_entity.entity_type = entity.type
    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  # type: ignore[assignment]
    alias_record.alias = _normalize_name(entity.name)
    alias_record.canonical_entity_key = key
    alias_record.canonical_entity_type = entity.type

    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,
            entity_key=key,
            display_name=entity.name,
            is_new=False,
            resolution_method="conflict_fallback",
        )

    return ResolvedEntity(
        entity_type=entity.type,
        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  # type: ignore[assignment]
            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,
) -> ResolvedEntity:
    """Resolve a single extracted entity to a canonical entity key.

    3 phases: exact match → fuzzy (string) → LLM fallback.
    """
    clean_name, hint = _parse_subject_hint(entity.name)
    entity_type = entity.type

    if _hint_implies_person(hint):
        entity_type = "person"

    clean_entity = ExtractedEntity(name=clean_name, type=entity_type)

    # 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_DIRECT_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
        ambiguous = [(ent, score) for ent, score in candidates if score >= FUZZY_LLM_THRESHOLD]
        if ambiguous:
            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)


# ---------------------------------------------------------------------------
# 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,
) -> 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.

    Returns:
        Mapping from entity name (lowercase) → ResolvedEntity.
    """
    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,
            )
            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

    # 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,
            )
            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
