"""Graph retrieval — BFS 2-hop traversal on the MemoryEntityRelationship KG.

Given a user_id + list of entity_keys, finds neighboring entities (2-hop default)
with relevance pruning on hop 2+, and collects MemoryFact records attached to
those neighbors. Scoring is composite: edge_strength * recency * edge_recency *
query_bonus.

Returns GraphRetrievalResult with scored facts for use in multi-signal merge.
"""

from __future__ import annotations

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

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

from arandu.models import MemoryEntityRelationship, MemoryFact

logger = logging.getLogger(__name__)

# Hard limit per spec: max 30 facts from graph signal
GRAPH_MAX_FACTS = 30
GRAPH_TIMEOUT_SEC = 2.0


@dataclass
class GraphRetrievalResult:
    """Result of graph-based retrieval.

    Attributes:
        facts: Scored fact dicts with source="graph".
        neighbor_keys: Entity keys discovered via BFS.
        edges_traversed: Total edges examined during BFS.
        edges: Deduplicated edge dicts with display names.
    """

    facts: list[dict[str, Any]] = field(default_factory=list)
    neighbor_keys: list[str] = field(default_factory=list)
    edges_traversed: int = 0
    edges: list[dict[str, Any]] = field(default_factory=list)


async def retrieve_graph_facts(
    session: AsyncSession,
    user_id: str,
    entity_keys: list[str],
    *,
    min_confidence: float = 0.3,
    as_of_start: datetime | None = None,
    as_of_end: datetime | None = None,
    broad_query: bool = False,
    max_facts: int | None = None,
    query_text: str = "",
    min_edge_strength: float = 0.5,
) -> GraphRetrievalResult:
    """BFS 2-hop retrieval on the knowledge graph with relevance pruning.

    1. Find neighbor entity_keys via MemoryEntityRelationship (2-hop,
       pruned by min_edge_strength on hop 2+).
    2. Fetch MemoryFact records for neighbor entities.
    3. Score facts with composite graph_score.
    4. Return scored facts with source="graph" and traversed edges.

    Args:
        session: Database session.
        user_id: User identifier.
        entity_keys: Seed entity_keys to start BFS from.
        min_confidence: Minimum fact confidence threshold.
        as_of_start: Start of temporal window (inclusive).
        as_of_end: End of temporal window (inclusive).
        broad_query: When True, allows expanded budget.
        max_facts: Override GRAPH_MAX_FACTS limit.
        query_text: Original query text for query_bonus scoring.
        min_edge_strength: Minimum edge strength for hop 2+ pruning.

    Returns:
        GraphRetrievalResult with scored facts and graph metadata.
    """
    if not entity_keys:
        return GraphRetrievalResult()

    try:
        # Step 1: BFS 2-hop neighbors with pruning
        neighbor_keys, edges_count, raw_edges, edge_strength_map, edge_valid_from_map = await _bfs_neighbors(
            session,
            user_id,
            entity_keys,
            hops=2,
            min_edge_strength=min_edge_strength,
        )

        # Step 2: Build edge dicts with display names
        edges = await _build_edge_dicts(session, user_id, raw_edges)

        if not neighbor_keys:
            return GraphRetrievalResult(edges_traversed=edges_count, edges=edges)

        # Step 3: Fetch facts for neighbor entities
        _limit = max_facts if max_facts is not None else GRAPH_MAX_FACTS
        facts = await _fetch_facts_for_entities(
            session,
            user_id,
            list(neighbor_keys),
            min_confidence=min_confidence,
            as_of_start=as_of_start,
            as_of_end=as_of_end,
            limit=_limit,
        )

        # Step 4: Composite scoring
        now = datetime.now(UTC)
        scored_facts: list[dict[str, Any]] = []
        for fact in facts:
            age_days = max(0.0, (now - (fact.valid_from or now)).total_seconds() / 86400)
            recency_factor = max(0.1, 1.0 / (1.0 + age_days / 30.0))

            edge_strength = edge_strength_map.get(fact.entity_key, 0.5)

            # Edge recency: decay based on edge valid_from (half-life 90 days)
            edge_vf = edge_valid_from_map.get(fact.entity_key)
            if edge_vf:
                edge_age_days = max(0.0, (now - edge_vf).total_seconds() / 86400)
            else:
                edge_age_days = 0.0
            edge_recency_factor = max(0.2, 1.0 / (1.0 + edge_age_days / 90.0))

            # Query bonus: 1.5x when entity_name appears in query_text
            query_bonus = 1.0
            if query_text and _entity_name_in_query(fact, query_text):
                query_bonus = 1.5

            graph_score = edge_strength * recency_factor * edge_recency_factor * query_bonus

            scored_facts.append(
                {
                    "fact": fact,
                    "score": graph_score,
                    "source": "graph",
                    "entity": f"{fact.entity_type}:{fact.entity_key}",
                    "attribute": fact.attribute_key,
                    "value": fact.value_text or str(fact.value_json),
                    "date": fact.valid_from.strftime("%Y-%m-%d") if fact.valid_from else "",
                }
            )

        logger.debug(
            "graph_retrieval: user=%s seeds=%s neighbors=%d facts=%d edges=%d",
            user_id,
            entity_keys,
            len(neighbor_keys),
            len(scored_facts),
            edges_count,
        )

        return GraphRetrievalResult(
            facts=scored_facts,
            neighbor_keys=list(neighbor_keys),
            edges_traversed=edges_count,
            edges=edges,
        )

    except Exception as e:
        logger.warning("graph_retrieval_error: user=%s error=%s", user_id, e)
        return GraphRetrievalResult()


