"""Spreading activation for adaptive retrieval.

Expands context from seed facts by following entity_key, cluster_id,
and knowledge graph relationship links. Uses dynamic importance scoring
to weight expanded facts with decay per hop.
"""

from __future__ import annotations

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

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

from arandu.config import MemoryConfig
from arandu.models import (
    MemoryEntityRelationship,
    MemoryFact,
    MemoryMetaObservation,
)
from arandu.read._helpers import is_fact_active, is_key_allowed
from arandu.read.importance import compute_dynamic_importance
from arandu.read.retrieval import RetrievalCandidate

logger = logging.getLogger(__name__)

__all__ = ["SpreadingActivationResult", "spread_activation"]

# Hop decay factors
HOP1_DECAY = 0.6
HOP2_DECAY = 0.35

# Per-hop limits
FACTS_PER_ENTITY_HOP1 = 5
FACTS_PER_CLUSTER_HOP1 = 3
META_OBSERVATIONS_LIMIT = 3
FACTS_PER_ENTITY_HOP2 = 3
FACTS_FROM_META_HOP2 = 3

# KG relationship defaults
KG_REL_HOP_DECAY = 0.5
KG_FACTS_PER_RELATED_ENTITY = 3
KG_MAX_RELATED_ENTITIES = 5


@dataclass
class SpreadingActivationResult:
    """Result of spreading activation expansion.

    Attributes:
        candidates: Expanded RetrievalCandidates from hop 1-2.
        meta_observations: Relevant meta-observations referencing seed facts.
        entities_explored: Entity keys explored during spreading.
        clusters_explored: Cluster IDs explored during spreading.
        hop1_count: Number of facts found in hop 1.
        hop2_count: Number of facts found in hop 2.
        kg_relationships_explored: Number of KG relationships traversed.
    """

    candidates: list[RetrievalCandidate] = field(default_factory=list)
    meta_observations: list[Any] = field(default_factory=list)
    entities_explored: list[str] = field(default_factory=list)
    clusters_explored: list[str] = field(default_factory=list)
    hop1_count: int = 0
    hop2_count: int = 0
    kg_relationships_explored: int = 0


def _entity_from_fact(fact: Any) -> str:
    """Derive entity key from a fact, avoiding double-prefix."""
    entity_key = getattr(fact, "entity_key", "")
    if ":" in entity_key:
        return entity_key
    entity_type = getattr(fact, "entity_type", "")
    return f"{entity_type}:{entity_key}"


async def spread_activation(
    session: AsyncSession,
    user_id: str,
    seed_fact_ids: list[str],
    config: MemoryConfig,
    *,
    seed_scores: dict[str, float] | None = None,
    allowed_keys: set[str] | None = None,
) -> list[RetrievalCandidate]:
    """Expand context from seed facts via entity_key, cluster_id, and KG relationships.

    Performs hop 1-2 expansion:
    - Hop 1: facts sharing entity_key or cluster_id with seed facts, plus KG neighbors.
    - Hop 2: facts from entities discovered in hop 1 (if configured).

    Each expanded fact is scored with dynamic importance * hop decay * avg seed score.

    Args:
        session: Database session.
        user_id: User identifier.
        seed_fact_ids: IDs of seed facts to expand from.
        config: Memory configuration with spreading activation params.
        seed_scores: Optional dict mapping seed fact ID → score (for decay calc).
        allowed_keys: Optional set of allowed attribute keys.

    Returns:
        List of RetrievalCandidates from spreading activation.
    """
    try:
        result = await _spread_activation_impl(
            session,
            user_id,
            seed_fact_ids,
            config,
            seed_scores=seed_scores,
            allowed_keys=allowed_keys,
        )
        return result.candidates
    except Exception as e:
        logger.warning("spreading_activation_failed: user=%s, error=%s", user_id, e)
        return []


