"""
Meta-Path Forward Push (MPFP) graph retrieval.

A sublinear graph traversal algorithm for memory retrieval over heterogeneous
graphs with multiple edge types (semantic, temporal, causal, entity).

Combines meta-path patterns from HIN literature with Forward Push local
propagation from Approximate PPR.

Key properties:
- Sublinear in graph size (threshold pruning bounds active nodes)
- Predefined patterns capture different retrieval intents
- All patterns run in parallel, results fused via RRF
- No LLM in the loop during traversal
"""

import asyncio
import logging
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple
from collections import defaultdict

from .types import RetrievalResult
from .graph_retrieval import GraphRetriever
from ..db_utils import acquire_with_retry

logger = logging.getLogger(__name__)


# -----------------------------------------------------------------------------
# Data Classes
# -----------------------------------------------------------------------------

@dataclass
class EdgeTarget:
    """A neighbor node with its edge weight."""
    node_id: str
    weight: float


@dataclass
class TypedAdjacency:
    """Adjacency lists split by edge type."""
    # edge_type -> from_node_id -> list of (to_node_id, weight)
    graphs: Dict[str, Dict[str, List[EdgeTarget]]] = field(default_factory=dict)

    def get_neighbors(self, edge_type: str, node_id: str) -> List[EdgeTarget]:
        """Get neighbors for a node via a specific edge type."""
        return self.graphs.get(edge_type, {}).get(node_id, [])

    def get_normalized_neighbors(
        self,
        edge_type: str,
        node_id: str,
        top_k: int
    ) -> List[EdgeTarget]:
        """Get top-k neighbors with weights normalized to sum to 1."""
        neighbors = self.get_neighbors(edge_type, node_id)[:top_k]
        if not neighbors:
            return []

        total = sum(n.weight for n in neighbors)
        if total == 0:
            return []

        return [
            EdgeTarget(node_id=n.node_id, weight=n.weight / total)
            for n in neighbors
        ]


@dataclass
class PatternResult:
    """Result from a single pattern traversal."""
    pattern: List[str]
    scores: Dict[str, float]  # node_id -> accumulated mass


@dataclass
class MPFPConfig:
    """Configuration for MPFP algorithm."""
    alpha: float = 0.15           # teleport/keep probability
    threshold: float = 1e-6       # mass pruning threshold (lower = explore more)
    top_k_neighbors: int = 20     # fan-out limit per node

    # Patterns from semantic seeds
    patterns_semantic: List[List[str]] = field(default_factory=lambda: [
        ['semantic', 'semantic'],   # topic expansion
        ['entity', 'temporal'],     # entity timeline
        ['semantic', 'causes'],     # reasoning chains (forward)
        ['semantic', 'caused_by'],  # reasoning chains (backward)
        ['entity', 'semantic'],     # entity context
    ])

    # Patterns from temporal seeds
    patterns_temporal: List[List[str]] = field(default_factory=lambda: [
        ['temporal', 'semantic'],   # what was happening then
        ['temporal', 'entity'],     # who was involved then
    ])


@dataclass
class SeedNode:
    """An entry point node with its initial score."""
    node_id: str
    score: float  # initial mass (e.g., similarity score)


# -----------------------------------------------------------------------------
# Core Algorithm
# -----------------------------------------------------------------------------

def mpfp_traverse(
    seeds: List[SeedNode],
    pattern: List[str],
    adjacency: TypedAdjacency,
    config: MPFPConfig,
) -> PatternResult:
    """
    Forward Push traversal following a meta-path pattern.

    Args:
        seeds: Entry point nodes with initial scores
        pattern: Sequence of edge types to follow
        adjacency: Typed adjacency structure
        config: Algorithm parameters

    Returns:
        PatternResult with accumulated scores per node
    """
    if not seeds:
        return PatternResult(pattern=pattern, scores={})

    scores: Dict[str, float] = {}

    # Initialize frontier with seed masses (normalized)
    total_seed_score = sum(s.score for s in seeds)
    if total_seed_score == 0:
        total_seed_score = len(seeds)  # fallback to uniform

    frontier: Dict[str, float] = {
        s.node_id: s.score / total_seed_score for s in seeds
    }

    # Follow pattern hop by hop
    for edge_type in pattern:
        next_frontier: Dict[str, float] = {}

        for node_id, mass in frontier.items():
            if mass < config.threshold:
                continue

            # Keep α portion for this node
            scores[node_id] = scores.get(node_id, 0) + config.alpha * mass

            # Push (1-α) to neighbors
            push_mass = (1 - config.alpha) * mass
            neighbors = adjacency.get_normalized_neighbors(
                edge_type, node_id, config.top_k_neighbors
            )

            for neighbor in neighbors:
                next_frontier[neighbor.node_id] = (
                    next_frontier.get(neighbor.node_id, 0) +
                    push_mass * neighbor.weight
                )

        frontier = next_frontier

    # Final frontier nodes get their remaining mass
    for node_id, mass in frontier.items():
        if mass >= config.threshold:
            scores[node_id] = scores.get(node_id, 0) + mass

    return PatternResult(pattern=pattern, scores=scores)


