"""Fact retrieval — semantic search + keyword matching + recency scoring.

Ported from advisor_api memory_retrieval.py (1945 lines → ~250 lines v0).
Only 3 signals: semantic (pgvector cosine), keyword (ILIKE), recency (exp decay).
All advanced features (graph, spreading activation, clusters, query expansion,
emotional trends, context compression) are anti-scope for v0.
"""

import logging
import math
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from arandu._fact_query import FactQuery
from arandu.config import MemoryConfig
from arandu.models import MemoryEntity, MemoryEntityAlias
from arandu.protocols import EmbeddingProvider

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Entity key resolution
# ---------------------------------------------------------------------------
async def _resolve_entity_keys(
    session: AsyncSession,
    agent_id: str,
    entity_keys: list[str],
) -> tuple[list[str], list[str]]:
    """Resolve a list of entity_keys through canonical + alias lookup.

    For each key in ``entity_keys``:
    1. If it matches a canonical key in ``memory_entities`` → keep as-is.
    2. Else, if it matches an alias in ``memory_entity_aliases`` → replace
       with the canonical key. Both the full key (e.g. ``"person:pedro"``)
       and the slug portion after the colon (``"pedro"``) are tried, since
       the write pipeline normalizes alias storage.
    3. Else, mark as invalid.

    The function is fail-safe: if the DB queries raise, it falls back to
    returning the original list as ``resolved`` with no ``invalid`` entries,
    so retrieval degrades gracefully to the pre-fix behavior.

    Args:
        session: Database session.
        agent_id: Agent identifier.
        entity_keys: List of user-provided entity keys (may include aliases).

    Returns:
        ``(resolved_keys, invalid_keys)``. ``resolved_keys`` is deduplicated
        and preserves first-seen order. ``invalid_keys`` preserves the
        original form of keys that matched neither canonical nor alias.
    """
    if not entity_keys:
        return [], []

    try:
        from sqlalchemy import select

        # Step 1: canonical lookup against memory_entities
        canonical_stmt = select(MemoryEntity.canonical_key).where(
            MemoryEntity.agent_id == agent_id,
            MemoryEntity.canonical_key.in_(entity_keys),
        )
        canonical_result = await session.execute(canonical_stmt)
        canonical_hits: set[str] = {row[0] for row in canonical_result.all()}

        # Keys that weren't canonical become alias candidates.
        # Also try the slug portion (after ":") as an alternate alias form.
        alias_candidates: list[str] = []
        key_to_alias_forms: dict[str, list[str]] = {}
        for key in entity_keys:
            if key in canonical_hits:
                continue
            forms: list[str] = [key]
            if ":" in key:
                slug = key.split(":", 1)[1]
                if slug and slug != key:
                    forms.append(slug)
            key_to_alias_forms[key] = forms
            alias_candidates.extend(forms)

        # Step 2: alias lookup against memory_entity_aliases
        alias_map: dict[str, str] = {}
        if alias_candidates:
            alias_stmt = select(
                MemoryEntityAlias.alias,
                MemoryEntityAlias.canonical_entity_key,
            ).where(
                MemoryEntityAlias.agent_id == agent_id,
                MemoryEntityAlias.alias.in_(alias_candidates),
            )
            alias_result = await session.execute(alias_stmt)
            alias_map = {row[0]: row[1] for row in alias_result.all()}

        # Step 3: assemble resolved + invalid, preserving order and dedup
        resolved_keys: list[str] = []
        invalid_keys: list[str] = []
        seen: set[str] = set()

        for key in entity_keys:
            if key in canonical_hits:
                if key not in seen:
                    resolved_keys.append(key)
                    seen.add(key)
                continue

            matched_canonical: str | None = None
            for form in key_to_alias_forms.get(key, []):
                if form in alias_map:
                    matched_canonical = alias_map[form]
                    break

            if matched_canonical:
                if matched_canonical not in seen:
                    resolved_keys.append(matched_canonical)
                    seen.add(matched_canonical)
            else:
                invalid_keys.append(key)

        return resolved_keys, invalid_keys

    except Exception as e:  # fail-safe: degrade to legacy behavior
        logger.warning(
            "resolve_entity_keys: failed, falling back to raw entity_keys: %s",
            e,
        )
        return list(entity_keys), []


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------
@dataclass
class RetrievalCandidate:
    """A fact candidate with scoring breakdown."""

    fact: Any  # MemoryFact ORM object
    scores: dict[str, float] = field(default_factory=dict)
    final_score: float = 0.0