def _entity_name_in_query(fact: Any, query_text: str) -> bool:
    """Check if the fact's entity_name appears in the query text."""
    entity_name = getattr(fact, "entity_name", None)
    if not entity_name or not isinstance(entity_name, str):
        return False
    return entity_name.lower() in query_text.lower()


async def _bfs_neighbors(
    session: AsyncSession,
    user_id: str,
    seed_keys: list[str],
    *,
    hops: int = 2,
    min_edge_strength: float = 0.5,
) -> tuple[set[str], int, list[tuple[str, str, str, float]], dict[str, float], dict[str, datetime]]:
    """BFS traversal on MemoryEntityRelationship with relevance pruning.

    Hop 1: no pruning (all direct neighbors).
    Hop 2+: only expand edges with strength >= min_edge_strength.

    Returns:
        Tuple of (neighbor_keys, edges_traversed, raw_edges,
        edge_strength_map, edge_valid_from_map).
    """
    seed_set = set(seed_keys)
    visited: set[str] = set(seed_keys)
    frontier: set[str] = set(seed_keys)
    all_neighbors: set[str] = set()
    total_edges = 0
    raw_edges: list[tuple[str, str, str, float]] = []
    edge_strength_map: dict[str, float] = {}
    edge_valid_from_map: dict[str, datetime] = {}

    for _hop in range(hops):
        if not frontier:
            break

        # Bidirectional: find neighbors where frontier is source OR target
        stmt = select(
            MemoryEntityRelationship.source_entity_key,
            MemoryEntityRelationship.target_entity_key,
            MemoryEntityRelationship.rel_type,
            MemoryEntityRelationship.strength,
            MemoryEntityRelationship.valid_from,
        ).where(
            MemoryEntityRelationship.user_id == user_id,
            MemoryEntityRelationship.invalidated_at.is_(None),
            or_(
                MemoryEntityRelationship.source_entity_key.in_(frontier),
                MemoryEntityRelationship.target_entity_key.in_(frontier),
            ),
        )

        result = await session.execute(stmt)
        rows = result.all()
        total_edges += len(rows)

        next_frontier: set[str] = set()
        for source_key, target_key, rel_type, strength, valid_from in rows:
            strength_val = float(strength or 0.8)

            # Hop 2+: prune edges below min_edge_strength threshold
            if _hop >= 1 and strength_val < min_edge_strength:
                continue

            raw_edges.append((source_key, target_key, rel_type, strength_val))

            # Add the OTHER side of the edge
            if source_key in frontier and target_key not in visited:
                next_frontier.add(target_key)
                all_neighbors.add(target_key)
                if target_key not in edge_strength_map or strength_val > edge_strength_map[target_key]:
                    edge_strength_map[target_key] = strength_val
                if valid_from and (
                    target_key not in edge_valid_from_map or valid_from > edge_valid_from_map[target_key]
                ):
                    edge_valid_from_map[target_key] = valid_from

            if target_key in frontier and source_key not in visited:
                next_frontier.add(source_key)
                all_neighbors.add(source_key)
                if source_key not in edge_strength_map or strength_val > edge_strength_map[source_key]:
                    edge_strength_map[source_key] = strength_val
                if valid_from and (
                    source_key not in edge_valid_from_map or valid_from > edge_valid_from_map[source_key]
                ):
                    edge_valid_from_map[source_key] = valid_from

        visited.update(next_frontier)
        frontier = next_frontier

    # Exclude seed keys from neighbors (only NEW entities)
    all_neighbors -= seed_set
    return all_neighbors, total_edges, raw_edges, edge_strength_map, edge_valid_from_map


