"""Core Doctor orchestrator — runs the agentic diagnosis loop."""
from __future__ import annotations
from typing import List, Optional
from .config import RagDoctorConfig
from .connectors.base import Document, PipelineConnector
from .tools import (
    ChunkAnalyzer, RetrievalAuditor, PositionTester,
    HallucinationTracer, ChunkOptimizer, QueryRewriter,
    ToolResult,
)
from .report import DiagnosisReport, SEVERITY_RANK


class Doctor:
    """
    Orchestrates the six diagnostic tools in a rule-based agent loop.

    Decision tree (in priority order):
      1. RetrievalAuditor   → RC-1 (retrieval_miss) or RC-5 (query_mismatch)
      2. PositionTester     → RC-2 (context_position_bias)
      3. ChunkAnalyzer      → RC-3 (chunk_fragmentation)
      4. HallucinationTracer→ RC-4 (hallucination)
      5. QueryRewriter      → sub-tool when RC-1 detected
      6. ChunkOptimizer     → sub-tool when RC-3 detected
    """

    def __init__(self, config: RagDoctorConfig, connector: Optional[PipelineConnector] = None):
        self.config    = config
        self.connector = connector
        cfg = config.diagnosis

        self._chunk_analyzer       = ChunkAnalyzer(coherence_threshold=cfg.coherence_threshold)
        self._retrieval_auditor    = RetrievalAuditor(recall_threshold=cfg.recall_threshold)
        self._position_tester      = PositionTester()
        self._hallucination_tracer = HallucinationTracer(faithfulness_threshold=cfg.faithfulness_threshold)
        self._chunk_optimizer      = ChunkOptimizer()
        self._query_rewriter       = QueryRewriter()

    @classmethod
    def from_config(cls, path: str, connector: Optional[PipelineConnector] = None) -> "Doctor":
        return cls(RagDoctorConfig.from_yaml(path), connector)

    @classmethod
    def default(cls, connector: Optional[PipelineConnector] = None) -> "Doctor":
        return cls(RagDoctorConfig.default(), connector)

    def diagnose(
        self,
        query: str,
        answer: str,
        docs: Optional[List[Document]] = None,
        expected: Optional[str] = None,
        corpus_texts: Optional[List[str]] = None,
    ) -> DiagnosisReport:
        """
        Run the full diagnosis loop.

        Args:
            query:        User's original query.
            answer:       Answer generated by the RAG pipeline.
            docs:         Retrieved documents (optional — fetched via connector if absent).
            expected:     Ground-truth answer (improves diagnosis accuracy).
            corpus_texts: Raw corpus for chunk optimization (optional).
        """
        tool_results: List[ToolResult] = []
        cfg = self.config.diagnosis

        # Fetch docs via connector if not provided
        if docs is None and self.connector:
            docs = self.connector.retrieve(query, top_k=self.config.retrieval.top_k)
        docs = docs or []

        # ── Step 1: Retrieval Audit ──────────────────────────────────────────
        retrieval_result = self._retrieval_auditor.run(
            docs=docs, expected=expected, query=query
        )
        tool_results.append(retrieval_result)
        retrieval_miss = not retrieval_result.passed

        # ── Step 2: Position Test (only when retrieval succeeded) ────────────
        position_result = None
        if not retrieval_miss and len(docs) >= 3 and cfg.position_tester:
            position_result = self._position_tester.run(
                docs=docs, query=query, expected=expected
            )
            tool_results.append(position_result)

        # ── Step 3: Chunk Analysis ───────────────────────────────────────────
        chunk_result = None
        if cfg.chunk_analyzer:
            chunk_result = self._chunk_analyzer.run(docs=docs, query=query)
            tool_results.append(chunk_result)

        # ── Step 4: Hallucination Tracing ────────────────────────────────────
        hallucination_result = None
        if cfg.hallucination_tracer:
            hallucination_result = self._hallucination_tracer.run(
                answer=answer, docs=docs
            )
            tool_results.append(hallucination_result)

        # ── Step 5: Query Rewriter (when retrieval miss) ─────────────────────
        rewrite_result = None
        if retrieval_miss and cfg.query_rewriter:
            rewrite_result = self._query_rewriter.run(
                query=query,
                original_docs=docs,
                connector=self.connector,
                expected=expected,
                top_k=self.config.retrieval.top_k,
            )
            tool_results.append(rewrite_result)

        # ── Step 6: Chunk Optimizer (when fragmentation detected) ────────────
        optimizer_result = None
        if (
            chunk_result and not chunk_result.passed
            and corpus_texts and cfg.chunk_optimizer
        ):
            test_pairs = (
                [{"query": query, "expected": expected}]
                if expected else []
            )
            if test_pairs:
                optimizer_result = self._chunk_optimizer.run(
                    corpus_texts=corpus_texts,
                    test_pairs=test_pairs,
                    current_strategy={
                        "strategy":      self.config.chunking.strategy,
                        "chunk_size":    self.config.chunking.chunk_size,
                        "chunk_overlap": self.config.chunking.chunk_overlap,
                    },
                )
                tool_results.append(optimizer_result)

        # ── Root Cause Attribution ───────────────────────────────────────────
        root_cause, severity, finding, fix, patch = self._root_cause(
            retrieval_result=retrieval_result,
            position_result=position_result,
            chunk_result=chunk_result,
            hallucination_result=hallucination_result,
            rewrite_result=rewrite_result,
            optimizer_result=optimizer_result,
        )

        return DiagnosisReport(
            query=query,
            answer=answer,
            expected=expected,
            root_cause=root_cause,
            severity=severity,
            finding=finding,
            tool_results=tool_results,
            fix_suggestion=fix,
            config_patch=patch,
            faithfulness_score=float(
                hallucination_result.details.get("faithfulness_score", 0.0)
                if hallucination_result else 0.0
            ),
            retrieval_score=float(
                retrieval_result.details.get("best_score", 0.0)
                if retrieval_result else 0.0
            ),
        )

    def _root_cause(
        self,
        retrieval_result, position_result, chunk_result,
        hallucination_result, rewrite_result, optimizer_result,
    ):
        fix, patch = None, {}

        # RC-1 / RC-5: retrieval miss
        if retrieval_result and not retrieval_result.passed:
            if rewrite_result and rewrite_result.details.get("improvement", 0) > 0.03:
                return (
                    "query_mismatch", "high",
                    f"Query vocabulary mismatch. "
                    f"'{rewrite_result.details['best_strategy']}' rewrite improved recall by "
                    f"+{rewrite_result.details['improvement']:.3f}.",
                    rewrite_result.recommendation, {},
                )
            return (
                "retrieval_miss", "high",
                retrieval_result.finding,
                retrieval_result.recommendation, {},
            )

        # RC-2: position bias
        if position_result and not position_result.passed:
            return (
                "context_position_bias",
                position_result.severity,
                position_result.finding,
                position_result.recommendation,
                {"retrieval.reranker": True},
            )

        # RC-3: chunk fragmentation
        if chunk_result and not chunk_result.passed and chunk_result.severity != "low":
            patch = {}
            if optimizer_result and not optimizer_result.passed:
                ranked = optimizer_result.details.get("ranked_strategies", [])
                if ranked:
                    b = ranked[0]
                    patch = {
                        "chunking.strategy":     b["strategy"],
                        "chunking.chunk_size":   b["chunk_size"],
                        "chunking.chunk_overlap": b["chunk_overlap"],
                    }
            return (
                "chunk_fragmentation",
                chunk_result.severity,
                chunk_result.finding,
                optimizer_result.recommendation if optimizer_result else chunk_result.recommendation,
                patch,
            )

        # RC-4: hallucination
        if hallucination_result and not hallucination_result.passed:
            return (
                "hallucination",
                hallucination_result.severity,
                hallucination_result.finding,
                hallucination_result.recommendation, {},
            )

        # RC-0: healthy
        worst = "low"
        for tr in [retrieval_result, position_result, chunk_result, hallucination_result]:
            if tr and SEVERITY_RANK.get(tr.severity, 0) > SEVERITY_RANK.get(worst, 0):
                worst = tr.severity
        return "healthy", worst, "No significant issues detected. Pipeline appears healthy.", None, {}

    def batch_diagnose(self, test_cases: List[dict]) -> List[DiagnosisReport]:
        return [
            self.diagnose(
                query=c.get("query", ""),
                answer=c.get("answer", ""),
                docs=c.get("docs"),
                expected=c.get("expected"),
                corpus_texts=c.get("corpus_texts"),
            )
            for c in test_cases
        ]