# ---------------------------------------------------------------------------
# Scoring helpers
# ---------------------------------------------------------------------------
def recency_score(
    item_time: datetime,
    half_life_days: float = 14.0,
    now: datetime | None = None,
) -> float:
    """Exponential decay based on age. Ported from memory_retrieval.py.

    Args:
        item_time: When the item became valid.
        half_life_days: Half-life for the decay curve.
        now: Anchor timestamp for "now". When ``None``, uses
            ``datetime.now(UTC)`` — but callers that need deterministic
            output across multiple ``recency_score`` invocations within
            a single request should pass an explicit ``now`` so all
            scores use the same reference point. This eliminates the
            ~6th-decimal drift caused by ``datetime.now()`` advancing
            milliseconds between back-to-back calls.
    """
    if now is None:
        now = datetime.now(UTC)
    if item_time.tzinfo is None:
        item_time = item_time.replace(tzinfo=UTC)
    age_days = (now - item_time).total_seconds() / 86400
    return math.exp(-math.log(2) * age_days / half_life_days)


def compute_combined_score(
    scores: dict[str, float],
    weights: dict[str, float],
) -> float:
    """Weighted combination of signal scores."""
    total = 0.0
    for key, weight in weights.items():
        total += weight * scores.get(key, 0.0)
    return total


# ---------------------------------------------------------------------------
# Signal 1: Semantic search (pgvector cosine similarity)
# ---------------------------------------------------------------------------
async def semantic_search(
    session: AsyncSession,
    agent_id: str,
    query_embedding: list[float],
    config: MemoryConfig,
    limit: int | None = None,
    entity_keys: list[str] | None = None,
) -> list[RetrievalCandidate]:
    """Search facts by cosine similarity using pgvector.

    Returns candidates sorted by similarity descending.
    Filters: agent_id, active facts (valid_to IS NULL), confidence >= min.

    Args:
        session: Database session.
        agent_id: Agent identifier.
        query_embedding: Query embedding vector.
        config: Memory configuration.
        limit: Override for pool size.
        entity_keys: Optional entity key filter. When provided, only facts
            linked to at least one of the specified entities are returned.
            The list is expected to be already resolved (canonical keys).
    """
    pool_size = (limit or config.topk_facts) * 5

    query = (
        FactQuery.active(agent_id)
        .with_min_confidence(config.min_confidence)
        .with_embedding()
        .with_deferred_embedding()
        .order_by_similarity(query_embedding)
        .limit(pool_size)
    )
    if entity_keys is not None:
        query = query.for_entities(entity_keys)

    result = await session.execute(query.build())
    rows = result.all()

    candidates = []
    for fact, similarity in rows:
        sim = float(similarity)
        if sim < config.min_similarity:
            continue
        candidates.append(
            RetrievalCandidate(
                fact=fact,
                scores={"semantic": sim},
            )
        )

    logger.info(
        "semantic_search: user=%s, pool=%d, after_threshold=%d",
        agent_id,
        len(rows),
        len(candidates),
    )
    return candidates


# ---------------------------------------------------------------------------
# Signal 2: Keyword search (SQL ILIKE fallback)
# ---------------------------------------------------------------------------
async def keyword_search(
    session: AsyncSession,
    agent_id: str,
    query_text: str,
    config: MemoryConfig,
    limit: int = 20,
    entity_keys: list[str] | None = None,
) -> list[RetrievalCandidate]:
    """Search facts by keyword matching using ILIKE on fact_text.

    Simple but effective fallback for queries with specific terms.
    No external dependencies — pure SQL.

    Args:
        session: Database session.
        agent_id: Agent identifier.
        query_text: Query text for keyword matching.
        config: Memory configuration.
        limit: Maximum number of candidates.
        entity_keys: Optional entity key filter.
    """
    if not query_text or not query_text.strip():
        return []

    # Extract significant words (>2 chars) for ILIKE matching
    words = [w for w in query_text.strip().split() if len(w) > 2]
    if not words:
        return []

    # Cap at 5 words to avoid expensive queries
    patterns = [f"%{word}%" for word in words[:5]]

    query = (
        FactQuery.active(agent_id)
        .with_min_confidence(config.min_confidence)
        .matching_text(patterns)
        .with_deferred_embedding()
        .order_by_recency()
        .limit(limit)
    )
    if entity_keys is not None:
        query = query.for_entities(entity_keys)

    result = await session.execute(query.build())
    facts = result.scalars().all()

    candidates = []
    for fact in facts:
        # Score: fraction of query words found in fact_text
        fact_lower = (fact.fact_text or "").lower()
        matches = sum(1 for w in words if w.lower() in fact_lower)
        kw_score = matches / len(words) if words else 0.0
        candidates.append(
            RetrievalCandidate(
                fact=fact,
                scores={"keyword": kw_score},
            )
        )

    logger.info("keyword_search: user=%s, hits=%d", agent_id, len(candidates))
    return candidates