def rrf_fusion(
    results: List[PatternResult],
    k: int = 60,
    top_k: int = 50,
) -> List[Tuple[str, float]]:
    """
    Reciprocal Rank Fusion to combine pattern results.

    Args:
        results: List of pattern results
        k: RRF constant (higher = more uniform weighting)
        top_k: Number of results to return

    Returns:
        List of (node_id, fused_score) tuples, sorted by score descending
    """
    fused: Dict[str, float] = {}

    for result in results:
        if not result.scores:
            continue

        # Rank nodes by their score in this pattern
        ranked = sorted(
            result.scores.keys(),
            key=lambda n: result.scores[n],
            reverse=True
        )

        for rank, node_id in enumerate(ranked):
            fused[node_id] = fused.get(node_id, 0) + 1.0 / (k + rank + 1)

    # Sort by fused score and return top-k
    sorted_results = sorted(
        fused.items(),
        key=lambda x: x[1],
        reverse=True
    )

    return sorted_results[:top_k]


# -----------------------------------------------------------------------------
# Database Loading
# -----------------------------------------------------------------------------

async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency:
    """
    Load all edges for a bank, split by edge type.

    Single query, then organize in-memory for fast traversal.
    """
    async with acquire_with_retry(pool) as conn:
        rows = await conn.fetch(
            """
            SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
            FROM memory_links ml
            JOIN memory_units mu ON ml.from_unit_id = mu.id
            WHERE mu.bank_id = $1
              AND ml.weight >= 0.1
            ORDER BY ml.from_unit_id, ml.weight DESC
            """,
            bank_id
        )

    graphs: Dict[str, Dict[str, List[EdgeTarget]]] = defaultdict(
        lambda: defaultdict(list)
    )

    for row in rows:
        from_id = str(row['from_unit_id'])
        to_id = str(row['to_unit_id'])
        link_type = row['link_type']
        weight = row['weight']

        graphs[link_type][from_id].append(
            EdgeTarget(node_id=to_id, weight=weight)
        )

    return TypedAdjacency(graphs=dict(graphs))


async def fetch_memory_units_by_ids(
    pool,
    node_ids: List[str],
    fact_type: str,
) -> List[RetrievalResult]:
    """Fetch full memory unit details for a list of node IDs."""
    if not node_ids:
        return []

    async with acquire_with_retry(pool) as conn:
        rows = await conn.fetch(
            """
            SELECT id, text, context, event_date, occurred_start, occurred_end,
                   mentioned_at, access_count, embedding, fact_type, document_id, chunk_id
            FROM memory_units
            WHERE id = ANY($1::uuid[])
              AND fact_type = $2
            """,
            node_ids,
            fact_type
        )

    return [RetrievalResult.from_db_row(dict(r)) for r in rows]


# -----------------------------------------------------------------------------
# Graph Retriever Implementation
# -----------------------------------------------------------------------------

