"""
Retrieval module for 4-way parallel search.

Implements:
1. Semantic retrieval (vector similarity)
2. BM25 retrieval (keyword/full-text search)
3. Graph retrieval (spreading activation)
4. Temporal retrieval (time-aware search with spreading)
"""

from typing import List, Dict, Any, Tuple, Optional
from datetime import datetime
import asyncio
from ..db_utils import acquire_with_retry


async def retrieve_semantic(
    conn,
    query_emb_str: str,
    agent_id: str,
    fact_type: str,
    limit: int
) -> List[Tuple[str, Dict[str, Any]]]:
    """
    Semantic retrieval via vector similarity.

    Args:
        conn: Database connection
        query_emb_str: Query embedding as string
        agent_id: Agent ID
        fact_type: Fact type to filter
        limit: Maximum results to return

    Returns:
        List of (doc_id, data) tuples
    """
    results = await conn.fetch(
        """
        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
               1 - (embedding <=> $1::vector) AS similarity
        FROM memory_units
        WHERE agent_id = $2
          AND embedding IS NOT NULL
          AND fact_type = $3
          AND (1 - (embedding <=> $1::vector)) >= 0.3
        ORDER BY embedding <=> $1::vector
        LIMIT $4
        """,
        query_emb_str, agent_id, fact_type, limit
    )
    return [(str(r["id"]), dict(r)) for r in results]


async def retrieve_bm25(
    conn,
    query_text: str,
    agent_id: str,
    fact_type: str,
    limit: int
) -> List[Tuple[str, Dict[str, Any]]]:
    """
    BM25 keyword retrieval via full-text search.

    Args:
        conn: Database connection
        query_text: Query text
        agent_id: Agent ID
        fact_type: Fact type to filter
        limit: Maximum results to return

    Returns:
        List of (doc_id, data) tuples
    """
    import re

    # Sanitize query text: remove special characters that have meaning in tsquery
    # Keep only alphanumeric characters and spaces
    sanitized_text = re.sub(r'[^\w\s]', ' ', query_text.lower())

    # Split and filter empty strings
    tokens = [token for token in sanitized_text.split() if token]

    if not tokens:
        # If no valid tokens, return empty results
        return []

    # Convert query to tsquery using OR for more flexible matching
    # This prevents empty results when some terms are missing
    query_tsquery = " | ".join(tokens)

    results = await conn.fetch(
        """
        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
               ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
        FROM memory_units
        WHERE agent_id = $2
          AND fact_type = $3
          AND search_vector @@ to_tsquery('english', $1)
        ORDER BY bm25_score DESC
        LIMIT $4
        """,
        query_tsquery, agent_id, fact_type, limit
    )
    return [(str(r["id"]), dict(r)) for r in results]


async def retrieve_graph(
    conn,
    query_emb_str: str,
    agent_id: str,
    fact_type: str,
    budget: int
) -> List[Tuple[str, Dict[str, Any]]]:
    """
    Graph retrieval via spreading activation.

    Args:
        conn: Database connection
        query_emb_str: Query embedding as string
        agent_id: Agent ID
        fact_type: Fact type to filter
        budget: Node budget for graph traversal

    Returns:
        List of (doc_id, data) tuples
    """
    # Find entry points
    entry_points = await conn.fetch(
        """
        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
               1 - (embedding <=> $1::vector) AS similarity
        FROM memory_units
        WHERE agent_id = $2
          AND embedding IS NOT NULL
          AND fact_type = $3
          AND (1 - (embedding <=> $1::vector)) >= 0.5
        ORDER BY embedding <=> $1::vector
        LIMIT 5
        """,
        query_emb_str, agent_id, fact_type
    )

    if not entry_points:
        return []

    # Simple BFS-style spreading activation
    visited = set()
    results = []
    queue = [(dict(r), r["similarity"]) for r in entry_points]
    budget_remaining = budget

    while queue and budget_remaining > 0:
        current, activation = queue.pop(0)
        unit_id = str(current["id"])

        if unit_id in visited:
            continue

        visited.add(unit_id)
        budget_remaining -= 1
        results.append((unit_id, current))

        # Get neighbors
        if budget_remaining > 0:
            neighbors = await conn.fetch(
                """
                SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
                       ml.weight, ml.link_type
                FROM memory_links ml
                JOIN memory_units mu ON ml.to_unit_id = mu.id
                WHERE ml.from_unit_id = $1
                  AND ml.weight >= 0.1
                  AND mu.fact_type = $2
                ORDER BY ml.weight DESC
                LIMIT 10
                """,
                current["id"], fact_type
            )

            for n in neighbors:
                neighbor_id = str(n["id"])
                if neighbor_id not in visited:
                    # Boost activation for causal links (they're high-value relationships)
                    link_type = n["link_type"]
                    base_weight = n["weight"]

                    # Causal links get 1.5-2.0x boost depending on type
                    if link_type in ("causes", "caused_by"):
                        # Direct causation - very strong relationship
                        causal_boost = 2.0
                    elif link_type in ("enables", "prevents"):
                        # Conditional causation - strong but not as direct
                        causal_boost = 1.5
                    else:
                        # Temporal, semantic, entity links - standard weight
                        causal_boost = 1.0

                    effective_weight = base_weight * causal_boost
                    new_activation = activation * effective_weight * 0.8
                    if new_activation > 0.1:
                        queue.append((dict(n), new_activation))

    return results