# ---------------------------------------------------------------------------
# Merge and rank candidates
# ---------------------------------------------------------------------------
def merge_and_rank(
    semantic_candidates: list[RetrievalCandidate],
    keyword_candidates: list[RetrievalCandidate],
    config: MemoryConfig,
    *,
    now: datetime | None = None,
) -> list[RetrievalCandidate]:
    """Merge candidates from semantic and keyword signals, apply recency, and rank.

    Deduplicates by fact ID. Computes combined score using config weights.

    Args:
        semantic_candidates: Candidates from semantic search.
        keyword_candidates: Candidates from keyword search.
        config: Memory configuration.
        now: Anchor timestamp for recency calculations. When ``None``,
            each ``recency_score`` call uses ``datetime.now(UTC)`` (which
            advances milliseconds between candidates and causes 6th-decimal
            score drift). Pass an explicit ``now`` to make all candidates
            in a single request share the same temporal reference.
    """
    weights = config.score_weights
    half_life = config.recency_half_life_days

    # Index by fact ID for deduplication
    by_id: dict[Any, RetrievalCandidate] = {}

    for c in semantic_candidates:
        by_id[c.fact.id] = c

    for c in keyword_candidates:
        fid = c.fact.id
        if fid in by_id:
            # Merge keyword score into existing candidate
            by_id[fid].scores["keyword"] = c.scores.get("keyword", 0.0)
        else:
            by_id[fid] = c

    # Apply recency + confidence decay to all candidates
    from arandu.read._helpers import compute_effective_confidence

    for candidate in by_id.values():
        candidate.scores.setdefault("semantic", 0.0)
        candidate.scores.setdefault("keyword", 0.0)

        valid_from = candidate.fact.valid_from
        if valid_from:
            candidate.scores["recency"] = recency_score(valid_from, half_life, now=now)
        else:
            candidate.scores["recency"] = 0.5

        # importance signal (direct from fact)
        candidate.scores["importance"] = candidate.fact.importance or 0.5

        # Apply confidence decay — older facts get lower effective confidence
        raw_conf = candidate.fact.confidence or 0.5
        eff_conf = compute_effective_confidence(raw_conf, valid_from)
        candidate.scores["confidence"] = eff_conf

        candidate.final_score = compute_combined_score(candidate.scores, weights)

    # Sort by final_score descending
    ranked = sorted(by_id.values(), key=lambda c: c.final_score, reverse=True)
    return ranked