async def _spread_activation_impl(
    session: AsyncSession,
    user_id: str,
    seed_fact_ids: list[str],
    config: MemoryConfig,
    *,
    seed_scores: dict[str, float] | None = None,
    allowed_keys: set[str] | None = None,
) -> SpreadingActivationResult:
    """Internal implementation — may raise."""
    hops = config.spreading_activation_hops
    now = datetime.now(UTC)

    # Fetch seed facts to get their entity_keys and cluster_ids
    seed_uuids = _parse_uuids(seed_fact_ids)
    if not seed_uuids:
        return SpreadingActivationResult()

    seed_facts = await _fetch_facts_by_ids(session, user_id, seed_uuids)
    if not seed_facts:
        return SpreadingActivationResult()

    seed_id_set: set[Any] = {f.id for f in seed_facts}
    avg_seed_score = _compute_avg_seed_score(seed_facts, seed_scores)

    # Collect entity_keys and cluster_ids from seeds
    entity_keys: set[str] = set()
    cluster_ids: set[Any] = set()
    for fact in seed_facts:
        entity_keys.add(_entity_from_fact(fact))
        cid = getattr(fact, "cluster_id", None)
        if cid is not None:
            cluster_ids.add(cid)

    entities_explored: list[str] = []
    clusters_explored: list[str] = []
    seen_fact_ids: set[Any] = set(seed_id_set)
    hop1_candidates: list[RetrievalCandidate] = []

    # -- Hop 1: by entity_key --
    for entity_key in entity_keys:
        entities_explored.append(entity_key)
        facts = await _fetch_neighbor_facts(
            session,
            user_id,
            entity_key,
            seen_fact_ids,
            limit=FACTS_PER_ENTITY_HOP1,
        )
        for fact in facts:
            if not is_key_allowed(fact.attribute_key, allowed_keys) or not is_fact_active(fact, now):
                continue
            seen_fact_ids.add(fact.id)
            score = _score_fact(fact, HOP1_DECAY, avg_seed_score, now)
            hop1_candidates.append(
                RetrievalCandidate(
                    fact=fact,
                    scores={"spreading_hop1": score},
                    final_score=score,
                )
            )

    # -- Hop 1: by cluster_id --
    for cid in cluster_ids:
        clusters_explored.append(str(cid))
        facts = await _fetch_cluster_facts(
            session,
            user_id,
            cid,
            seen_fact_ids,
            limit=FACTS_PER_CLUSTER_HOP1,
        )
        for fact in facts:
            if not is_key_allowed(fact.attribute_key, allowed_keys) or not is_fact_active(fact, now):
                continue
            seen_fact_ids.add(fact.id)
            score = _score_fact(fact, HOP1_DECAY, avg_seed_score, now)
            hop1_candidates.append(
                RetrievalCandidate(
                    fact=fact,
                    scores={"spreading_hop1_cluster": score},
                    final_score=score,
                )
            )

    # -- Hop 1: by KG relationships --
    kg_rels_explored = 0
    entity_list = list(entity_keys)
    max_related = config.spreading_max_related_entities

    related_entities = await _fetch_kg_related(session, user_id, entity_list, max_related)
    for rel_entity_key, rel_strength in related_entities:
        if rel_entity_key in entity_keys:
            continue
        kg_rels_explored += 1
        entities_explored.append(f"kg:{rel_entity_key}")
        facts = await _fetch_neighbor_facts(
            session,
            user_id,
            rel_entity_key,
            seen_fact_ids,
            limit=KG_FACTS_PER_RELATED_ENTITY,
        )
        for fact in facts:
            if not is_key_allowed(fact.attribute_key, allowed_keys) or not is_fact_active(fact, now):
                continue
            seen_fact_ids.add(fact.id)
            dyn = _compute_fact_importance(fact, now, is_in_active_pattern=False)
            score = KG_REL_HOP_DECAY * avg_seed_score * dyn * float(rel_strength)
            hop1_candidates.append(
                RetrievalCandidate(
                    fact=fact,
                    scores={"spreading_kg": score},
                    final_score=score,
                )
            )

    # -- Hop 1: meta_observations referencing seed facts --
    meta_observations = await _fetch_meta_observations(session, user_id, seed_id_set)

    # -- Hop 2 (if configured) --
    hop2_candidates: list[RetrievalCandidate] = []
    if hops >= 2:
        hop1_entities = {_entity_from_fact(c.fact) for c in hop1_candidates}
        new_entities = hop1_entities - entity_keys
        excluded_ids = seed_id_set | {c.fact.id for c in hop1_candidates}

        for ek in new_entities:
            entities_explored.append(ek)
            facts = await _fetch_neighbor_facts(
                session,
                user_id,
                ek,
                excluded_ids,
                limit=config.spreading_facts_per_entity,
            )
            for fact in facts:
                if fact.id in seen_fact_ids:
                    continue
                if not is_key_allowed(fact.attribute_key, allowed_keys) or not is_fact_active(fact, now):
                    continue
                seen_fact_ids.add(fact.id)
                score = _score_fact(fact, HOP2_DECAY, avg_seed_score, now)
                hop2_candidates.append(
                    RetrievalCandidate(
                        fact=fact,
                        scores={"spreading_hop2": score},
                        final_score=score,
                    )
                )

    all_candidates = hop1_candidates + hop2_candidates

    logger.info(
        "spreading_activation: user=%s, seeds=%d, hop1=%d, hop2=%d, kg_rels=%d",
        user_id,
        len(seed_facts),
        len(hop1_candidates),
        len(hop2_candidates),
        kg_rels_explored,
    )

    return SpreadingActivationResult(
        candidates=all_candidates,
        meta_observations=meta_observations,
        entities_explored=entities_explored,
        clusters_explored=clusters_explored,
        hop1_count=len(hop1_candidates),
        hop2_count=len(hop2_candidates),
        kg_relationships_explored=kg_rels_explored,
    )


