"""
QppgLangChainCallback — drop-in LangChain agent monitor.

Wraps BaseCallbackHandler to capture ReAct steps automatically and score
every agent run with LabelFreeScorer / QppgMonitor.

Requirements:
    pip install "llm-guard-kit[langchain]"
    # installs langchain-core>=0.1

Usage:
    from langchain.agents import AgentExecutor
    from qppg_service.integrations.langchain_callback import QppgLangChainCallback

    callback = QppgLangChainCallback(threshold=0.65)
    agent_executor.invoke({"input": question}, config={"callbacks": [callback]})

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

    # Or with a pre-built monitor for persistence:
    from qppg_service import QppgMonitor
    monitor  = QppgMonitor(db_path="~/.qppg/chains.db")
    callback = QppgLangChainCallback(monitor=monitor)
"""

from __future__ import annotations

import json
import re
from typing import Any, Optional, Union

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


def _lazy_base() -> type:
    try:
        from langchain_core.callbacks.base import BaseCallbackHandler
        return BaseCallbackHandler
    except ImportError:
        raise ImportError(
            "LangChain integration requires langchain-core. "
            "Install with: pip install 'llm-guard-kit[langchain]'"
        )


def _normalize_tool_name(tool_name: str) -> str:
    """Map LangChain tool names to QPPG action_type strings."""
    lower = tool_name.lower()
    if any(k in lower for k in ("search", "retriev", "lookup", "wikipedia",
                                 "tavily", "duck", "serp", "bing", "google")):
        return "Search"
    if any(k in lower for k in ("finish", "final", "answer")):
        return "Finish"
    return tool_name  # preserve custom tool names as-is


def _extract_thought(action_log: str) -> str:
    """Extract the Thought from a LangChain action log string."""
    m = re.search(r"Thought[:\s]+(.+?)(?=Action[:\s]|$)", action_log, re.DOTALL | re.IGNORECASE)
    if m:
        return m.group(1).strip()
    # Fallback: return whatever is before the first "Action:" line
    parts = re.split(r"\nAction[:\s]", action_log, maxsplit=1)
    return parts[0].strip() if parts else action_log.strip()


class QppgLangChainCallback:
    """
    LangChain BaseCallbackHandler that records agent steps and scores runs.

    Parameters
    ----------
    monitor    : existing QppgMonitor to use; if None a new LabelFreeScorer is created
    scorer     : existing LabelFreeScorer; used if monitor is None
    threshold  : risk threshold for alerts (used when creating a default scorer)
    question   : override the question string (auto-detected from input otherwise)
    """

    def __new__(cls, *args, **kwargs):
        # Dynamically inherit from BaseCallbackHandler at instantiation time
        base = _lazy_base()
        if base not in cls.__bases__:
            cls = type(cls.__name__, (base,), dict(cls.__dict__))
        instance = object.__new__(cls)
        return instance

    def __init__(
        self,
        monitor: Optional[QppgMonitor] = None,
        scorer: Optional[LabelFreeScorer] = None,
        threshold: float = 0.65,
        question: str = "",
    ):
        super().__init__()
        self._monitor   = monitor
        self._scorer    = scorer or (None if monitor else LabelFreeScorer(
            review_threshold=threshold))
        self.threshold  = threshold
        self._question  = question

        self._current_steps: list[dict] = []
        self._last_result: Optional[LabelFreeResult] = None
        self._last_alert:  Optional[QppgAlert]       = None

    # ── LangChain callback hooks ───────────────────────────────────────────────

    def on_chain_start(self, serialized: dict, inputs: dict, **kwargs: Any) -> None:
        """Detect the question from chain inputs."""
        self._current_steps = []
        self._last_result   = None
        self._last_alert    = None
        if not self._question:
            # Common input keys used by different agent types
            for key in ("input", "question", "query", "human_input"):
                if key in inputs:
                    self._question = str(inputs[key])
                    break

    def on_agent_action(self, action: Any, **kwargs: Any) -> None:
        """Capture thought and action from AgentAction."""
        thought    = _extract_thought(getattr(action, "log", ""))
        tool_name  = getattr(action, "tool", "")
        tool_input = getattr(action, "tool_input", "")

        # tool_input may be a dict (structured args) or a string
        if isinstance(tool_input, dict):
            action_arg = json.dumps(tool_input)
        else:
            action_arg = str(tool_input)

        step = {
            "thought":     thought,
            "action_type": _normalize_tool_name(tool_name),
            "action_arg":  action_arg,
            "observation": "",  # filled in by on_tool_end
        }
        self._current_steps.append(step)

    def on_tool_end(self, output: str, **kwargs: Any) -> None:
        """Fill the observation of the most recent step."""
        if self._current_steps:
            self._current_steps[-1]["observation"] = str(output)[:1000]

    def on_agent_finish(self, finish: Any, **kwargs: Any) -> None:
        """Score the completed chain."""
        final_answer = ""
        rv = getattr(finish, "return_values", {})
        if isinstance(rv, dict):
            final_answer = str(rv.get("output", rv.get("answer", "")))
        else:
            final_answer = str(rv)

        # Add a Finish step
        log_text = getattr(finish, "log", "")
        thought  = _extract_thought(log_text) if log_text else ""
        self._current_steps.append({
            "thought":     thought,
            "action_type": "Finish",
            "action_arg":  final_answer,
            "observation": "",
        })

        q = self._question or "unknown"
        steps = [s for s in self._current_steps if s["action_type"] != "Finish"
                 or s == self._current_steps[-1]]

        if self._monitor is not None:
            alert = self._monitor.track(q, steps, final_answer, finished=True)
            self._last_alert  = alert
            # Retrieve result from last history entry
            if self._monitor._history:
                entry = self._monitor._history[-1]
                # Reconstruct a minimal result proxy for get_last_result()
                self._last_result = self._monitor.scorer.score(
                    q, steps, final_answer, True)
        else:
            scorer = self._scorer or LabelFreeScorer(review_threshold=self.threshold)
            self._last_result = scorer.score(q, steps, final_answer, True)

        # Reset for next run
        self._question = ""

    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
        """Clear state on chain error."""
        self._current_steps = []

    # ── Public accessors ───────────────────────────────────────────────────────

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

    def get_last_alert(self) -> Optional[QppgAlert]:
        """Return the QppgAlert from the most recent run (None if no alert)."""
        return self._last_alert

    def reset(self) -> None:
        """Manually reset state between runs (auto-reset on on_chain_start)."""
        self._current_steps = []
        self._last_result   = None
        self._last_alert    = None
        self._question      = ""
