from typing import Dict, Any, Optional, List
from uuid import UUID

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
except ImportError:
    # Dummy class if langchain is not installed
    class BaseCallbackHandler:
        pass

from ..client import Anchor8Client

class Anchor8CallbackHandler(BaseCallbackHandler):
    """
    LangChain Callback Handler that enforces Anchor8 Guard checks
    before tool execution and logs activities.
    """

    def __init__(self, agent_id: str, client: Optional[Anchor8Client] = None, abort_on_block: bool = True):
        self.agent_id = agent_id
        self.client = client or Anchor8Client() # Will raise if no env var
        self.abort_on_block = abort_on_block

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when tool starts running."""
        tool_name = serialized.get("name")

        # Parse inputs - LangChain inputs can be string or dict
        # We try to make it a dict for the API
        try:
            # Often input_str is a JSON string or just a string
            import json
            arguments = json.loads(input_str)
            if not isinstance(arguments, dict):
                arguments = {"input": arguments}
        except:
            arguments = {"input": input_str}

        allowed, reason, modified = self.client.check(
            agent_id=self.agent_id,
            tool_name=tool_name,
            arguments=arguments
        )

        if not allowed:
            msg = f"Anchor8 Security Blocked Tool '{tool_name}': {reason}"
            if self.abort_on_block:
                raise PermissionError(msg)
            else:
                print(msg)

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Run when tool ends running."""
        # We could log success here
        pass

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> Any:
        """Run when LLM starts running."""
        # Could log LLM prompts to Observer
        pass
