"""
StrivioCallbackHandler — LangChain callback that sends every agent action
to the Strivio AI Firewall automatically.

Usage:
    from langchain.agents import initialize_agent, AgentType
    from langchain_openai import ChatOpenAI
    from langchain.tools import DuckDuckGoSearchRun
    from strivio import StrivioCallbackHandler

    strivio = StrivioCallbackHandler(
        api_url="http://localhost:4000/api/v1",
        token="your-jwt-token",
        agent_id="my-research-agent",
        block_on_high_risk=True,  # raises StrivioBlockedError if firewall blocks
    )

    llm   = ChatOpenAI(model="gpt-4o", callbacks=[strivio])
    tools = [DuckDuckGoSearchRun()]
    agent = initialize_agent(tools, llm, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                             callbacks=[strivio])

    agent.run("What is the latest news on AI security?")
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Union

from .client import StrivioClient

logger = logging.getLogger("strivio.langchain")


class StrivioBlockedError(Exception):
    """Raised when Strivio firewall blocks an action."""
    def __init__(self, action: str, target: str, risk_score: int):
        self.action     = action
        self.target     = target
        self.risk_score = risk_score
        super().__init__(
            f"[Strivio] Action BLOCKED by firewall: action={action} target={target} riskScore={risk_score}"
        )


class StrivioCallbackHandler:
    """
    LangChain-compatible callback handler for Strivio AI Firewall.

    Drop-in integration: pass as callbacks=[strivio] to any LangChain
    LLM, chain, agent, or tool.

    Intercepts:
      - on_llm_start        → logs LLM call event
      - on_tool_start       → logs tool invocation + checks firewall
      - on_tool_end         → logs tool result
      - on_agent_action     → logs agent decision
      - on_chain_start      → logs chain entry
    """

    def __init__(
        self,
        api_url:            Optional[str] = None,
        token:              Optional[str] = None,
        agent_id:           Optional[str] = None,
        block_on_high_risk: bool          = False,
        risk_threshold:     int           = 80,
        async_mode:         bool          = True,
    ):
        self.client = StrivioClient(
            api_url=api_url,
            token=token,
            agent_id=agent_id or "langchain-agent",
            async_mode=False,  # sync for block_on_high_risk to work
        )
        self.block_on_high_risk = block_on_high_risk
        self.risk_threshold     = risk_threshold

    # ── LangChain callback interface ──────────────────────────────────────

    def on_llm_start(self, serialized: Dict, prompts: List[str], **kwargs) -> None:
        model = serialized.get("name") or serialized.get("id", ["unknown"])[-1]
        prompt = prompts[0] if prompts else ""
        self.client.llm_call(prompt=prompt, model=str(model))

    def on_llm_end(self, response: Any, **kwargs) -> None:
        pass  # Could log token usage here

    def on_llm_error(self, error: Exception, **kwargs) -> None:
        self.client.send_event(
            action="llm_error",
            target="llm",
            metadata={"error": str(error)[:300]},
        )

    def on_tool_start(self, serialized: Dict, input_str: str, **kwargs) -> None:
        tool_name = serialized.get("name") or "unknown_tool"
        result    = self.client.send_event(
            action="tool_call",
            target=tool_name,
            metadata={"tool": tool_name, "input": input_str[:500]},
        )
        if self.block_on_high_risk and result:
            score = result.get("riskScore", 0)
            if result.get("blocked") or score >= self.risk_threshold:
                raise StrivioBlockedError(
                    action="tool_call",
                    target=tool_name,
                    risk_score=score,
                )

    def on_tool_end(self, output: str, **kwargs) -> None:
        # Log that the tool completed — useful for audit trail
        self.client.send_event(
            action="tool_result",
            target="tool_output",
            metadata={"output_preview": str(output)[:300]},
        )

    def on_tool_error(self, error: Exception, **kwargs) -> None:
        self.client.send_event(
            action="tool_error",
            target="tool",
            metadata={"error": str(error)[:300]},
        )

    def on_agent_action(self, action: Any, **kwargs) -> None:
        tool  = getattr(action, "tool", "unknown")
        input_= getattr(action, "tool_input", "")
        self.client.send_event(
            action="agent_decision",
            target=tool,
            metadata={"tool": tool, "input": str(input_)[:500]},
        )

    def on_agent_finish(self, finish: Any, **kwargs) -> None:
        output = getattr(finish, "return_values", {})
        self.client.send_event(
            action="agent_finish",
            target="final_output",
            metadata={"output": str(output)[:300]},
        )

    def on_chain_start(self, serialized: Dict, inputs: Dict, **kwargs) -> None:
        chain_name = serialized.get("name") or serialized.get("id", ["chain"])[-1]
        self.client.send_event(
            action="chain_start",
            target=str(chain_name),
            metadata={"inputs_preview": str(inputs)[:300]},
        )

    def on_chain_end(self, outputs: Dict, **kwargs) -> None:
        pass

    def on_chain_error(self, error: Exception, **kwargs) -> None:
        self.client.send_event(
            action="chain_error",
            target="chain",
            metadata={"error": str(error)[:300]},
        )

    # Support new LangChain BaseCallbackHandler interface
    def __call__(self, *args, **kwargs):
        return self
