"""Retrieval Agent — LLM decides retrieval strategy for memory facts.

Before running cosine similarity, the agent analyzes the user query and chooses:
- multi_signal: all signals run in parallel (default)
- skip: no retrieval needed (greetings, casual)

Produces a ``RetrievalPlan`` with entities for graph signal, reformulated
similarity_query, broad_query flag, and optional as_of_range for time-travel.
"""

from __future__ import annotations

import json
import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from typing import Any

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

from arandu.models import MemoryFact
from arandu.protocols import LLMProvider

logger = logging.getLogger(__name__)

RETRIEVAL_AGENT_MAX_TOKENS = 250
RETRIEVAL_AGENT_TIMEOUT_SEC = 3.0
SCHEMA_CACHE_TTL_SEC = 300  # 5 minutes
ANAPHORA_CONTEXT_CONFIDENCE_MIN = 0.55

_schema_cache: dict[str, tuple[list[tuple[str, str]], float]] = {}


@dataclass
class PatternQuery:
    """One pattern-based query: entity_key LIKE pattern, optional attribute filter.

    Attributes:
        entity_pattern: SQL LIKE pattern for entity_key matching.
        attribute_filter: Optional attribute key filter (always None in V5).
    """

    entity_pattern: str
    attribute_filter: str | None = None


@dataclass
class RetrievalPlan:
    """Output of the retrieval agent.

    V5 simplified: the pipeline no longer selects a single strategy.
    Instead, all signals (semantic, graph, keyword) run in parallel.
    The plan provides entities for graph signal, reformulated query for
    semantic signal, and budget/time-travel parameters.

    Attributes:
        strategy: "multi_signal" (default) or "skip".
        entities: Detected entity_keys for graph signal.
        pattern_queries: Pattern queries for keyword signal.
        similarity_query: Reformulated query for semantic signal.
        max_facts: Budget per signal.
        reason: Why this plan was chosen.
        latency_ms: Time spent planning.
        as_of_range: Optional time-travel window (start, end).
        broad_query: True for comprehensive queries.
    """

    strategy: str = "multi_signal"
    entities: list[str] = field(default_factory=list)
    pattern_queries: list[PatternQuery] = field(default_factory=list)
    similarity_query: str | None = None
    max_facts: int = 50
    reason: str = ""
    latency_ms: float = 0.0
    as_of_range: tuple[datetime, datetime] | None = None
    broad_query: bool = False


def _get_schema_prefixes(pairs: list[tuple[str, str]]) -> list[str]:
    """Return sorted list of group prefixes from entity_key pairs."""
    prefixes: set[str] = set()
    for entity_key, _ in pairs:
        if not entity_key:
            continue
        if ":" in entity_key:
            prefixes.add(entity_key.rsplit(":", 1)[0])
        else:
            prefixes.add(entity_key)
    return sorted(prefixes)


async def _get_user_schema(session: AsyncSession, user_id: str) -> list[tuple[str, str]]:
    """Fetch distinct (entity_key, attribute_key) pairs for user. Cached 5 min per user_id."""
    now = time.time()
    if user_id in _schema_cache:
        pairs, expiry = _schema_cache[user_id]
        if now < expiry:
            return pairs
    stmt = (
        select(MemoryFact.entity_key, MemoryFact.attribute_key)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.valid_to.is_(None),
        )
        .distinct()
    )
    result = await session.execute(stmt)
    pairs = [(r[0], r[1]) for r in result.all() if r[0] and r[1]]
    _schema_cache[user_id] = (pairs, now + SCHEMA_CACHE_TTL_SEC)
    return pairs


def _format_grouped_schema(pairs: list[tuple[str, str]]) -> str:
    """Group (entity_key, attribute_key) by entity prefix and format for prompt."""
    if not pairs:
        return "(no data)"
    groups: dict[str, tuple[set[str], set[str]]] = defaultdict(lambda: (set(), set()))
    for entity_key, attribute_key in pairs:
        prefix = entity_key.rsplit(":", 1)[0] if ":" in entity_key else entity_key
        ents, attrs = groups[prefix]
        ents.add(entity_key)
        attrs.add(attribute_key)
    lines: list[str] = []
    for prefix in sorted(groups.keys()):
        ents, attrs = groups[prefix]
        ent_list = sorted(ents)
        if len(ent_list) <= 4:
            ent_label = ", ".join(ent_list)
        else:
            ent_label = ", ".join(ent_list[:3]) + f", ... ({len(ent_list)} entities)"
        attr_label = ", ".join(sorted(attrs))
        n = len(ent_list)
        lines.append(f'Group "{prefix}" ({n} entity(ies): {ent_label}):')
        lines.append(f"  Attributes: {attr_label}")
    return "\n".join(lines)