async def retrieve_temporal(
    conn,
    query_emb_str: str,
    agent_id: str,
    fact_type: str,
    start_date: datetime,
    end_date: datetime,
    budget: int,
    semantic_threshold: float = 0.4
) -> List[Tuple[str, Dict[str, Any]]]:
    """
    Temporal retrieval with spreading activation.

    Strategy:
    1. Find entry points (facts in date range with semantic relevance)
    2. Spread through temporal links to related facts
    3. Score by temporal proximity + semantic similarity + link weight

    Args:
        conn: Database connection
        query_emb_str: Query embedding as string
        agent_id: Agent ID
        fact_type: Fact type to filter
        start_date: Start of time range
        end_date: End of time range
        budget: Node budget for spreading
        semantic_threshold: Minimum semantic similarity to include

    Returns:
        List of (doc_id, data) tuples with temporal_score
    """
    from datetime import timezone

    # Ensure start_date and end_date are timezone-aware (UTC) to match database datetimes
    if start_date.tzinfo is None:
        start_date = start_date.replace(tzinfo=timezone.utc)
    if end_date.tzinfo is None:
        end_date = end_date.replace(tzinfo=timezone.utc)

    entry_points = await conn.fetch(
        """
        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
               1 - (embedding <=> $1::vector) AS similarity
        FROM memory_units
        WHERE agent_id = $2
          AND fact_type = $3
          AND embedding IS NOT NULL
          AND (
              -- Match if occurred range overlaps with query range
              (occurred_start IS NOT NULL AND occurred_end IS NOT NULL
               AND occurred_start <= $5 AND occurred_end >= $4)
              OR
              -- Match if mentioned_at falls within query range
              (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
              OR
              -- Match if any occurred date is set and overlaps (even if only start or end is set)
              (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
              OR
              (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
          )
          AND (1 - (embedding <=> $1::vector)) >= $6
        ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, (embedding <=> $1::vector) ASC
        LIMIT 10
        """,
        query_emb_str, agent_id, fact_type, start_date, end_date, semantic_threshold
    )

    if not entry_points:
        # Check if there are ANY memories with temporal metadata for this agent
        total_with_dates = await conn.fetchval(
            """SELECT COUNT(*) FROM memory_units
               WHERE agent_id = $1 AND fact_type = $2
               AND (occurred_start IS NOT NULL OR occurred_end IS NOT NULL OR mentioned_at IS NOT NULL)""",
            agent_id, fact_type
        )
        return []

    # Calculate temporal scores for entry points
    total_days = (end_date - start_date).total_seconds() / 86400
    mid_date = start_date + (end_date - start_date) / 2  # Calculate once for all comparisons
    results = []
    visited = set()

    for ep in entry_points:
        unit_id = str(ep["id"])
        visited.add(unit_id)

        # Calculate temporal proximity using the most relevant date
        # Priority: occurred_start/end (event time) > mentioned_at (mention time)
        best_date = None
        if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
            # Use midpoint of occurred range
            best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
        elif ep["occurred_start"] is not None:
            best_date = ep["occurred_start"]
        elif ep["occurred_end"] is not None:
            best_date = ep["occurred_end"]
        elif ep["mentioned_at"] is not None:
            best_date = ep["mentioned_at"]

        # Temporal proximity score (closer to range center = higher score)
        if best_date:
            days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
            temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
        else:
            temporal_proximity = 0.5  # Fallback if no dates (shouldn't happen due to WHERE clause)

        data = dict(ep)
        data["temporal_score"] = temporal_proximity
        data["temporal_proximity"] = temporal_proximity
        results.append((unit_id, data))

    # Spread through temporal links
    queue = [(dict(ep), ep["similarity"], 1.0) for ep in entry_points]  # (unit, semantic_sim, temporal_score)
    budget_remaining = budget - len(entry_points)

    while queue and budget_remaining > 0:
        current, semantic_sim, temporal_score = queue.pop(0)
        current_id = str(current["id"])

        # Get neighbors via temporal and causal links
        if budget_remaining > 0:
            neighbors = await conn.fetch(
                """
                SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
                       ml.weight, ml.link_type,
                       1 - (mu.embedding <=> $1::vector) AS similarity
                FROM memory_links ml
                JOIN memory_units mu ON ml.to_unit_id = mu.id
                WHERE ml.from_unit_id = $2
                  AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
                  AND ml.weight >= 0.1
                  AND mu.fact_type = $3
                  AND mu.embedding IS NOT NULL
                  AND (1 - (mu.embedding <=> $1::vector)) >= $4
                ORDER BY ml.weight DESC
                LIMIT 10
                """,
                query_emb_str, current["id"], fact_type, semantic_threshold
            )

            for n in neighbors:
                neighbor_id = str(n["id"])
                if neighbor_id in visited:
                    continue

                visited.add(neighbor_id)
                budget_remaining -= 1

                # Calculate temporal score for neighbor using best available date
                neighbor_best_date = None
                if n["occurred_start"] is not None and n["occurred_end"] is not None:
                    neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
                elif n["occurred_start"] is not None:
                    neighbor_best_date = n["occurred_start"]
                elif n["occurred_end"] is not None:
                    neighbor_best_date = n["occurred_end"]
                elif n["mentioned_at"] is not None:
                    neighbor_best_date = n["mentioned_at"]

                if neighbor_best_date:
                    days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
                    neighbor_temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
                else:
                    neighbor_temporal_proximity = 0.3  # Lower score if no temporal data

                # Boost causal links (same as graph retrieval)
                link_type = n["link_type"]
                if link_type in ("causes", "caused_by"):
                    causal_boost = 2.0
                elif link_type in ("enables", "prevents"):
                    causal_boost = 1.5
                else:
                    causal_boost = 1.0

                # Propagate temporal score through links (decay, with causal boost)
                propagated_temporal = temporal_score * n["weight"] * causal_boost * 0.7

                # Combined temporal score
                combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)

                neighbor_data = dict(n)
                neighbor_data["temporal_score"] = combined_temporal
                neighbor_data["temporal_proximity"] = neighbor_temporal_proximity
                results.append((neighbor_id, neighbor_data))

                # Add to queue for further spreading
                if budget_remaining > 0 and combined_temporal > 0.2:
                    queue.append((dict(n), n["similarity"], combined_temporal))

                if budget_remaining <= 0:
                    break

    return results