# ---------------------------------------------------------------------------
# Scoring helpers
# ---------------------------------------------------------------------------


def _compute_fact_importance(fact: Any, now: datetime, *, is_in_active_pattern: bool = False) -> float:
    """Compute dynamic importance for a single fact."""
    return compute_dynamic_importance(
        float(getattr(fact, "importance", 0.5) or 0.5),
        getattr(fact, "times_retrieved", 0) or 0,
        getattr(fact, "last_retrieved_at", None),
        getattr(fact, "user_correction_count", 0) or 0,
        is_in_active_pattern,
        now,
    )


def _score_fact(fact: Any, decay: float, avg_seed_score: float, now: datetime) -> float:
    """Score an expanded fact with decay * avg_seed_score * dynamic importance."""
    dyn = _compute_fact_importance(fact, now)
    return decay * avg_seed_score * dyn


def _compute_avg_seed_score(
    seed_facts: list[Any],
    seed_scores: dict[str, float] | None,
) -> float:
    """Compute average seed score from provided scores or default 0.5."""
    if seed_scores:
        scores = [seed_scores.get(str(f.id), 0.5) for f in seed_facts]
        return sum(scores) / len(scores) if scores else 0.5
    return 0.5


def _parse_uuids(ids: list[str]) -> list[Any]:
    """Parse string IDs to UUIDs, skipping invalid ones."""
    result: list[Any] = []
    for id_str in ids:
        try:
            result.append(uuid.UUID(str(id_str)))
        except (ValueError, TypeError):
            continue
    return result


# ---------------------------------------------------------------------------
# DB query helpers
# ---------------------------------------------------------------------------


async def _fetch_facts_by_ids(
    session: AsyncSession,
    user_id: str,
    fact_ids: list[Any],
) -> list[Any]:
    """Fetch MemoryFact records by IDs."""
    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.id.in_(fact_ids),
        )
        .options(defer(MemoryFact.embedding))
    )
    result = await session.execute(stmt)
    return list(result.scalars().all())