def _build_system_prompt(schema_grouped_text: str, schema_prefixes: list[str]) -> str:
    """Build system prompt from real schema; no hardcoded entity_key examples."""
    prefixes_help = (
        f"Available groups for THIS user: {', '.join(schema_prefixes)}."
        if schema_prefixes
        else "This user has no groups in the schema yet (empty list)."
    )
    example_prefix = schema_prefixes[0] if schema_prefixes else "group"
    pattern_example = (
        f'{{"strategy":"pattern","pattern_queries":[{{"entity_pattern":"{example_prefix}:%","attribute_filter":null}}],'
        '"similarity_query":null,"max_facts":50,"reason":"aggregation"}'
    )
    return f"""You are a retrieval agent. Your task is to decide the best strategy for searching the user's memory.

The user has the following stored data (each group shows entities and their attributes). This is the REAL schema for this user — use ONLY these groups.

{schema_grouped_text}

{prefixes_help}
NEVER invent prefixes; use only those listed above. For aggregation queries ("who are my X?", "my friends", "my team"), choose the prefix that best matches the intent and use entity_pattern = prefix + ":%".

MANDATORY RULES:
- For strategy="pattern": attribute_filter MUST always be null. Fetch all attributes for matched entities.
- Greetings or casual conversation with no memory need: use strategy="skip".
- If the query asks for ALL items of a type, use strategy="pattern" with entity_pattern = one of the prefixes above + ":%".
- If the query is about ONE specific person/thing, use strategy="similarity".
- entity_key uses ":" as separator. Wildcard at end: "<prefix>:%".

TIME-TRAVEL (temporal queries):
If the query mentions a past period, add "as_of_range": ["start_iso", "end_iso"] to the JSON.
- Use ISO-8601 dates with timezone offset.
- Vague temporal references without concrete dates: omit as_of_range.

MANDATORY FORMAT (valid JSON, no markdown):

Aggregation (example with prefix from this schema):
{pattern_example}

Specific question:
{{"strategy":"similarity","pattern_queries":[],"similarity_query":"search terms describing the fact","max_facts":20,"reason":"specific search"}}

Greeting:
{{"strategy":"skip","pattern_queries":[],"similarity_query":null,"max_facts":0,"reason":"greeting"}}

QUERY REFORMULATION (MANDATORY):
similarity_query is NOT a copy of the user's question. It is a DESCRIPTION of the type of fact that answers the question.

Respond ONLY with valid JSON, no markdown."""


def _parse_plan(raw: str | None, query_text: str) -> RetrievalPlan | None:
    """Parse LLM JSON into RetrievalPlan. Returns None on parse error."""
    raw = (raw or "").strip()
    if not raw:
        return None
    # Strip markdown code block if present
    if raw.startswith("```"):
        lines = raw.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        raw = "\n".join(lines)
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None
    if not isinstance(data, dict):
        return None

    # V5: map all non-skip strategies to "multi_signal"
    raw_strategy = data.get("strategy") or "similarity"
    strategy = "skip" if raw_strategy == "skip" else "multi_signal"

    # Accept "pattern_queries" or "queries"; accept list or single dict
    raw_pq = data.get("pattern_queries") or data.get("queries") or []
    if isinstance(raw_pq, dict):
        raw_pq = [raw_pq]
    pattern_queries: list[PatternQuery] = []
    for pq in raw_pq:
        if not isinstance(pq, dict):
            continue
        entity_pattern = (pq.get("entity_pattern") or pq.get("entityPattern") or "").strip()
        if not entity_pattern:
            continue
        pattern_queries.append(
            PatternQuery(
                entity_pattern=entity_pattern,
                attribute_filter=None,
            )
        )

    # Fallback: flat format (entity_key at top level)
    if not pattern_queries and data.get("entity_key"):
        entity_pattern = str(data["entity_key"]).strip()
        if entity_pattern:
            pattern_queries = [PatternQuery(entity_pattern=entity_pattern, attribute_filter=None)]

    similarity_query = data.get("similarity_query")
    if similarity_query is not None:
        similarity_query = str(similarity_query).strip() or query_text
    else:
        similarity_query = query_text if strategy not in ("pattern", "skip") else None

    max_facts = data.get("max_facts")
    if max_facts is None:
        max_facts = 50
    try:
        max_facts = max(5, min(200, int(max_facts)))
    except (TypeError, ValueError):
        max_facts = 50

    reason = str(data.get("reason") or "").strip() or "agent"

    # Parse as_of_range for time-travel queries
    as_of_range: tuple[datetime, datetime] | None = None
    raw_range = data.get("as_of_range")
    if isinstance(raw_range, list) and len(raw_range) == 2:
        try:
            start_dt = datetime.fromisoformat(str(raw_range[0]))
            end_dt = datetime.fromisoformat(str(raw_range[1]))
            if start_dt.tzinfo is None:
                start_dt = start_dt.replace(tzinfo=UTC)
            if end_dt.tzinfo is None:
                end_dt = end_dt.replace(tzinfo=UTC)
            now_utc = datetime.now(UTC)
            if start_dt <= end_dt and start_dt <= now_utc + timedelta(days=1):
                as_of_range = (start_dt, end_dt)
        except (ValueError, TypeError):
            pass

    # V5: extract entity_keys from pattern queries for graph retrieval
    entities: list[str] = []
    raw_entities = data.get("entities") or []
    if isinstance(raw_entities, list):
        entities = [str(e).strip() for e in raw_entities if isinstance(e, str) and e.strip()]
    for pq in pattern_queries:
        ep = pq.entity_pattern.rstrip(":%").rstrip("%")
        if ep and ep not in entities:
            entities.append(ep)

    return RetrievalPlan(
        strategy=strategy,
        entities=entities,
        pattern_queries=pattern_queries,
        similarity_query=similarity_query,
        max_facts=max_facts,
        reason=reason,
        latency_ms=0.0,
        as_of_range=as_of_range,
    )


