"""Read pipeline orchestrator — agent → retrieve (multi-signal) → rerank → format.

Multi-signal retrieval: semantic + keyword + graph run in parallel via
``asyncio.gather()``. The retrieval agent (LLM planner) decides entities
for graph signal, reformulates the query, and detects broad/temporal queries.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from arandu.config import MemoryConfig
from arandu.protocols import EmbeddingProvider, LLMProvider
from arandu.read.context_compression import compress_context
from arandu.read.graph_retrieval import GraphRetrievalResult, retrieve_graph_facts
from arandu.read.query_expansion import expand_query, resolve_query_entities
from arandu.read.reranker import rerank_with_llm
from arandu.read.retrieval import (
    RetrievalCandidate,
    compute_pattern_signal,
    retrieve_candidates,
    retrieve_relevant_events,
)
from arandu.read.retrieval_agent import plan_retrieval
from arandu.read.spreading_activation import spread_activation
from arandu.tracing import PipelineTrace

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result types (also exported from client.py as public API)
# ---------------------------------------------------------------------------
@dataclass
class ScoredFact:
    """A fact with retrieval scores."""

    fact_id: str = ""
    entity_name: str = ""
    attribute_key: str = ""
    fact_text: str = ""
    score: float = 0.0
    scores: dict[str, float] = field(default_factory=dict)


@dataclass
class ReadResult:
    """Result of the read pipeline (internal). Mapped to RetrieveResult by client."""

    facts: list[ScoredFact] = field(default_factory=list)
    context: str = ""
    total_candidates: int = 0
    duration_ms: float = 0.0


# ---------------------------------------------------------------------------
# Context formatting
# ---------------------------------------------------------------------------
def format_context(facts: list[ScoredFact]) -> str:
    """Format scored facts as a prompt-ready context string.

    Output format:
        ## Known facts about the user:
        - Lives in São Paulo (confidence: 0.95)
        - Wife's name is Ana (confidence: 0.90)
    """
    if not facts:
        return ""

    lines = ["## Known facts about the user:"]
    for f in facts:
        conf = f.scores.get("confidence", 0.0)
        line = f"- {f.fact_text}"
        if conf > 0:
            line += f" (confidence: {conf:.2f})"
        lines.append(line)

    return "\n".join(lines)


def _candidate_to_scored_fact(candidate: RetrievalCandidate) -> ScoredFact:
    """Convert a RetrievalCandidate to a ScoredFact for the public API."""
    fact = candidate.fact
    scores = dict(candidate.scores)
    scores["confidence"] = fact.confidence or 0.0

    return ScoredFact(
        fact_id=str(fact.id),
        entity_name=fact.entity_name or "",
        attribute_key=fact.attribute_key or "",
        fact_text=fact.fact_text or "",
        score=candidate.final_score,
        scores=scores,
    )


# ---------------------------------------------------------------------------
# Graph → candidate merging
# ---------------------------------------------------------------------------
def _merge_graph_into_candidates(
    candidates: list[RetrievalCandidate],
    graph_result: GraphRetrievalResult,
) -> list[RetrievalCandidate]:
    """Merge graph-sourced facts into the candidate list.

    Deduplicates by fact ID. Graph facts get a "graph" score.
    """
    existing_ids: set[Any] = {c.fact.id for c in candidates}

    for gf in graph_result.facts:
        fact = gf["fact"]
        if fact.id in existing_ids:
            # Add graph score to existing candidate
            for c in candidates:
                if c.fact.id == fact.id:
                    c.scores["graph"] = gf["score"]
                    break
        else:
            existing_ids.add(fact.id)
            candidates.append(
                RetrievalCandidate(
                    fact=fact,
                    scores={"graph": gf["score"]},
                    final_score=gf["score"],
                )
            )

    return candidates


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
async def run_read_pipeline(
    session: AsyncSession,
    user_id: str,
    query: str,
    llm: LLMProvider,
    embeddings: EmbeddingProvider,
    config: MemoryConfig,
    trace: PipelineTrace | None = None,
    reranker_llm: LLMProvider | None = None,
) -> ReadResult:
    """Execute the read pipeline: agent → retrieve (multi-signal) → rerank → format.

    Multi-signal retrieval runs semantic + keyword + graph in parallel via
    ``asyncio.gather()``. The retrieval agent plans which entities to use
    for the graph signal and reformulates the query.

    Args:
        session: Database session (caller manages transaction).
        user_id: User identifier.
        query: The query to search memory for.
        llm: Injected LLM provider.
        embeddings: Injected embedding provider.
        config: Memory configuration.

    Returns:
        ReadResult with scored facts and formatted context.
    """
    t_start = time.perf_counter()

    # 0. Retrieval agent: plan strategy (entities, query reformulation)
    _t_plan = time.perf_counter() if trace else 0.0
    plan = await plan_retrieval(session, user_id, query, llm)
    if trace:
        trace.add_step(
            "planning",
            {
                "strategy": plan.strategy,
                "entities": plan.entities,
                "similarity_query": plan.similarity_query,
                "broad_query": plan.broad_query,
            },
            trace.elapsed(_t_plan),
        )

    # Short-circuit: skip retrieval for greetings
    if plan.strategy == "skip":
        duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
        return ReadResult(facts=[], context="", total_candidates=0, duration_ms=duration_ms)

    # 0b. Query expansion: entity priming + temporal anchoring
    expanded = await expand_query(session, user_id, query, plan, llm)

    # Use reformulated query if agent provided one
    effective_query = plan.similarity_query or query

    # Enrich query with expanded terms from entity priming
    if expanded.expanded_terms:
        effective_query = f"{effective_query} {' '.join(expanded.expanded_terms)}"

    # Use temporal range from expansion if planner didn't provide one
    as_of_range = plan.as_of_range or expanded.temporal_range

    # 1. Multi-signal retrieval: semantic+keyword, graph, and events in parallel
    retrieve_task = retrieve_candidates(
        session,
        user_id,
        effective_query,
        embeddings,
        config,
    )

    # Event retrieval (parallel with fact retrieval)
    event_task: asyncio.Task[list[dict]] | None = None
    try:
        query_embedding = await embeddings.embed_one(effective_query)
        if query_embedding:
            event_task = asyncio.ensure_future(
                retrieve_relevant_events(
                    session,
                    user_id,
                    query_embedding,
                    config,
                )
            )
    except Exception:
        pass  # Embedding failure — skip event retrieval

    # Unify entities from all sources for graph signal
    _entities_llm = list(plan.entities)
    _entities_expansion = [ek for ek in expanded.primed_entities if ek not in _entities_llm]

    _entities_deterministic: list[str] = []
    try:
        _det_result = await resolve_query_entities(session, user_id, query)
        _already = set(_entities_llm) | set(_entities_expansion)
        _entities_deterministic = [ek for ek in _det_result if ek not in _already]
    except Exception as e:
        logger.warning("read_pipeline: deterministic entity resolution failed: %s", e)

    all_entities = _entities_llm + _entities_expansion + _entities_deterministic

    graph_task: asyncio.Task[GraphRetrievalResult] | None = None
    if all_entities:
        graph_task = asyncio.ensure_future(
            retrieve_graph_facts(
                session,
                user_id,
                all_entities,
                query_text=effective_query,
                broad_query=plan.broad_query,
                as_of_start=as_of_range[0] if as_of_range else None,
                as_of_end=as_of_range[1] if as_of_range else None,
            )
        )

    candidates = await retrieve_task

    # Merge graph results if available
    graph_result = GraphRetrievalResult()
    if graph_task is not None:
        try:
            graph_result = await graph_task
        except Exception as e:
            logger.warning("read_pipeline: graph retrieval failed: %s", e)

    if graph_result.facts:
        candidates = _merge_graph_into_candidates(candidates, graph_result)
        # Re-sort by final_score after merge
        candidates.sort(key=lambda c: c.final_score, reverse=True)

    # 1b. Spreading activation: expand context from seed facts
    _spreading_total = 0
    _spreading_unique = 0
    if candidates:
        seed_fact_ids = [str(c.fact.id) for c in candidates[: config.topk_facts]]
        seed_scores = {str(c.fact.id): c.final_score for c in candidates[: config.topk_facts]}
        spreading_candidates = await spread_activation(
            session,
            user_id,
            seed_fact_ids,
            config,
            seed_scores=seed_scores,
        )
        _spreading_total = len(spreading_candidates)
        if spreading_candidates:
            existing_ids = {c.fact.id for c in candidates}
            for sc in spreading_candidates:
                if sc.fact.id not in existing_ids:
                    existing_ids.add(sc.fact.id)
                    candidates.append(sc)
                    _spreading_unique += 1
            candidates.sort(key=lambda c: c.final_score, reverse=True)

    # 1c. Pattern signal: boost recurring facts
    if candidates:
        candidates = compute_pattern_signal(candidates)

    if trace:
        trace.add_step(
            "retrieval",
            {
                "total_candidates": len(candidates),
                "graph_facts": len(graph_result.facts),
                "graph_edges_traversed": graph_result.edges_traversed,
                "spreading_hops_configured": config.spreading_activation_hops,
                "spreading_candidates_returned": _spreading_total,
                "spreading_candidates_unique": _spreading_unique,
                "entities_unified": all_entities,
                "entities_sources": {
                    "llm": _entities_llm,
                    "expansion": _entities_expansion,
                    "deterministic": _entities_deterministic,
                },
            },
            trace.elapsed(t_start),  # retrieval is the bulk of time
        )

    # 1d. Collect relevant events
    events: list[dict] = []
    if event_task is not None:
        try:
            events = await event_task
        except Exception as e:
            logger.warning("read_pipeline: event retrieval failed: %s", e)

    total_candidates = len(candidates)

    if not candidates:
        duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
        return ReadResult(
            facts=[],
            context="",
            total_candidates=0,
            duration_ms=duration_ms,
        )

    # 2. Rerank via LLM (optional, graceful degradation)
    _t_rerank = time.perf_counter() if trace else 0.0
    reranker_applied = False
    if config.enable_reranker:
        candidates = await rerank_with_llm(query, candidates, reranker_llm or llm, config)
        reranker_applied = True
    if trace:
        trace.add_step(
            "reranking",
            {
                "applied": reranker_applied,
                "candidates_count": len(candidates),
            },
            trace.elapsed(_t_rerank),
        )

    # 3. Take top-K and filter by min_score
    top_candidates = candidates[: config.topk_facts]
    if config.min_score > 0:
        top_candidates = [c for c in top_candidates if c.final_score >= config.min_score]

    # 4. Convert to ScoredFact
    scored_facts = [_candidate_to_scored_fact(c) for c in top_candidates]

    if trace:
        trace.add_step(
            "scoring",
            {
                "top_k": [
                    {
                        "fact_id": f.fact_id,
                        "value": f.fact_text,
                        "score": f.score,
                        "scores": f.scores,
                    }
                    for f in scored_facts
                ],
            },
            0.0,  # scoring is part of conversion, negligible time
        )

    # 5. Context compression: tiered hot/warm/cold within token budget
    fact_dicts = [
        {
            "fact": c.fact,
            "entity": c.fact.entity_name or "",
            "attribute": c.fact.attribute_key or "",
            "value": c.fact.fact_text or "",
            "date": str(getattr(c.fact, "valid_from", "") or ""),
            "score": c.final_score,
        }
        for c in top_candidates
    ]
    _t_compress = time.perf_counter() if trace else 0.0
    try:
        compressed = await compress_context(fact_dicts, events, config)
        context = compressed.context_text or format_context(scored_facts)
        if trace:
            trace.add_step(
                "compression",
                {
                    "context_length": len(context),
                    "facts_included": len(scored_facts),
                },
                trace.elapsed(_t_compress),
            )
    except Exception as exc:
        logger.warning("read_pipeline: context compression failed, falling back: %s", exc)
        context = format_context(scored_facts)
        if trace:
            trace.add_step(
                "compression",
                {
                    "context_length": len(context),
                    "facts_included": len(scored_facts),
                    "fallback": True,
                },
                trace.elapsed(_t_compress),
            )

    duration_ms = round((time.perf_counter() - t_start) * 1000, 1)

    logger.info(
        "read_pipeline: user=%s, plan=%s, total_candidates=%d, graph_facts=%d, returned=%d, latency_ms=%.1f",
        user_id,
        plan.strategy,
        total_candidates,
        len(graph_result.facts),
        len(scored_facts),
        duration_ms,
    )

    return ReadResult(
        facts=scored_facts,
        context=context,
        total_candidates=total_candidates,
        duration_ms=duration_ms,
    )
