"""
QppgHaystackMonitor — Haystack pipeline wrapper for llm-guard-kit.

Wraps a Haystack Pipeline and intercepts component outputs to build QPPG
chain dicts and score each run automatically.

Requirements:
    pip install "llm-guard-kit[haystack]"
    # installs haystack-ai>=2.0

Usage:
    from haystack import Pipeline
    from qppg_service.integrations.haystack_callback import QppgHaystackMonitor

    pipeline = Pipeline()
    # ... add components ...

    qppg = QppgHaystackMonitor(pipeline)
    outputs, result = qppg.run({"query": {"query": "What year was X founded?"}})

    if result and result.needs_review:
        print(f"HIGH RISK: {result.risk_score:.2f}")

    # With a monitor:
    from qppg_service import QppgMonitor
    monitor = QppgMonitor(db_path="~/.qppg/chains.db")
    qppg    = QppgHaystackMonitor(pipeline, monitor=monitor)
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

from ..label_free_scorer import LabelFreeScorer, LabelFreeResult
from ..monitor import QppgMonitor


def _check_haystack():
    try:
        import haystack  # noqa: F401
    except ImportError:
        raise ImportError(
            "Haystack integration requires haystack-ai. "
            "Install with: pip install 'llm-guard-kit[haystack]'"
        )


class QppgHaystackMonitor:
    """
    Wraps a Haystack Pipeline to capture intermediate outputs and score each run.

    Detects the following standard Haystack component types:
    - Retrievers (InMemoryBM25Retriever, BM25Retriever, etc.) → Search steps
    - Generators (OpenAIGenerator, etc.) → final answer
    - PromptBuilders → thought extraction

    Parameters
    ----------
    pipeline : Haystack Pipeline instance to wrap
    monitor  : optional QppgMonitor for persistence
    threshold: alert threshold (when creating default scorer)
    """

    def __init__(
        self,
        pipeline: Any,
        monitor: Optional[QppgMonitor] = None,
        threshold: float = 0.65,
    ):
        _check_haystack()
        self._pipeline  = pipeline
        self._monitor   = monitor
        self._scorer    = monitor.scorer if monitor else LabelFreeScorer(
            review_threshold=threshold)
        self._last_result: Optional[LabelFreeResult] = None

    def run(
        self,
        inputs: Dict[str, Any],
        **kwargs: Any,
    ) -> Tuple[Dict[str, Any], Optional[LabelFreeResult]]:
        """
        Run the pipeline and score the resulting chain.

        Parameters
        ----------
        inputs  : pipeline inputs dict (same format as pipeline.run())
        **kwargs: passed through to pipeline.run()

        Returns
        -------
        (pipeline_outputs, LabelFreeResult | None)
        """
        # Extract question from inputs
        question = self._extract_question(inputs)

        # Run the pipeline
        outputs = self._pipeline.run(inputs, **kwargs)

        # Build chain steps from outputs
        steps, final_answer = self._extract_steps(inputs, outputs)

        if not steps and not final_answer:
            self._last_result = None
            return outputs, None

        finished = bool(final_answer)

        if self._monitor is not None:
            self._monitor.track(question, steps, final_answer, finished=finished)
        self._last_result = self._scorer.score(question, steps, final_answer, finished)
        return outputs, self._last_result

    def get_last_result(self) -> Optional[LabelFreeResult]:
        """Return the most recent LabelFreeResult."""
        return self._last_result

    # ── Extraction helpers ─────────────────────────────────────────────────────

    def _extract_question(self, inputs: Dict[str, Any]) -> str:
        """Find the question string in pipeline inputs."""
        for component_inputs in inputs.values():
            if isinstance(component_inputs, dict):
                for key in ("query", "question", "input", "prompt"):
                    if key in component_inputs:
                        val = component_inputs[key]
                        if isinstance(val, str):
                            return val
        return ""

    def _extract_steps(
        self,
        inputs: Dict[str, Any],
        outputs: Dict[str, Any],
    ) -> Tuple[List[dict], str]:
        """
        Map Haystack component outputs to QPPG step dicts.

        Heuristics:
        - Component name contains "retriever" or "retriev" → Search step
        - Component name contains "generator" or "llm"    → Finish step / thought
        - Component name contains "prompt"                → thought
        """
        steps: List[dict] = []
        final_answer = ""
        query = self._extract_question(inputs)

        for comp_name, comp_output in outputs.items():
            if not isinstance(comp_output, dict):
                continue

            lower = comp_name.lower()

            if any(k in lower for k in ("retriever", "retriev", "bm25", "embedding")):
                # Extract documents as observation
                docs = (comp_output.get("documents")
                        or comp_output.get("results")
                        or [])
                obs_parts = []
                for doc in docs[:3]:
                    content = ""
                    if hasattr(doc, "content"):
                        content = str(doc.content or "")[:300]
                    elif isinstance(doc, dict):
                        content = str(doc.get("content", ""))[:300]
                    if content:
                        obs_parts.append(content)

                steps.append({
                    "thought":     "",
                    "action_type": "Search",
                    "action_arg":  query,
                    "observation": " | ".join(obs_parts) if obs_parts else "[retrieved]",
                })

            elif any(k in lower for k in ("generator", "llm", "openai", "gemini",
                                           "anthropic", "cohere")):
                # Extract generated text as final answer
                replies = (comp_output.get("replies")
                           or comp_output.get("answers")
                           or comp_output.get("results")
                           or [])
                if replies:
                    first = replies[0]
                    if isinstance(first, str):
                        final_answer = first
                    elif hasattr(first, "data"):
                        final_answer = str(first.data or "")
                    elif isinstance(first, dict):
                        final_answer = str(first.get("answer", first.get("content", "")))

                if final_answer:
                    steps.append({
                        "thought":     "",
                        "action_type": "Finish",
                        "action_arg":  final_answer,
                        "observation": "",
                    })

            elif "prompt" in lower:
                # PromptBuilder output → treat rendered prompt as thought for next step
                prompt_val = comp_output.get("prompt", "")
                if prompt_val and steps:
                    steps[-1]["thought"] = str(prompt_val)[:300]

        return steps, final_answer