async def retrieve_parallel(
    pool,
    query_text: str,
    query_embedding_str: str,
    agent_id: str,
    fact_type: str,
    thinking_budget: int,
    question_date: Optional[datetime] = None,
    query_analyzer: Optional["QueryAnalyzer"] = None
) -> Tuple[List, List, List, Optional[List]]:
    """
    Run 3-way or 4-way parallel retrieval (adds temporal if detected).

    Args:
        pool: Database connection pool
        query_text: Query text
        query_embedding_str: Query embedding as string
        agent_id: Agent ID
        fact_type: Fact type to filter
        thinking_budget: Budget for graph traversal and retrieval limits
        question_date: Optional date when question was asked (for temporal filtering)
        query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)

    Returns:
        Tuple of (semantic_results, bm25_results, graph_results, temporal_results)
        temporal_results is None if no temporal constraint detected
    """
    # Detect temporal constraint
    from .temporal_extraction import extract_temporal_constraint
    import logging
    logger = logging.getLogger(__name__)

    temporal_constraint = extract_temporal_constraint(
        query_text, reference_date=question_date, analyzer=query_analyzer
    )

    async def run_semantic():
        async with acquire_with_retry(pool) as conn:
            return await retrieve_semantic(conn, query_embedding_str, agent_id, fact_type, limit=thinking_budget)

    async def run_bm25():
        async with acquire_with_retry(pool) as conn:
            return await retrieve_bm25(conn, query_text, agent_id, fact_type, limit=thinking_budget)

    async def run_graph():
        async with acquire_with_retry(pool) as conn:
            return await retrieve_graph(conn, query_embedding_str, agent_id, fact_type, budget=thinking_budget)

    async def run_temporal(start_date, end_date):
        async with acquire_with_retry(pool) as conn:
            return await retrieve_temporal(
                conn, query_embedding_str, agent_id, fact_type,
                start_date, end_date, budget=thinking_budget, semantic_threshold=0.4
            )

    # Run retrievals in parallel
    if temporal_constraint:
        start_date, end_date = temporal_constraint
        semantic_results, bm25_results, graph_results, temporal_results = await asyncio.gather(
            run_semantic(), run_bm25(), run_graph(), run_temporal(start_date, end_date)
        )
    else:
        semantic_results, bm25_results, graph_results = await asyncio.gather(
            run_semantic(), run_bm25(), run_graph()
        )
        temporal_results = None

    return semantic_results, bm25_results, graph_results, temporal_results