"""Query Expansion — Entity Priming.

Post-processes a ``RetrievalPlan`` produced by the retrieval agent to:

1. **Entity priming**: resolve entities mentioned in the query via KG
   (aliases + relationships) and inject context into the query.

Temporal anchoring is handled by the retrieval agent LLM (``as_of_range``).
Called from ``run_read_pipeline()`` between ``plan_retrieval()`` and the
actual retrieval execution.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from datetime import datetime

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

from arandu.models import MemoryEntityAlias, MemoryEntityRelationship, MemoryFact
from arandu.read.retrieval_agent import RetrievalPlan

logger = logging.getLogger(__name__)

__all__ = ["ExpandedQuery", "expand_query"]

# Maximum entities to expand and max chars to add to query
_MAX_ENTITIES_EXPAND = 5
_MAX_CHARS_ADDED = 200


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------


@dataclass
class ExpandedQuery:
    """Result of query expansion.

    Attributes:
        primed_entities: Entity keys discovered via alias + KG priming.
        temporal_range: Resolved date range (from retrieval agent as_of_range).
        expanded_terms: Additional context terms from entity facts.
    """

    primed_entities: list[str]
    temporal_range: tuple[datetime, datetime] | None
    expanded_terms: list[str]


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


async def expand_query(
    session: AsyncSession,
    user_id: str,
    query: str,
    plan: RetrievalPlan,
    llm: object,
) -> ExpandedQuery:
    """Expand a retrieval plan with entity priming and temporal anchoring.

    Fail-safe: any exception returns an empty ExpandedQuery.

    Args:
        session: Database session.
        user_id: User identifier.
        query: Original user query text.
        plan: RetrievalPlan from the retrieval agent.
        llm: LLM provider (reserved for future use).

    Returns:
        ExpandedQuery with primed entities, temporal range, and expanded terms.
    """
    try:
        return await _expand_query_impl(session, user_id, query, plan)
    except Exception as e:
        logger.warning("query_expansion_failed: user=%s, error=%s", user_id, e)
        return ExpandedQuery(primed_entities=[], temporal_range=None, expanded_terms=[])


async def _expand_query_impl(
    session: AsyncSession,
    user_id: str,
    query: str,
    plan: RetrievalPlan,
) -> ExpandedQuery:
    """Internal implementation — may raise."""
    query_text = plan.similarity_query or query or ""
    primed_entities: list[str] = []
    expanded_terms: list[str] = []
    temporal_range: tuple[datetime, datetime] | None = None

    # 1. Entity priming (only for similarity / multi_signal strategies)
    if plan.strategy in ("similarity", "multi_signal", "hybrid") and query_text:
        primed_entities, expanded_terms = await _entity_priming(session, user_id, query_text)

    # Temporal anchoring is handled by the retrieval agent LLM (as_of_range).

    return ExpandedQuery(
        primed_entities=primed_entities,
        temporal_range=temporal_range,
        expanded_terms=expanded_terms,
    )


# ---------------------------------------------------------------------------
# Entity priming
# ---------------------------------------------------------------------------


async def _entity_priming(
    session: AsyncSession,
    user_id: str,
    query_text: str,
) -> tuple[list[str], list[str]]:
    """Resolve entities mentioned in the query and return context.

    Steps:
    1. Search aliases matching words in the query → canonical entity_keys.
    2. Fetch 1-hop KG relationships for matched entities.
    3. Fetch top facts (value_text) for matched entities as context.

    Returns:
        Tuple of (entity_keys, context_terms).
    """
    words = _extract_query_words(query_text)
    if not words:
        return [], []

    matched_entities = await _resolve_aliases(session, user_id, words)
    if not matched_entities:
        return [], []

    entity_keys = set(matched_entities.values())

    related_keys = await _fetch_related_entities(session, user_id, entity_keys)
    all_keys = entity_keys | related_keys

    context_parts = await _fetch_entity_context(session, user_id, all_keys)

    expansion = " ".join(context_parts)
    if len(expansion) > _MAX_CHARS_ADDED:
        expansion = expansion[:_MAX_CHARS_ADDED].rsplit(" ", 1)[0]

    terms = [t.strip() for t in expansion.split() if t.strip()]

    if terms:
        logger.info(
            "query_expansion: entity priming added %d terms for user %s (%d entities)",
            len(terms),
            user_id,
            len(all_keys),
        )

    return list(all_keys), terms


def _extract_query_words(query_text: str) -> list[str]:
    """Extract candidate words from query for alias matching."""
    words = re.findall(r"\b[A-Za-zÀ-ÿ]{2,}\b", query_text)
    return list({w.lower() for w in words})


async def _resolve_aliases(
    session: AsyncSession,
    user_id: str,
    words: list[str],
) -> dict[str, str]:
    """Resolve words against MemoryEntityAlias.

    Returns:
        Dict mapping alias → canonical_entity_key for matches found.
    """
    if not words:
        return {}

    stmt = (
        select(MemoryEntityAlias.alias, MemoryEntityAlias.canonical_entity_key)
        .where(
            MemoryEntityAlias.user_id == user_id,
            MemoryEntityAlias.alias.in_(words),
        )
        .limit(_MAX_ENTITIES_EXPAND * 2)
    )
    result = await session.execute(stmt)
    return {row[0]: row[1] for row in result.all()}


async def _fetch_related_entities(
    session: AsyncSession,
    user_id: str,
    entity_keys: set[str],
) -> set[str]:
    """Fetch 1-hop related entity_keys from KG relationships."""
    if not entity_keys:
        return set()

    keys_list = list(entity_keys)
    stmt = (
        select(
            MemoryEntityRelationship.source_entity_key,
            MemoryEntityRelationship.target_entity_key,
        )
        .where(
            MemoryEntityRelationship.user_id == user_id,
            or_(
                MemoryEntityRelationship.source_entity_key.in_(keys_list),
                MemoryEntityRelationship.target_entity_key.in_(keys_list),
            ),
        )
        .limit(_MAX_ENTITIES_EXPAND * 3)
    )
    result = await session.execute(stmt)

    related: set[str] = set()
    for source, target in result.all():
        if source not in entity_keys:
            related.add(source)
        if target not in entity_keys:
            related.add(target)

    if len(related) > _MAX_ENTITIES_EXPAND:
        related = set(list(related)[:_MAX_ENTITIES_EXPAND])

    return related


async def _fetch_entity_context(
    session: AsyncSession,
    user_id: str,
    entity_keys: set[str],
) -> list[str]:
    """Fetch top facts for entities to build context terms."""
    if not entity_keys:
        return []

    keys_list = list(entity_keys)[:_MAX_ENTITIES_EXPAND]
    stmt = (
        select(MemoryFact.entity_key, MemoryFact.value_text)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.entity_key.in_(keys_list),
            MemoryFact.valid_to.is_(None),
            MemoryFact.value_text.isnot(None),
        )
        .order_by(MemoryFact.confidence.desc())
        .limit(_MAX_ENTITIES_EXPAND * 3)
    )
    result = await session.execute(stmt)

    parts: list[str] = []
    for _entity_key, value_text in result.all():
        if value_text and value_text.strip():
            parts.append(value_text.strip())

    return parts


# Temporal anchoring was removed — the retrieval agent LLM handles temporal
# resolution via as_of_range in the plan JSON. This approach is language-agnostic.