class MPFPGraphRetriever(GraphRetriever):
    """
    Graph retrieval using Meta-Path Forward Push.

    Runs predefined patterns in parallel from semantic and temporal seeds,
    then fuses results via RRF.
    """

    def __init__(self, config: Optional[MPFPConfig] = None):
        """
        Initialize MPFP retriever.

        Args:
            config: Algorithm configuration (uses defaults if None)
        """
        self.config = config or MPFPConfig()
        self._adjacency_cache: Dict[str, TypedAdjacency] = {}

    @property
    def name(self) -> str:
        return "mpfp"

    async def retrieve(
        self,
        pool,
        query_embedding_str: str,
        bank_id: str,
        fact_type: str,
        budget: int,
        query_text: Optional[str] = None,
        semantic_seeds: Optional[List[RetrievalResult]] = None,
        temporal_seeds: Optional[List[RetrievalResult]] = None,
    ) -> List[RetrievalResult]:
        """
        Retrieve facts using MPFP algorithm.

        Args:
            pool: Database connection pool
            query_embedding_str: Query embedding (used for fallback seed finding)
            bank_id: Memory bank ID
            fact_type: Fact type to filter
            budget: Maximum results to return
            query_text: Original query text (optional)
            semantic_seeds: Pre-computed semantic entry points
            temporal_seeds: Pre-computed temporal entry points

        Returns:
            List of RetrievalResult with activation scores
        """
        # Load typed adjacency (could cache per bank_id with TTL)
        adjacency = await load_typed_adjacency(pool, bank_id)

        # Convert seeds to SeedNode format
        semantic_seed_nodes = self._convert_seeds(semantic_seeds, 'similarity')
        temporal_seed_nodes = self._convert_seeds(temporal_seeds, 'temporal_score')

        # If no semantic seeds provided, fall back to finding our own
        if not semantic_seed_nodes:
            semantic_seed_nodes = await self._find_semantic_seeds(
                pool, query_embedding_str, bank_id, fact_type
            )

        # Run all patterns in parallel
        tasks = []

        # Patterns from semantic seeds
        for pattern in self.config.patterns_semantic:
            if semantic_seed_nodes:
                tasks.append(
                    asyncio.to_thread(
                        mpfp_traverse,
                        semantic_seed_nodes,
                        pattern,
                        adjacency,
                        self.config,
                    )
                )

        # Patterns from temporal seeds
        for pattern in self.config.patterns_temporal:
            if temporal_seed_nodes:
                tasks.append(
                    asyncio.to_thread(
                        mpfp_traverse,
                        temporal_seed_nodes,
                        pattern,
                        adjacency,
                        self.config,
                    )
                )

        if not tasks:
            return []

        # Gather pattern results
        pattern_results = await asyncio.gather(*tasks)

        # Fuse results
        fused = rrf_fusion(pattern_results, top_k=budget)

        if not fused:
            return []

        # Get top result IDs (don't exclude seeds - they may be highly relevant)
        result_ids = [node_id for node_id, score in fused][:budget]

        # Fetch full details
        results = await fetch_memory_units_by_ids(pool, result_ids, fact_type)

        # Add activation scores from fusion
        score_map = {node_id: score for node_id, score in fused}
        for result in results:
            result.activation = score_map.get(result.id, 0.0)

        # Sort by activation
        results.sort(key=lambda r: r.activation or 0, reverse=True)

        return results

    def _convert_seeds(
        self,
        seeds: Optional[List[RetrievalResult]],
        score_attr: str,
    ) -> List[SeedNode]:
        """Convert RetrievalResult seeds to SeedNode format."""
        if not seeds:
            return []

        result = []
        for seed in seeds:
            score = getattr(seed, score_attr, None)
            if score is None:
                score = seed.activation or seed.similarity or 1.0
            result.append(SeedNode(node_id=seed.id, score=score))

        return result

    async def _find_semantic_seeds(
        self,
        pool,
        query_embedding_str: str,
        bank_id: str,
        fact_type: str,
        limit: int = 20,
        threshold: float = 0.3,
    ) -> List[SeedNode]:
        """Fallback: find semantic seeds via embedding search."""
        async with acquire_with_retry(pool) as conn:
            rows = await conn.fetch(
                """
                SELECT id, 1 - (embedding <=> $1::vector) AS similarity
                FROM memory_units
                WHERE bank_id = $2
                  AND embedding IS NOT NULL
                  AND fact_type = $3
                  AND (1 - (embedding <=> $1::vector)) >= $4
                ORDER BY embedding <=> $1::vector
                LIMIT $5
                """,
                query_embedding_str, bank_id, fact_type, threshold, limit
            )

        return [
            SeedNode(node_id=str(r['id']), score=r['similarity'])
            for r in rows
        ]