"""Implement the Agentic RAG pipeline for multi-step reasoning (P0-1).

Architecture
------------
The pipeline adds four explicit reasoning steps before answer generation:

  1. Intent Analysis  — classify query type (simple_qa / compare / multi_hop / ambiguous)
  2. Query Planning   — for complex intents, decompose into focused sub-queries
  3. Iterative Retrieval — retrieve for each sub-query, merge with deduplication
  4. Grounding Check  — verify retrieved context is sufficient; refine query when not
  5. Answer Generation — stream final answer with citations (reuses AnswerGenerator)

Each step emits SSE ``thinking`` events so the frontend can render the live
reasoning trace.  The pipeline is entirely synchronous and returns a generator so
it plugs into the same ``iter_in_thread`` pattern used by the existing chat routes.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Generator

from loguru import logger

from app.application.knowledge import KnowledgeRetrievalService
from app.application.agent.hyde_expander import HyDEExpander
from app.config.settings import settings
from app.domain.conversation import ConversationStore
from app.domain.retrieval import RetrievedChunk
from app.infrastructure.llm.openai_compatible_answer_generator import OpenAICompatibleAnswerGenerator
from app.services.llm.llm_factory import get_llm_client

# ── Prompts ───────────────────────────────────────────────────────────────────
# Each prompt is kept module-level for easy review and fine-tuning.

_INTENT_SYSTEM = (
    "You are a query classifier for a Chinese regulatory compliance knowledge base.\n\n"
    "Classify the query into exactly one of:\n"
    '- "simple_qa"   : Single-hop, factual question about one regulation or clause\n'
    '- "compare"     : Comparison between two or more regulations, standards, or versions\n'
    '- "multi_hop"   : Requires chaining facts across multiple regulations to answer\n'
    '- "ambiguous"   : Too vague or broad to retrieve effectively\n\n'
    "Return ONLY valid JSON — no markdown, no extra text:\n"
    '{"type": "...", "reason": "one sentence", "requires_decomposition": true/false}\n\n'
    '"requires_decomposition" must be true for compare and multi_hop types.'
)

_PLAN_SYSTEM = (
    "You are a query planner for a Chinese regulatory compliance knowledge base.\n\n"
    "Decompose the query into 2-4 focused, self-contained sub-queries that together fully "
    "address the original question.  Each sub-query must target one specific regulation, "
    "clause, or concept and be independently searchable.\n\n"
    "Return ONLY a valid JSON array — no markdown, no extra text:\n"
    '["sub-query 1", "sub-query 2", ...]'
)

_GROUNDING_SYSTEM = (
    "You are a grounding verifier for a regulatory compliance QA system.\n\n"
    "Given a query and retrieved regulation passages, decide whether the passages contain "
    "sufficient, accurate information to answer the query.\n\n"
    "Return ONLY valid JSON — no markdown, no extra text:\n"
    '{"sufficient": true/false, "confidence": 0.0-1.0, "reason": "one sentence", '
    '"refined_query": "a more specific search query if not sufficient, else null"}'
)


# ── Result dataclasses ────────────────────────────────────────────────────────

@dataclass
class IntentResult:
    """Capture the output of the intent-analysis step."""

    type: str = "simple_qa"
    reason: str = ""
    requires_decomposition: bool = False


@dataclass
class GroundingResult:
    """Capture the output of the grounding-check step."""

    sufficient: bool = True
    confidence: float = 1.0
    reason: str = ""
    refined_query: str | None = None


# ── Service ───────────────────────────────────────────────────────────────────

class AgenticConversationService:
    """Multi-step Agentic RAG pipeline with live reasoning trace via SSE.

    The service is intentionally synchronous so it can be wrapped in
    ``iter_in_thread`` by the route layer without any async boilerplate.
    """

    def __init__(
        self,
        *,
        retrieval_service: KnowledgeRetrievalService,
        answer_generator: OpenAICompatibleAnswerGenerator,
        conversation_store: ConversationStore,
    ) -> None:
        """Initialise with injected dependencies from the composition root."""
        self.retrieval_service = retrieval_service
        self.answer_generator = answer_generator
        self.conversation_store = conversation_store
        # HyDE expander is stateless — one instance shared for all requests.
        self._hyde = HyDEExpander()

    # ── Private helpers ───────────────────────────────────────────────────────

    def _llm_json(
        self,
        system: str,
        user: str,
        provider: str | None,
        model: str | None,
        max_tokens: int = 300,
    ) -> dict | list | None:
        """Call the LLM with a JSON-only prompt and return the parsed result.

        Returns ``None`` on any API or parse failure so callers can degrade
        gracefully without raising.
        """
        client = get_llm_client(
            provider=provider or settings.llm_provider,
            model=model or settings.llm_model,
        )
        resp = client.chat(
            [{"role": "system", "content": system}, {"role": "user", "content": user}],
            max_tokens=max_tokens,
            temperature=0.1,
        )
        if not resp.is_success:
            logger.warning("AgenticService LLM call failed: {}", resp.error)
            return None
        try:
            raw = resp.content.strip()
            # Strip accidental markdown code fences the model may add.
            if raw.startswith("```"):
                parts = raw.split("```")
                raw = parts[1] if len(parts) > 1 else raw
                if raw.startswith("json"):
                    raw = raw[4:]
            return json.loads(raw.strip())
        except (json.JSONDecodeError, IndexError) as exc:
            logger.debug("AgenticService JSON parse failed: {} | raw={}", exc, resp.content[:200])
            return None

    def _analyze_intent(
        self, query: str, provider: str | None, model: str | None
    ) -> IntentResult:
        """Classify query intent to select the appropriate retrieval strategy."""
        data = self._llm_json(
            _INTENT_SYSTEM,
            f"Query: {query}",
            provider,
            model,
            max_tokens=settings.agentic_intent_max_tokens,
        )
        if isinstance(data, dict):
            return IntentResult(
                type=str(data.get("type", "simple_qa")),
                reason=str(data.get("reason", "")),
                requires_decomposition=bool(data.get("requires_decomposition", False)),
            )
        return IntentResult(type="simple_qa", reason="fallback — classifier returned no JSON", requires_decomposition=False)

    def _plan_queries(
        self, query: str, intent_type: str, provider: str | None, model: str | None
    ) -> list[str]:
        """Decompose a complex query into focused, independently-retrievable sub-queries."""
        data = self._llm_json(
            _PLAN_SYSTEM,
            f"Original query ({intent_type}): {query}",
            provider,
            model,
            max_tokens=settings.agentic_plan_max_tokens,
        )
        if isinstance(data, list) and data:
            # Cap at configured maximum to keep latency predictable.
            return [str(q) for q in data[:settings.agentic_max_sub_queries] if q]
        return [query]

    def _check_grounding(
        self,
        query: str,
        chunks: list[RetrievedChunk],
        provider: str | None,
        model: str | None,
    ) -> GroundingResult:
        """Verify whether retrieved chunks are sufficient to ground an accurate answer.

        Uses a fast score-threshold heuristic first; falls back to an LLM call only
        when scores are borderline so that the happy-path adds no extra latency.
        """
        if not chunks:
            return GroundingResult(
                sufficient=False,
                confidence=0.0,
                reason="未检索到相关内容",
                refined_query=None,
            )

        avg_score = sum(c.score for c in chunks) / len(chunks)
        # Fast path: high-confidence retrieval → skip extra LLM call.
        if avg_score > settings.agentic_grounding_threshold and len(chunks) >= 3:
            return GroundingResult(
                sufficient=True,
                confidence=round(avg_score, 3),
                reason="检索置信度充足，无需二次查询",
                refined_query=None,
            )

        # LLM-based grounding check for borderline retrievals.
        context_preview = "\n".join(
            f"[{i + 1}] (score={c.score:.2f}) {c.text[:200]}" for i, c in enumerate(chunks[:5])
        )
        data = self._llm_json(
            _GROUNDING_SYSTEM,
            f"Query: {query}\n\nRetrieved passages:\n{context_preview}",
            provider,
            model,
            max_tokens=settings.agentic_grounding_max_tokens,
        )
        if isinstance(data, dict):
            return GroundingResult(
                sufficient=bool(data.get("sufficient", True)),
                confidence=float(data.get("confidence", 0.5)),
                reason=str(data.get("reason", "")),
                refined_query=data.get("refined_query") or None,
            )
        return GroundingResult(sufficient=True, confidence=0.5, reason="grounding check skipped (parse error)", refined_query=None)

    @staticmethod
    def _intent_to_template(intent_type: str) -> str:
        """Map an intent type to the best prompt template name for answer generation."""
        mapping = {
            "compare": "comparison",
            "multi_hop": "compliance_qa",
            "simple_qa": "compliance_qa",
            "ambiguous": "compliance_qa",
        }
        return mapping.get(intent_type, "compliance_qa")

    @staticmethod
    def _deduplicate(chunks: list[RetrievedChunk], max_chunks: int) -> list[RetrievedChunk]:
        """Remove duplicate chunk IDs, preserving first-occurrence order up to max_chunks."""
        seen: set[str] = set()
        result: list[RetrievedChunk] = []
        for chunk in chunks:
            if chunk.chunk_id not in seen:
                seen.add(chunk.chunk_id)
                result.append(chunk)
                if len(result) >= max_chunks:
                    break
        return result

    # ── Public interface ──────────────────────────────────────────────────────

    def stream_agentic_chat(
        self,
        *,
        query: str,
        session_id: str | None = None,
        filters: str | None = None,
        provider: str | None = None,
        model: str | None = None,
        top_k: int = 5,
        context_text: str | None = None,
        context_filename: str | None = None,
    ) -> tuple[str, Generator[dict, None, None]]:
        """Run the full Agentic RAG pipeline and return ``(session_id, event_generator)``.

        When context_text is provided (user-attached document) it is:
        - Summarised and prepended to the intent-analysis prompt so the classifier
          understands what kind of question is being asked.
        - Treated as baseline grounding so the pipeline skips unnecessary retries
          when the document itself is the primary source.
        - Passed to the answer generator so the LLM sees the full document alongside
          retrieved regulation chunks.

        The generator yields SSE event dicts compatible with the route's
        ``iter_in_thread`` pattern.
        """
        session = self.conversation_store.get_session(session_id) if session_id else None
        if session is None:
            session = self.conversation_store.create_session()
        self.conversation_store.save_message(session.session_id, role="user", content=query)
        history = [{"role": msg.role, "content": msg.content} for msg in session.messages[-10:]]
        active_session_id = session.session_id

        # Build a brief document summary for classifier/planner prompts (avoid
        # passing the full text which could overwhelm small-context LLMs).
        _doc_summary: str = ""
        if context_text and context_text.strip():
            _doc_label = context_filename or "document"
            _preview = context_text.strip()[:400]
            _doc_summary = f"[User has attached document: {_doc_label}]\nDocument preview: {_preview}…\n\n"

        def event_stream() -> Generator[dict, None, None]:
            """Execute all pipeline steps and yield SSE events."""
            # ── Step 1: Intent Analysis ──────────────────────────────────────
            yield {"event": "thinking", "data": {"step": "intent_analysis", "status": "running"}}
            # Prepend doc summary so the classifier knows what the user is asking about
            intent_user_msg = f"{_doc_summary}Query: {query}" if _doc_summary else f"Query: {query}"
            data = self._llm_json(
                _INTENT_SYSTEM, intent_user_msg, provider, model,
                max_tokens=settings.agentic_intent_max_tokens,
            )
            if isinstance(data, dict):
                intent = IntentResult(
                    type=str(data.get("type", "simple_qa")),
                    reason=str(data.get("reason", "")),
                    requires_decomposition=bool(data.get("requires_decomposition", False)),
                )
            else:
                intent = IntentResult(type="simple_qa", reason="fallback", requires_decomposition=False)
            logger.debug("Agentic intent: type={} decompose={}", intent.type, intent.requires_decomposition)
            yield {
                "event": "thinking",
                "data": {
                    "step": "intent_analysis",
                    "status": "done",
                    "intent_type": intent.type,
                    "reason": intent.reason,
                    "requires_decomposition": intent.requires_decomposition,
                },
            }

            # ── Step 2: Query Planning ───────────────────────────────────────
            sub_queries: list[str] = [query]
            if intent.requires_decomposition:
                yield {"event": "thinking", "data": {"step": "query_planning", "status": "running"}}
                plan_user_msg = f"{_doc_summary}Original query ({intent.type}): {query}" if _doc_summary else f"Original query ({intent.type}): {query}"
                data_plan = self._llm_json(
                    _PLAN_SYSTEM, plan_user_msg, provider, model,
                    max_tokens=settings.agentic_plan_max_tokens,
                )
                if isinstance(data_plan, list) and data_plan:
                    sub_queries = [str(q) for q in data_plan[:settings.agentic_max_sub_queries] if q]
                logger.debug("Agentic sub-queries ({}): {}", len(sub_queries), sub_queries)
                yield {
                    "event": "thinking",
                    "data": {"step": "query_planning", "status": "done", "sub_queries": sub_queries},
                }

            # ── Step 3: Iterative Retrieval ──────────────────────────────────
            # Always retrieve using the user's original question (NOT the document
            # text) so embedding quality is preserved for regulation matching.
            # HyDE enriches the retrieval query with a short hypothetical answer
            # to close the vocabulary gap between terse queries and long documents.
            candidate_k = max(top_k * 3, 15)
            all_chunks: list[RetrievedChunk] = []

            # For simple_qa with a single query, HyDE gives the biggest benefit
            # (bridging vague/colloquial questions to formal document language).
            # For compare/multi_hop, the planner already decomposed into precise
            # sub-queries, so HyDE is less critical but still applied per sub-query.
            for idx, sq in enumerate(sub_queries, start=1):
                yield {
                    "event": "thinking",
                    "data": {"step": "retrieving", "status": "running", "query": sq, "index": idx, "total": len(sub_queries)},
                }
                # HyDE expansion: generate hypothetical answer, embed it for retrieval.
                # Falls back to original sub-query if LLM call fails.
                retrieval_query = self._hyde.expand(sq)
                chunks = self.retrieval_service.retrieve(query=retrieval_query, top_k=candidate_k, filters=filters)
                all_chunks.extend(chunks)
                yield {
                    "event": "thinking",
                    "data": {"step": "retrieving", "status": "done", "query": sq, "index": idx, "total": len(sub_queries), "found": len(chunks)},
                }

            unique_chunks = self._deduplicate(all_chunks, max_chunks=top_k * 4)

            # ── Step 4: Grounding Check ──────────────────────────────────────
            yield {"event": "thinking", "data": {"step": "grounding_check", "status": "running"}}

            # When the user has attached a document, the document itself provides
            # baseline grounding — skip the re-query loop to avoid the LLM asking
            # "please provide the document text" as a refined query.
            if context_text and context_text.strip():
                grounding = GroundingResult(
                    sufficient=True,
                    confidence=0.95,
                    reason="用户已附件上传文档，以文档内容为基础作答",
                    refined_query=None,
                )
            else:
                grounding = self._check_grounding(query, unique_chunks, provider, model)

            yield {
                "event": "thinking",
                "data": {
                    "step": "grounding_check",
                    "status": "done",
                    "sufficient": grounding.sufficient,
                    "confidence": grounding.confidence,
                    "reason": grounding.reason,
                },
            }

            # Only retry from vector store when no document is attached and grounding failed
            if not grounding.sufficient and grounding.refined_query and not context_text:
                logger.info("Grounding insufficient — re-querying: {}", grounding.refined_query)
                yield {
                    "event": "thinking",
                    "data": {"step": "retrieving", "status": "running", "query": grounding.refined_query, "index": 1, "total": 1, "retry": True},
                }
                # Apply HyDE to the refined query as well for better retrieval.
                refined_hyde_query = self._hyde.expand(grounding.refined_query)
                refined_chunks = self.retrieval_service.retrieve(query=refined_hyde_query, top_k=candidate_k, filters=filters)
                all_chunks.extend(refined_chunks)
                unique_chunks = self._deduplicate(all_chunks, max_chunks=top_k * 4)
                yield {
                    "event": "thinking",
                    "data": {"step": "retrieving", "status": "done", "query": grounding.refined_query, "index": 1, "total": 1, "found": len(refined_chunks), "retry": True},
                }

            final_chunks = unique_chunks[:top_k]

            # ── Step 5: Answer Generation ────────────────────────────────────
            sources_payload = [s.__dict__ for s in self.answer_generator._sources(final_chunks)]
            yield {"event": "sources", "data": sources_payload}

            answer_parts: list[str] = []
            for event in self.answer_generator.stream_generate(
                query=query,
                retrieved_chunks=final_chunks,
                history=history,
                provider=provider,
                model=model,
                prompt_template=self._intent_to_template(intent.type),
                context_text=context_text,
                context_filename=context_filename,
            ):
                if event.get("event") == "content":
                    answer_parts.append(str(event.get("data", "")))
                yield event

            full_answer = "".join(answer_parts)
            self.conversation_store.save_message(
                active_session_id,
                role="assistant",
                content=full_answer,
                sources=sources_payload,
            )

        return active_session_id, event_stream()

