"""LangChain callback handler for scandar-guard."""
from __future__ import annotations

from typing import Any, Optional, Union
from uuid import UUID

from ..audit import AuditLogger
from ..config import GuardConfig
from ..inspector import MessageInspector
from ..session import SessionManager
from ..tool_inspector import ToolCallInspector


class ScandarCallback:
    """
    LangChain callback handler for Scandar Guard.

    Attach to any LangChain chain or agent to enable runtime security monitoring.

    Usage::

        from scandar_guard.integrations.langchain import ScandarCallback

        chain = MyChain(
            llm=ChatAnthropic(),
            callbacks=[ScandarCallback()]
        )
    """

    def __init__(self, config: Optional[GuardConfig] = None) -> None:
        self._config = config or GuardConfig()
        self._session = SessionManager(session_id=self._config.session_id)
        self._msg_inspector = MessageInspector()
        self._tool_inspector = ToolCallInspector()
        self._logger = AuditLogger(log_file=self._config.log_file)

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        **kwargs: Any,
    ) -> None:
        for prompt in prompts:
            result = self._msg_inspector.inspect_message("user", prompt, self._session.session_id)
            if result.findings:
                self._logger.log_event("llm_input_inspected", self._session.session_id, result.findings, {"prompt_count": len(prompts)})

    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
        if hasattr(response, "generations"):
            for gen_list in response.generations:
                for gen in gen_list:
                    text = getattr(gen, "text", "") or ""
                    if text:
                        result = self._msg_inspector.inspect_response(text, self._session.session_id)
                        if result.findings:
                            self._logger.log_event("llm_output_inspected", self._session.session_id, result.findings, {})

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        tool_name = serialized.get("name", "unknown")
        result = self._tool_inspector.inspect_tool_call(tool_name, {"input": input_str}, self._session.session_id)
        if result.findings:
            self._logger.log_event("tool_call_inspected", self._session.session_id, result.findings, {"tool_name": tool_name})

    def on_tool_end(self, output: str, **kwargs: Any) -> None:
        tool_name = kwargs.get("name", "unknown")
        result = self._tool_inspector.inspect_tool_result(tool_name, output, self._session.session_id)
        if result.findings:
            self._logger.log_event("tool_result_inspected", self._session.session_id, result.findings, {"tool_name": tool_name})

    def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
        pass  # Tool errors are not a security concern for Guard