# ---------------------------------------------------------------------------
# Signal 3: Event retrieval (semantic similarity + recency)
# ---------------------------------------------------------------------------
async def retrieve_relevant_events(
    session: AsyncSession,
    agent_id: str,
    query_embedding: list[float],
    config: MemoryConfig,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    """Retrieve relevant events by embedding similarity + recency.

    Args:
        session: Database session.
        agent_id: User identifier.
        query_embedding: Query embedding vector.
        config: Memory configuration.
        limit: Max events to return.

    Returns:
        List of event dicts with ``date``, ``text``, ``score``.
    """
    from sqlalchemy import select as sa_select

    from arandu.models import MemoryEvent

    max_events = limit or config.topk_events

    stmt = (
        sa_select(
            MemoryEvent,
            (1 - MemoryEvent.embedding_vec.cosine_distance(query_embedding)).label("similarity"),
        )
        .where(
            MemoryEvent.agent_id == agent_id,
            MemoryEvent.embedding_vec.isnot(None),
        )
        .order_by(MemoryEvent.embedding_vec.cosine_distance(query_embedding))
        .limit(config.event_max_scan)
    )

    result = await session.execute(stmt)
    rows = result.all()

    events: list[dict[str, Any]] = []
    for event, similarity in rows:
        sim = float(similarity)
        if sim < config.min_similarity:
            continue
        rec = recency_score(event.occurred_at, config.recency_half_life_days)
        combined = 0.6 * sim + 0.3 * rec + 0.1 * (event.importance or 0.5)
        events.append(
            {
                "date": str(event.occurred_at.date()) if event.occurred_at else "",
                "text": event.text or "",
                "score": combined,
                "event_id": str(event.id),
            }
        )

    events.sort(key=lambda e: e["score"], reverse=True)
    logger.info(
        "retrieve_relevant_events: user=%s, scanned=%d, returned=%d", agent_id, len(rows), min(len(events), max_events)
    )
    return events[:max_events]


# ---------------------------------------------------------------------------
# Pattern signal: boost facts that recur across multiple events
# ---------------------------------------------------------------------------
def compute_pattern_signal(
    candidates: list[RetrievalCandidate],
) -> list[RetrievalCandidate]:
    """Boost facts with high reinforcement counts (pattern signal).

    Facts that have been confirmed/reinforced multiple times get a small
    additive score boost.

    Args:
        candidates: Current ranked candidates.

    Returns:
        Candidates with updated scores (sorted by final_score).
    """
    for c in candidates:
        reinforcement = getattr(c.fact, "reinforcement_count", 0) or 0
        if reinforcement > 0:
            # Bounded pattern boost: up to 0.1 extra
            pattern_boost = min(0.1, reinforcement * 0.02)
            c.scores["pattern"] = pattern_boost
            c.final_score += pattern_boost

    candidates.sort(key=lambda x: x.final_score, reverse=True)
    return candidates


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
async def retrieve_candidates(
    session: AsyncSession,
    agent_id: str,
    query_text: str,
    embeddings: EmbeddingProvider,
    config: MemoryConfig,
    entity_keys: list[str] | None = None,
) -> tuple[list[RetrievalCandidate], list[str]]:
    """Retrieve and rank fact candidates using semantic + keyword + recency.

    Args:
        session: Database session.
        agent_id: User identifier.
        query_text: The query to search for.
        embeddings: Injected embedding provider.
        config: Memory configuration.
        entity_keys: Optional entity key filter. When provided, only facts
            linked to at least one of the specified entities are returned.
            Aliases are resolved automatically against ``memory_entity_aliases``
            before the SQL filter is applied.

    Returns:
        Tuple of ``(ranked_candidates, invalid_entity_keys)``. The second
        element is the list of user-provided keys that matched neither a
        canonical entity nor a registered alias; callers should surface
        these as warnings.
    """
    # Anchor "now" once at the start of the request so every recency score
    # in this retrieve uses the same temporal reference. Without this, each
    # ``recency_score()`` call would invoke ``datetime.now()`` independently
    # and milliseconds of drift would cause 6th-decimal score variation
    # between back-to-back runs.
    request_now = datetime.now(UTC)

    # Resolve aliases once upstream — semantic_search and keyword_search
    # receive the canonical list, so there is a single resolution point.
    resolved_entity_keys: list[str] | None = None
    invalid_entity_keys: list[str] = []
    if entity_keys is not None:
        resolved_entity_keys, invalid_entity_keys = await _resolve_entity_keys(
            session,
            agent_id,
            entity_keys,
        )

    # Get query embedding
    query_embedding = await embeddings.embed_one(query_text)

    # Run semantic search (primary signal)
    semantic_candidates: list[RetrievalCandidate] = []
    if query_embedding:
        semantic_candidates = await semantic_search(
            session,
            agent_id,
            query_embedding,
            config,
            entity_keys=resolved_entity_keys,
        )
    else:
        logger.warning("retrieve_candidates: embedding failed, skipping semantic search")

    # Run keyword search (complementary signal)
    keyword_candidates = await keyword_search(
        session,
        agent_id,
        query_text,
        config,
        entity_keys=resolved_entity_keys,
    )

    # Merge, apply recency, rank — pass the frozen ``request_now`` so all
    # candidates share the same temporal reference for deterministic scores.
    ranked = merge_and_rank(semantic_candidates, keyword_candidates, config, now=request_now)

    logger.info(
        "retrieve_candidates: user=%s, semantic=%d, keyword=%d, merged=%d",
        agent_id,
        len(semantic_candidates),
        len(keyword_candidates),
        len(ranked),
    )
    return ranked, invalid_entity_keys