def _display_name_from_key(entity_key: str) -> str:
    """Derive a readable display name from an entity_key slug.

    Args:
        entity_key: The entity key (e.g. "person:kevin_nascimento").

    Returns:
        Human-readable name (e.g. "Kevin Nascimento").
    """
    if entity_key == "user:self":
        return "User"
    name_part = entity_key.split(":", 1)[-1] if ":" in entity_key else entity_key
    return name_part.replace("_", " ").title()


async def _build_edge_dicts(
    session: AsyncSession,
    user_id: str,
    raw_edges: list[tuple[str, str, str, float]],
) -> list[dict[str, Any]]:
    """Build edge dicts with display names from raw BFS edge tuples.

    Fetches entity_name from MemoryFact for each unique entity_key.
    Falls back to _display_name_from_key() when DB name is unavailable.

    Args:
        session: Database session.
        user_id: User identifier.
        raw_edges: List of (source_key, target_key, rel_type, strength) tuples.

    Returns:
        Deduplicated edge dicts sorted by strength descending.
    """
    if not raw_edges:
        return []

    # Collect unique entity keys
    all_keys: set[str] = set()
    for source_key, target_key, _rt, _st in raw_edges:
        all_keys.add(source_key)
        all_keys.add(target_key)

    # Fetch display names: pick the most recent entity_name per entity_key
    name_map: dict[str, str] = {}
    if all_keys:
        stmt = (
            select(
                MemoryFact.entity_key,
                MemoryFact.entity_name,
            )
            .where(
                MemoryFact.user_id == user_id,
                MemoryFact.entity_key.in_(all_keys),
                MemoryFact.entity_name.isnot(None),
            )
            .distinct(MemoryFact.entity_key)
            .order_by(MemoryFact.entity_key, MemoryFact.valid_from.desc())
        )
        try:
            result = await session.execute(stmt)
            for row in result.all():
                if row[1]:
                    name_map[row[0]] = row[1]
        except Exception as e:
            logger.warning("edge_display_names: error fetching names: %s", e)

    # Build edge dicts, deduplicated by (source, target, rel_type)
    seen: set[tuple[str, str, str]] = set()
    edges: list[dict[str, Any]] = []
    for source_key, target_key, rel_type, strength in raw_edges:
        dedup = (source_key, target_key, rel_type)
        if dedup in seen:
            continue
        seen.add(dedup)
        edges.append(
            {
                "source_key": source_key,
                "target_key": target_key,
                "rel_type": rel_type,
                "strength": round(float(strength or 0.8), 2),
                "source_name": name_map.get(source_key, _display_name_from_key(source_key)),
                "target_name": name_map.get(target_key, _display_name_from_key(target_key)),
            }
        )

    # Sort by strength desc for context priority
    edges.sort(key=lambda e: e["strength"], reverse=True)
    return edges


async def _fetch_facts_for_entities(
    session: AsyncSession,
    user_id: str,
    entity_keys: list[str],
    *,
    min_confidence: float = 0.3,
    as_of_start: datetime | None = None,
    as_of_end: datetime | None = None,
    limit: int = GRAPH_MAX_FACTS,
) -> list[Any]:
    """Fetch active MemoryFact records for a list of entity_keys.

    Args:
        session: Database session.
        user_id: User identifier.
        entity_keys: Entity keys to fetch facts for.
        min_confidence: Minimum confidence threshold.
        as_of_start: Start of temporal window.
        as_of_end: End of temporal window.
        limit: Maximum number of facts to return.

    Returns:
        List of MemoryFact ORM objects.
    """
    if not entity_keys:
        return []

    # Temporal filter
    temporal_filters: list[Any] = []
    if as_of_start is not None and as_of_end is not None:
        temporal_filters.append(MemoryFact.valid_from <= as_of_end)
        temporal_filters.append(or_(MemoryFact.valid_to.is_(None), MemoryFact.valid_to > as_of_start))
    else:
        temporal_filters.append(MemoryFact.valid_to.is_(None))

    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.entity_key.in_(entity_keys),
            MemoryFact.confidence >= min_confidence,
            *temporal_filters,
        )
        .order_by(MemoryFact.valid_from.desc())
        .limit(limit)
        .options(defer(MemoryFact.embedding))
    )

    result = await session.execute(stmt)
    return list(result.scalars().all())
