"""
QppgLlamaIndexCallback — LlamaIndex agent monitor for llm-guard-kit.

Hooks into LlamaIndex's callback system to capture retrieval steps and
score every query run automatically.

Requirements:
    pip install "llm-guard-kit[llamaindex]"
    # installs llama-index-core>=0.10

Usage:
    from llama_index.core import Settings
    from llama_index.core.callbacks import CallbackManager
    from qppg_service.integrations.llamaindex_callback import QppgLlamaIndexCallback

    qppg_cb = QppgLlamaIndexCallback()
    Settings.callback_manager = CallbackManager([qppg_cb])

    # After querying:
    response = query_engine.query("What year was X founded?")
    result   = qppg_cb.get_last_result()
    if result and result.needs_review:
        print(f"HIGH RISK: {result.risk_score:.2f}")
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

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


def _lazy_imports():
    try:
        from llama_index.core.callbacks.base_handler import BaseCallbackHandler
        from llama_index.core.callbacks.schema import CBEventType, EventPayload
        return BaseCallbackHandler, CBEventType, EventPayload
    except ImportError:
        raise ImportError(
            "LlamaIndex integration requires llama-index-core. "
            "Install with: pip install 'llm-guard-kit[llamaindex]'"
        )


class QppgLlamaIndexCallback:
    """
    LlamaIndex callback handler that monitors query quality.

    Captures RETRIEVE events as Search steps and QUERY end events as
    the chain completion trigger.

    Parameters
    ----------
    monitor   : existing QppgMonitor; if None, a new LabelFreeScorer is used
    threshold : alert threshold (used when creating a default scorer)
    """

    def __new__(cls, *args, **kwargs):
        base, CBEventType, EventPayload = _lazy_imports()
        if base not in cls.__bases__:
            cls = type(cls.__name__, (base,), dict(cls.__dict__))
        return object.__new__(cls)

    def __init__(
        self,
        monitor: Optional[QppgMonitor] = None,
        threshold: float = 0.65,
    ):
        _, CBEventType, _ = _lazy_imports()
        event_starts_to_ignore: List[Any] = []
        event_ends_to_ignore: List[Any] = []
        super().__init__(
            event_starts_to_ignore=event_starts_to_ignore,
            event_ends_to_ignore=event_ends_to_ignore,
        )
        self._monitor   = monitor
        self._scorer    = monitor.scorer if monitor else LabelFreeScorer(
            review_threshold=threshold)
        self._CBEventType = CBEventType

        self._current_question: str = ""
        self._current_steps: List[dict] = []
        self._last_result: Optional[LabelFreeResult] = None
        self._pending_retrieve_query: str = ""

    # ── LlamaIndex event hooks ─────────────────────────────────────────────────

    def on_event_start(
        self,
        event_type: Any,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> str:
        CBEventType = self._CBEventType
        payload = payload or {}

        if event_type == CBEventType.QUERY:
            self._current_steps = []
            self._last_result   = None
            query_str = payload.get("query_str", "")
            if query_str:
                self._current_question = str(query_str)

        elif event_type == CBEventType.RETRIEVE:
            # Store query string for the upcoming retrieve step
            self._pending_retrieve_query = str(payload.get("query_str", ""))

        return event_id

    def on_event_end(
        self,
        event_type: Any,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> None:
        CBEventType = self._CBEventType
        payload = payload or {}

        if event_type == CBEventType.RETRIEVE:
            # Build a Search step from retrieved nodes
            nodes = payload.get("nodes", [])
            obs_parts = []
            for node in nodes[:3]:  # cap at 3 nodes for observation
                text = ""
                if hasattr(node, "node") and hasattr(node.node, "text"):
                    text = node.node.text[:300]
                elif hasattr(node, "text"):
                    text = node.text[:300]
                if text:
                    obs_parts.append(text)

            self._current_steps.append({
                "thought":     "",
                "action_type": "Search",
                "action_arg":  self._pending_retrieve_query,
                "observation": " | ".join(obs_parts) if obs_parts else "[retrieved]",
            })
            self._pending_retrieve_query = ""

        elif event_type == CBEventType.LLM:
            # Capture LLM response as thought for the last step
            response = payload.get("response", None)
            if response and self._current_steps:
                text = ""
                if hasattr(response, "message"):
                    text = str(getattr(response.message, "content", ""))
                elif hasattr(response, "text"):
                    text = str(response.text)
                else:
                    text = str(response)
                if text and not self._current_steps[-1]["thought"]:
                    self._current_steps[-1]["thought"] = text[:300]

        elif event_type == CBEventType.QUERY:
            # Query completed — score the chain
            response = payload.get("response", None)
            final_answer = ""
            if response is not None:
                if hasattr(response, "response"):
                    final_answer = str(response.response or "")
                elif hasattr(response, "answer"):
                    final_answer = str(response.answer or "")
                else:
                    final_answer = str(response)

            # Add Finish step
            self._current_steps.append({
                "thought":     "",
                "action_type": "Finish",
                "action_arg":  final_answer,
                "observation": "",
            })

            q     = self._current_question or "unknown"
            steps = self._current_steps
            if self._monitor is not None:
                self._monitor.track(q, steps, final_answer, finished=True)
            self._last_result = self._scorer.score(q, steps, final_answer, True)

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        pass

    def end_trace(
        self,
        trace_id: Optional[str] = None,
        trace_map: Optional[Dict[str, List[str]]] = None,
    ) -> None:
        pass

    # ── Accessor ───────────────────────────────────────────────────────────────

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