async def plan_retrieval(
    session: AsyncSession,
    user_id: str,
    query_text: str,
    llm: LLMProvider,
    *,
    session_context: Any | None = None,
) -> RetrievalPlan:
    """Call LLM to decide retrieval strategy.

    On timeout/parse/API error returns fallback plan (strategy=multi_signal,
    reason=agent_fallback).

    Args:
        session: Database session.
        user_id: User identifier.
        query_text: The user's query.
        llm: Injected LLM provider (via Protocol).
        session_context: Optional session digest with anaphora context.

    Returns:
        RetrievalPlan with strategy, entities, and query parameters.
    """
    if not query_text or not query_text.strip():
        return RetrievalPlan(
            strategy="multi_signal",
            similarity_query=query_text,
            reason="empty_query",
            latency_ms=0.0,
        )

    start = time.perf_counter()

    try:
        pairs = await _get_user_schema(session, user_id)
        schema_grouped_text = _format_grouped_schema(pairs)
    except Exception as e:
        logger.warning("retrieval_agent: schema fetch failed: %s", e)
        return RetrievalPlan(
            strategy="multi_signal",
            similarity_query=query_text,
            reason="agent_fallback",
            latency_ms=(time.perf_counter() - start) * 1000,
        )

    schema_prefixes = _get_schema_prefixes(pairs)
    system = _build_system_prompt(schema_grouped_text, schema_prefixes)
    user_msg = f'User query: "{query_text[:500]}"'

    # Inject session context only when anaphora gate passes.
    if session_context and getattr(session_context, "turn_count", 0) > 1:
        summary = getattr(session_context, "session_summary", "") or ""
        is_anaphoric = bool(getattr(session_context, "is_anaphoric", False))
        confidence = float(getattr(session_context, "anaphora_confidence", 1.0) or 0.0)
        if summary and is_anaphoric and confidence >= ANAPHORA_CONTEXT_CONFIDENCE_MIN:
            user_msg += f"\n\nSession context (recent user messages): {summary[:300]}"

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user_msg},
    ]

    try:
        response_content = await llm.complete(
            messages,
            temperature=0,
            max_tokens=RETRIEVAL_AGENT_MAX_TOKENS,
        )
    except Exception as e:
        logger.warning("retrieval_agent: llm error: %s", e)
        return RetrievalPlan(
            strategy="multi_signal",
            similarity_query=query_text,
            reason="agent_fallback",
            latency_ms=(time.perf_counter() - start) * 1000,
        )

    latency_ms = (time.perf_counter() - start) * 1000
    logger.info("retrieval_agent raw response: %s", repr(response_content))

    plan = _parse_plan(response_content, query_text)
    logger.info("retrieval_agent parsed plan: %s", plan)

    if plan is None:
        logger.warning("retrieval_agent: parse failed, fallback to multi_signal")
        return RetrievalPlan(
            strategy="multi_signal",
            similarity_query=query_text,
            reason="agent_fallback",
            latency_ms=latency_ms,
        )

    plan.latency_ms = latency_ms

    # Detect broad queries from LLM reason field
    reason_lower = plan.reason.lower()
    if "aggreg" in reason_lower or "broad" in reason_lower or "overview" in reason_lower:
        plan.broad_query = True

    return plan
