"""
cortex/dialectic/engine.py — Dialectic Engine (M14 Fast Path).

Implements the "ask your memory" feature: given a natural-language question,
retrieve relevant memories and synthesise a grounded answer via an LLM.

Two paths:
    ask()              — Fast path: single-shot retrieval + synthesis.
    ask_with_history() — History-aware: passes prior conversation turns to LLM.

DialecticResponse carries:
    answer      — synthesised answer string
    sources     — the memory dicts that were used
    confidence  — float (0-1) derived from retrieval scores
    reasoning   — optional short explanation of how the answer was reached

Design:
    - retrieval_engine is expected to expose .retrieve(query, owner_id, top_k)
      (compatible with both RetrievalPipeline and CortexPlugin.retrieve())
    - llm_provider.complete() is called once per ask() call
    - Graceful degradation: if retrieval returns nothing, returns a polite
      "not enough information" response rather than raising.

Dependencies:
    cortex.dialectic.prompts  — SYNTHESIS_PROMPT, HISTORY_SYSTEM_PROMPT, format_*
    cortex.llm.provider       — CompletionRequest (for typed LLM calls)
"""

from __future__ import annotations

import asyncio
import logging
import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from cortex.dialectic.prompts import (
    SYNTHESIS_PROMPT,
    HISTORY_SYSTEM_PROMPT,
    format_memories,
    format_history,
)
from cortex.dialectic.planner import QueryPlanner

logger = logging.getLogger(__name__)

# Sentinel answer when memories are empty / unhelpful
_NO_INFO = "I don't have enough information in the provided memories to answer that question."

# ---------------------------------------------------------------------------
# Response dataclass
# ---------------------------------------------------------------------------


@dataclass
class DialecticResponse:
    """Result of a DialecticEngine.ask() call.

    Attributes:
        answer:     LLM-synthesised answer grounded in retrieved memories.
        sources:    The memory dicts that were passed to the LLM as context.
                    Each dict has at minimum: subject/predicate/object_value
                    (or content) + category + confidence/score.
        confidence: Float 0–1 derived from the mean retrieval score of sources.
                    0.0 if no sources were found.
        reasoning:  Optional — short explanation of how the answer was derived.
                    Populated when the LLM provides a reasoning preamble.
    """

    answer: str
    sources: List[Dict[str, Any]] = field(default_factory=list)
    confidence: float = 0.0
    reasoning: str = ""


# ---------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------


