"""
cortex/dialectic/planner.py — QueryPlanner for the Dialectic Agentic Path (M14).

The QueryPlanner orchestrates multi-step reasoning:
    1. Decompose a complex question into simpler sub-queries (LLM call).
    2. Decide if the accumulated evidence is sufficient to answer (LLM call).
    3. Synthesise a final grounded answer from all collected evidence (LLM call).

Design goals:
    - Each method is a single focused LLM call — easy to test and mock.
    - Stateless: all context is passed in via arguments; the planner holds no
      per-query state between calls.
    - Graceful degradation: on any LLM failure, sensible defaults are returned
      so the engine can still attempt a best-effort answer.

Dependencies:
    cortex.dialectic.prompts  — DECOMPOSE_PROMPT, CONTINUE_PROMPT,
                                FINAL_SYNTHESIS_PROMPT, format_evidence_chain
    cortex.llm.provider       — CompletionRequest (typed LLM request)
"""

from __future__ import annotations

import asyncio
import logging
import re
from typing import Any, Dict, List, Optional

from cortex.dialectic.prompts import (
    DECOMPOSE_PROMPT,
    CONTINUE_PROMPT,
    FINAL_SYNTHESIS_PROMPT,
    format_evidence_chain,
)

logger = logging.getLogger(__name__)

# Fallback when decomposition fails — treat original query as the only sub-query
_DECOMPOSE_FALLBACK = 1

# ---------------------------------------------------------------------------
# QueryPlanner
# ---------------------------------------------------------------------------


class QueryPlanner:
    """Multi-step query planner for the agentic dialectic path.

    Args:
        llm_provider: LLMProvider instance (has async ``.complete()`` method).
        config:       CortexConfig — used for model selection.
    """

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

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

    def _model(self) -> str:
        """Return the configured model name, preferring dialectic_model over extraction_model."""
        base_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 query planning through tier system
        from cortex.llm.router import resolve_routed_model
        return resolve_routed_model(self._llm, "query_planning", base_model) or base_model

    async def _complete(self, system: str, user: str) -> Optional[str]:
        """Call the LLM and return the raw text response.

        Returns ``None`` on any error (caller handles degradation).
        """
        try:
            from cortex.llm.provider import CompletionRequest
        except ImportError:
            CompletionRequest = None  # type: ignore[assignment, misc]

        try:
            if CompletionRequest is not None:
                req = CompletionRequest(
                    system=system,
                    user=user,
                    model=self._model(),
                    response_format="text",
                    temperature=0.1,
                )
                response = await asyncio.wait_for(
                    self._llm.complete(req), timeout=60.0
                )
            else:
                response = await asyncio.wait_for(
                    self._llm.complete(
                        system=system,
                        user=user,
                        model=self._model(),
                    ),
                    timeout=60.0,
                )

            if response is None:
                return None

            text = getattr(response, "text", None) or str(response)
            return text.strip() if text else None

        except asyncio.TimeoutError:
            logger.warning("QueryPlanner LLM call timed out after 60s")
            return None
        except Exception as exc:
            logger.warning("QueryPlanner LLM call failed: %s", exc)
            return None

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

    async def decompose(self, query: str) -> List[str]:
        """Break a complex question into 2-4 simpler sub-queries.

        Calls the LLM with DECOMPOSE_PROMPT and parses numbered list output.
        Falls back to ``[query]`` (the original question as the sole sub-query)
        if parsing fails or the LLM is unavailable.

        Args:
            query: The complex natural-language question to decompose.

        Returns:
            List of sub-query strings (1–4 items). Never empty.
        """
        prompt = DECOMPOSE_PROMPT.format(query=query)
        raw = await self._complete(
            system=(
                "You are a question decomposition assistant. "
                "Your only job is to break complex questions into simpler sub-questions."
            ),
            user=prompt,
        )

        if not raw:
            logger.warning("decompose() got no LLM response — using original query")
            return [query]

        # Parse numbered list: "1. ...", "2. ..." etc.
        sub_queries: List[str] = []
        for line in raw.splitlines():
            line = line.strip()
            # Match lines like "1. What..." or "1) What..."
            match = re.match(r"^\d+[\.\)]\s*(.+)$", line)
            if match:
                sub_queries.append(match.group(1).strip())

        if not sub_queries:
            # LLM returned something but not a numbered list — use raw lines
            sub_queries = [ln.strip() for ln in raw.splitlines() if ln.strip()]

        # Cap at 4 sub-queries, always return at least the original
        sub_queries = sub_queries[:4] if sub_queries else [query]
        logger.debug("decompose() produced %d sub-queries for: %r", len(sub_queries), query)
        return sub_queries

    async def should_continue(
        self,
        evidence: List[Dict[str, Any]],
        query: str,
    ) -> bool:
        """Decide if more retrieval is needed given current evidence.

        Calls the LLM with CONTINUE_PROMPT. Returns ``False`` (stop) on any
        LLM failure to avoid infinite loops.

        Args:
            evidence: List of evidence step dicts accumulated so far.
                      Each dict has ``sub_query``, ``answer``, ``sources``.
            query:    The original complex question.

        Returns:
            ``True`` if more retrieval steps should be attempted,
            ``False`` if the evidence is sufficient or if the LLM is unavailable.
        """
        if not evidence:
            return True  # No evidence yet — always continue

        evidence_block = format_evidence_chain(evidence)
        prompt_text = CONTINUE_PROMPT.format(
            query=query,
            evidence_block=evidence_block,
        )

        raw = await self._complete(
            system=(
                "You are an evidence sufficiency evaluator. "
                "Answer only 'yes' or 'no'."
            ),
            user=prompt_text,
        )

        if not raw:
            logger.warning("should_continue() got no LLM response — stopping")
            return False

        answer = raw.strip().lower()
        if answer.startswith("yes"):
            return False  # Evidence is sufficient → stop
        return True  # "no" or anything else → continue

    async def synthesize_final(
        self,
        query: str,
        evidence_chain: List[Dict[str, Any]],
    ) -> str:
        """Synthesise a final answer from all collected evidence.

        Calls the LLM with FINAL_SYNTHESIS_PROMPT combining all evidence steps.
        Returns a polite fallback string on LLM failure.

        Args:
            query:          The original complex question.
            evidence_chain: All evidence steps from the agentic loop.

        Returns:
            Synthesised answer string.
        """
        _NO_INFO = (
            "I don't have enough information in the provided memories "
            "to answer that question."
        )

        if not evidence_chain:
            return _NO_INFO

        evidence_block = format_evidence_chain(evidence_chain)
        system_prompt = FINAL_SYNTHESIS_PROMPT.format(
            query=query,
            evidence_block=evidence_block,
        )

        raw = await self._complete(
            system=system_prompt,
            user=f"Question: {query}",
        )

        if not raw:
            logger.warning("synthesize_final() got no LLM response — returning fallback")
            return _NO_INFO

        return raw