async def _fetch_neighbor_facts(
    session: AsyncSession,
    user_id: str,
    entity_key: str,
    exclude_ids: set[Any],
    *,
    limit: int = 5,
) -> list[Any]:
    """Fetch active facts for an entity_key, excluding known IDs."""
    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.valid_to.is_(None),
            MemoryFact.entity_key == entity_key,
        )
        .order_by(MemoryFact.valid_from.desc())
        .limit(limit)
        .options(defer(MemoryFact.embedding))
    )
    if exclude_ids:
        stmt = stmt.where(MemoryFact.id.notin_(exclude_ids))
    result = await session.execute(stmt)
    return list(result.scalars().all())


async def _fetch_cluster_facts(
    session: AsyncSession,
    user_id: str,
    cluster_id: Any,
    exclude_ids: set[Any],
    *,
    limit: int = 3,
) -> list[Any]:
    """Fetch active facts for a cluster_id, excluding known IDs."""
    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.valid_to.is_(None),
            MemoryFact.cluster_id == cluster_id,
        )
        .order_by(MemoryFact.valid_from.desc())
        .limit(limit)
        .options(defer(MemoryFact.embedding))
    )
    if exclude_ids:
        stmt = stmt.where(MemoryFact.id.notin_(exclude_ids))
    result = await session.execute(stmt)
    return list(result.scalars().all())


async def _fetch_kg_related(
    session: AsyncSession,
    user_id: str,
    entity_keys: list[str],
    max_related: int,
) -> list[tuple[str, float]]:
    """Fetch KG-related entities (both directions)."""
    if not entity_keys:
        return []

    # Forward
    fwd_stmt = (
        select(MemoryEntityRelationship.target_entity_key, MemoryEntityRelationship.strength)
        .where(
            MemoryEntityRelationship.user_id == user_id,
            MemoryEntityRelationship.source_entity_key.in_(entity_keys),
        )
        .limit(max_related)
    )
    fwd_result = await session.execute(fwd_stmt)
    related: list[tuple[str, float]] = [(str(r[0]), float(r[1])) for r in fwd_result.all()]

    # Reverse
    rev_stmt = (
        select(MemoryEntityRelationship.source_entity_key, MemoryEntityRelationship.strength)
        .where(
            MemoryEntityRelationship.user_id == user_id,
            MemoryEntityRelationship.target_entity_key.in_(entity_keys),
        )
        .limit(max_related)
    )
    rev_result = await session.execute(rev_stmt)
    related.extend((str(r[0]), float(r[1])) for r in rev_result.all())

    # Deduplicate
    seen: set[str] = set(entity_keys)
    deduped: list[tuple[str, float]] = []
    for ek, strength in related:
        if ek not in seen:
            seen.add(ek)
            deduped.append((ek, float(strength or 0.8)))
    return deduped[:max_related]


async def _fetch_meta_observations(
    session: AsyncSession,
    user_id: str,
    seed_ids: set[Any],
) -> list[Any]:
    """Fetch meta-observations that reference seed facts."""
    if not seed_ids:
        return []

    stmt = (
        select(MemoryMetaObservation)
        .where(
            MemoryMetaObservation.user_id == user_id,
            MemoryMetaObservation.is_active.is_(True),
        )
        .limit(50)
    )
    result = await session.execute(stmt)
    meta_observations: list[Any] = []
    seed_str = {str(x) for x in seed_ids}
    for obs in result.scalars().all():
        supporting = getattr(obs, "supporting_fact_ids", None) or []
        if not supporting:
            continue
        obs_fact_ids = {str(x) for x in supporting}
        if obs_fact_ids & seed_str:
            meta_observations.append(obs)
            if len(meta_observations) >= META_OBSERVATIONS_LIMIT:
                break

    return meta_observations