class DialecticEngine:
    """
    Dialectic Engine — synthesises answers from retrieved memories.

    Args:
        retrieval_engine: Any object with a ``.retrieve(query, owner_id, ...)``
            async method. Compatible with CortexPlugin (used via plugin.retrieve)
            and RetrievalPipeline directly.
        llm_provider:     LLMProvider instance (has async ``.complete()`` method).
        config:           CortexConfig (used for model name, owner_id defaults).
    """

    def __init__(
        self,
        retrieval_engine: Any,
        llm_provider: Any,
        config: Any,
    ) -> None:
        self._retrieval = retrieval_engine
        self._llm = llm_provider
        self._config = config

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    async def _retrieve_memories(
        self,
        query: str,
        owner_id: str,
        top_k: int = 10,
    ) -> List[Dict[str, Any]]:
        """Call retrieval and normalise the result to a list of memory dicts.

        Handles both RetrievalPipeline (returns RetrievalResult dataclass) and
        CortexPlugin.retrieve() (same return type) gracefully.

        Returns an empty list on any retrieval failure.
        """
        try:
            result = await self._retrieval.retrieve(
                query=query,
                owner_id=owner_id,
            )
        except TypeError:
            # Some implementations don't accept owner_id as keyword
            try:
                result = await self._retrieval.retrieve(query=query)
            except Exception as exc:
                logger.warning("Retrieval failed: %s", exc)
                return []
        except Exception as exc:
            logger.warning("Retrieval failed: %s", exc)
            return []

        if result is None:
            return []

        # Normalise: RetrievalResult may have .items (list of dicts/objects)
        items = getattr(result, "items", None)
        if items is None:
            # Try dict-style access
            if hasattr(result, "get"):
                items = result.get("items", [])
            elif isinstance(result, list):
                items = result
            else:
                items = []

        # Trim to top_k
        items = list(items)[:top_k]

        # Convert dataclass / object items to plain dicts
        normalised: List[Dict[str, Any]] = []
        for item in items:
            if isinstance(item, dict):
                normalised.append(item)
            elif hasattr(item, "__dict__"):
                normalised.append(item.__dict__)
            else:
                normalised.append({"content": str(item)})

        return normalised

    @staticmethod
    def _compute_confidence(memories: List[Dict[str, Any]]) -> float:
        """Derive a confidence score from retrieval scores.

        Takes the mean of ``score`` or ``confidence`` values in the memory
        dicts. Returns 0.0 when no scores are available.
        """
        scores = []
        for m in memories:
            s = m.get("score") or m.get("confidence")
            if s is not None:
                try:
                    scores.append(float(s))
                except (TypeError, ValueError):
                    pass
        if not scores:
            return 0.0
        mean_score = sum(scores) / len(scores)
        # Clamp to [0, 1]
        return max(0.0, min(1.0, mean_score))

    async def _synthesise(
        self,
        system_prompt: str,
        query: str,
        history: Optional[List[Dict[str, str]]] = None,
    ) -> str:
        """Call the LLM to synthesise an answer.

        Builds the user turn from *query* and optional *history*.
        Returns the raw text response, or _NO_INFO on failure.
        """
        try:
            from cortex.llm.provider import CompletionRequest
        except ImportError:
            CompletionRequest = None  # type: ignore[assignment, misc]

        # Build user message
        user_parts = []
        if history:
            user_parts.append(format_history(history))
            user_parts.append("")  # blank line separator
        user_parts.append(f"Question: {query}")
        user_message = "\n".join(user_parts)

        # Choose model from config: prefer dialectic_model, fall back to extraction_model
        model = (
            getattr(self._config, "dialectic_model", None)
            or getattr(self._config, "extraction_model", None)
            or "gpt-4.1-nano-2025-04-14"
        )
        # Model routing (#375): route dialectic synthesis through tier system
        from cortex.llm.router import resolve_routed_model
        model = resolve_routed_model(self._llm, "dialectic_synthesis", model)

        try:
            if CompletionRequest is not None:
                req = CompletionRequest(
                    system=system_prompt,
                    user=user_message,
                    model=model,
                    response_format="text",
                    temperature=0.2,
                )
                response = await asyncio.wait_for(
                    self._llm.complete(req), timeout=60.0
                )
            else:
                # Fallback: call with positional args if dataclass not importable
                response = await asyncio.wait_for(
                    self._llm.complete(
                        system=system_prompt,
                        user=user_message,
                        model=model,
                    ),
                    timeout=60.0,
                )

            if response is None:
                return _NO_INFO

            # CompletionResponse has .text; some providers return str directly
            text = getattr(response, "text", None) or str(response)
            return text.strip() if text else _NO_INFO

        except asyncio.TimeoutError:
            logger.error("LLM synthesis timed out after 60s")
            return _NO_INFO
        except Exception as exc:
            logger.error("LLM synthesis failed: %s", exc, exc_info=True)
            return _NO_INFO

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def ask(
        self,
        query: str,
        owner_id: str,
        context: Optional[Dict[str, Any]] = None,
    ) -> DialecticResponse:
        """Fast path: single-shot retrieval + LLM synthesis.

        1. Retrieve top_k=10 relevant memories for *query* / *owner_id*.
        2. Format memories as a numbered context block.
        3. Send to LLM with SYNTHESIS_PROMPT.
        4. Return DialecticResponse with answer + sources + confidence.

        Args:
            query:    Natural-language question.
            owner_id: Memory owner namespace (user/agent UUID).
            context:  Optional extra context dict (reserved for future use).

        Returns:
            DialecticResponse with answer, sources, confidence, reasoning.
        """
        memories = await self._retrieve_memories(query=query, owner_id=owner_id, top_k=10)

        # Short-circuit: no memories → polite not-enough-info response
        if not memories:
            return DialecticResponse(
                answer=_NO_INFO,
                sources=[],
                confidence=0.0,
                reasoning="No relevant memories were retrieved for this query.",
            )

        confidence = self._compute_confidence(memories)
        memories_block = format_memories(memories)

        system_prompt = SYNTHESIS_PROMPT.format(
            owner_id=owner_id,
            memories_block=memories_block,
        )

        answer = await self._synthesise(system_prompt=system_prompt, query=query)

        return DialecticResponse(
            answer=answer,
            sources=memories,
            confidence=confidence,
            reasoning=(
                f"Retrieved {len(memories)} relevant memories "
                f"(mean confidence: {confidence:.2f}) and synthesised an answer."
            ),
        )

    async def ask_with_history(
        self,
        query: str,
        owner_id: str,
        history: List[Dict[str, str]],
    ) -> DialecticResponse:
        """Ask with conversation history for follow-up questions.

        Same as ask() but prepends *history* turns to the LLM user message.
        The retrieval query is still *query* only (not history-augmented).

        Args:
            query:    The latest user question.
            owner_id: Memory owner namespace.
            history:  Prior conversation turns — list of {role, content} dicts.
                      Older messages first; the caller's new query is appended
                      as the final user turn.

        Returns:
            DialecticResponse with answer, sources, confidence, reasoning.
        """
        memories = await self._retrieve_memories(query=query, owner_id=owner_id, top_k=10)

        if not memories:
            return DialecticResponse(
                answer=_NO_INFO,
                sources=[],
                confidence=0.0,
                reasoning="No relevant memories were retrieved for this query.",
            )

        confidence = self._compute_confidence(memories)
        memories_block = format_memories(memories)

        system_prompt = HISTORY_SYSTEM_PROMPT.format(
            owner_id=owner_id,
            memories_block=memories_block,
        )

        answer = await self._synthesise(
            system_prompt=system_prompt,
            query=query,
            history=history,
        )

        return DialecticResponse(
            answer=answer,
            sources=memories,
            confidence=confidence,
            reasoning=(
                f"Retrieved {len(memories)} relevant memories "
                f"(mean confidence: {confidence:.2f}) and synthesised an answer "
                f"with {len(history)} history turns."
            ),
        )

    async def ask_agentic(
        self,
        query: str,
        owner_id: str,
        max_steps: int = 5,
    ) -> DialecticResponse:
        """Agentic path: multi-step retrieval + reasoning for complex questions.

        Unlike the fast path (ask()), the agentic path:
          1. Decomposes the question into sub-queries via the QueryPlanner.
          2. Retrieves relevant memories for each sub-query.
          3. Synthesises an intermediate answer per sub-query.
          4. Checks if the accumulated evidence is sufficient (should_continue).
          5. If not sufficient and budget remains, generates follow-up sub-queries.
          6. Produces a final synthesis from all evidence when done.

        The loop runs at most *max_steps* retrieval iterations (each iteration
        processes one sub-query) to prevent runaway LLM costs.

        Args:
            query:     Natural-language question (may be complex / multi-part).
            owner_id:  Memory owner namespace (user/agent UUID).
            max_steps: Maximum number of retrieval+synthesis iterations.
                       Defaults to 5.  Minimum clamped to 1.

        Returns:
            DialecticResponse with:
                answer     — final synthesised answer from all evidence
                sources    — deduplicated list of all memory dicts used
                confidence — mean confidence across all retrieved memories
                reasoning  — human-readable summary of the agentic reasoning chain
        """
        max_steps = max(1, max_steps)
        planner = QueryPlanner(llm_provider=self._llm, config=self._config)

        # --- Step 1: Decompose the question into sub-queries ----------------
        sub_queries = await planner.decompose(query)
        logger.debug(
            "ask_agentic: decomposed %r into %d sub-queries: %s",
            query,
            len(sub_queries),
            sub_queries,
        )

        # --- Steps 2-4: Retrieval + intermediate synthesis loop -------------
        evidence_chain: List[Dict[str, Any]] = []
        all_sources: List[Dict[str, Any]] = []
        steps_used = 0

        query_queue = list(sub_queries)  # mutable queue for follow-up queries

        while query_queue and steps_used < max_steps:
            sub_query = query_queue.pop(0)
            steps_used += 1

            # Retrieve memories for this sub-query
            memories = await self._retrieve_memories(
                query=sub_query,
                owner_id=owner_id,
                top_k=10,
            )

            # Track all sources (deduplicate by id if available)
            seen_ids = {s.get("id") for s in all_sources if s.get("id")}
            for mem in memories:
                mem_id = mem.get("id")
                if mem_id is None or mem_id not in seen_ids:
                    all_sources.append(mem)
                    if mem_id is not None:
                        seen_ids.add(mem_id)

            # Synthesise an intermediate answer for this sub-query
            if memories:
                memories_block = format_memories(memories)
                from cortex.dialectic.prompts import SYNTHESIS_PROMPT as _SP
                system_prompt = _SP.format(
                    owner_id=owner_id,
                    memories_block=memories_block,
                )
                intermediate_answer = await self._synthesise(
                    system_prompt=system_prompt,
                    query=sub_query,
                )
            else:
                intermediate_answer = (
                    "No relevant memories found for this sub-question."
                )

            evidence_chain.append(
                {
                    "sub_query": sub_query,
                    "answer": intermediate_answer,
                    "sources": memories,
                    "step": steps_used,
                }
            )

            # --- Step 4: Check if evidence is sufficient --------------------
            if steps_used < max_steps and not query_queue:
                # Budget remains and queue is empty — ask if we should continue
                still_needed = await planner.should_continue(
                    evidence=evidence_chain,
                    query=query,
                )
                if not still_needed:
                    logger.debug(
                        "ask_agentic: planner says evidence sufficient after %d steps",
                        steps_used,
                    )
                    break
                # Generate follow-up sub-queries if evidence is insufficient
                # (only if we still have budget)
                if steps_used < max_steps:
                    follow_ups = await planner.decompose(
                        f"What additional information is needed to answer: {query}"
                    )
                    # Add new unique sub-queries to the queue (cap to remaining budget)
                    existing = {e["sub_query"] for e in evidence_chain}
                    for fq in follow_ups:
                        if fq not in existing and len(query_queue) < (max_steps - steps_used):
                            query_queue.append(fq)

        # --- Step 5: Final synthesis -----------------------------------------
        if not evidence_chain:
            return DialecticResponse(
                answer=(
                    "I don't have enough information in the provided memories "
                    "to answer that question."
                ),
                sources=[],
                confidence=0.0,
                reasoning="No evidence was collected during the agentic loop.",
            )

        final_answer = await planner.synthesize_final(
            query=query,
            evidence_chain=evidence_chain,
        )

        confidence = self._compute_confidence(all_sources)
        reasoning = (
            f"Agentic path: {steps_used} retrieval step(s), "
            f"{len(all_sources)} total source(s), "
            f"mean confidence {confidence:.2f}. "
            f"Sub-queries: {[e['sub_query'] for e in evidence_chain]}"
        )

        return DialecticResponse(
            answer=final_answer,
            sources=all_sources,
            confidence=confidence,
            reasoning=reasoning,
        )
