"""
Trust-gated chain wrapper for LangChain.

Wraps any chain to verify agents before allowing execution.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from langchain_core.runnables import Runnable, RunnableConfig

from agent_trust import AgentTrustClient, Verdict


class UntrustedAgentError(Exception):
    """Raised when an agent fails trust verification."""
    
    def __init__(
        self,
        message: str,
        agent_name: str,
        agent_url: str,
        verdict: Verdict,
        trust_score: Optional[float] = None,
    ):
        super().__init__(message)
        self.agent_name = agent_name
        self.agent_url = agent_url
        self.verdict = verdict
        self.trust_score = trust_score


class TrustGatedChain(Runnable):
    """
    A chain wrapper that verifies agent trust before execution.
    
    Use this to wrap chains that interact with external agents,
    ensuring only trusted agents can participate.
    
    Example:
        from langchain_openai import ChatOpenAI
        from agent_trust_langchain import TrustGatedChain
        
        llm = ChatOpenAI()
        
        # Wrap the chain with trust verification
        gated = TrustGatedChain(
            chain=llm,
            agent_name="External Service",
            agent_url="https://external.ai/agent",
            min_trust_score=60,
            block_verdict=Verdict.BLOCK
        )
        
        try:
            result = gated.invoke("Hello")
        except UntrustedAgentError as e:
            print(f"Agent not trusted: {e}")
    """
    
    def __init__(
        self,
        chain: Runnable,
        agent_name: str,
        agent_url: str,
        agent_description: Optional[str] = None,
        min_trust_score: float = 50.0,
        block_on_block_verdict: bool = True,
        block_on_caution_verdict: bool = False,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
        cache_verification: bool = True,
    ):
        """
        Initialize the trust-gated chain.
        
        Args:
            chain: The underlying chain to wrap
            agent_name: Name of the agent being verified
            agent_url: URL/identifier of the agent
            agent_description: Description of the agent
            min_trust_score: Minimum trust score required (0-100)
            block_on_block_verdict: Block if verdict is BLOCK
            block_on_caution_verdict: Block if verdict is CAUTION
            api_url: Custom API URL
            api_key: API key for authentication
            cache_verification: Cache verification result for chain lifetime
        """
        self.chain = chain
        self.agent_name = agent_name
        self.agent_url = agent_url
        self.agent_description = agent_description
        self.min_trust_score = min_trust_score
        self.block_on_block_verdict = block_on_block_verdict
        self.block_on_caution_verdict = block_on_caution_verdict
        self.cache_verification = cache_verification
        
        # Initialize client
        kwargs = {}
        if api_url:
            kwargs["api_url"] = api_url
        if api_key:
            kwargs["api_key"] = api_key
        self.client = AgentTrustClient(**kwargs)
        
        # Cached verification result
        self._cached_result = None
        self._is_verified = False
    
    def _verify_agent(self) -> None:
        """Verify the agent and raise if untrusted."""
        # Use cache if available
        if self.cache_verification and self._cached_result is not None:
            result = self._cached_result
        else:
            result = self.client.verify_agent(
                name=self.agent_name,
                url=self.agent_url,
                description=self.agent_description,
            )
            if self.cache_verification:
                self._cached_result = result
        
        # Check verdict
        if self.block_on_block_verdict and result.verdict == Verdict.BLOCK:
            raise UntrustedAgentError(
                f"Agent '{self.agent_name}' is blocked: {result.reasoning}",
                agent_name=self.agent_name,
                agent_url=self.agent_url,
                verdict=result.verdict,
                trust_score=result.trust_score,
            )
        
        if self.block_on_caution_verdict and result.verdict == Verdict.CAUTION:
            raise UntrustedAgentError(
                f"Agent '{self.agent_name}' flagged with caution: {result.reasoning}",
                agent_name=self.agent_name,
                agent_url=self.agent_url,
                verdict=result.verdict,
                trust_score=result.trust_score,
            )
        
        # Check trust score
        if result.trust_score is not None and result.trust_score < self.min_trust_score:
            raise UntrustedAgentError(
                f"Agent '{self.agent_name}' trust score ({result.trust_score}) "
                f"below minimum ({self.min_trust_score})",
                agent_name=self.agent_name,
                agent_url=self.agent_url,
                verdict=result.verdict,
                trust_score=result.trust_score,
            )
        
        self._is_verified = True
    
    def invoke(
        self,
        input: Any,
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> Any:
        """Verify agent trust, then invoke the underlying chain."""
        self._verify_agent()
        return self.chain.invoke(input, config, **kwargs)
    
    async def ainvoke(
        self,
        input: Any,
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> Any:
        """Async version of invoke."""
        self._verify_agent()  # Still sync for now
        return await self.chain.ainvoke(input, config, **kwargs)
    
    def batch(
        self,
        inputs: List[Any],
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> List[Any]:
        """Verify once, then batch process."""
        self._verify_agent()
        return self.chain.batch(inputs, config, **kwargs)
    
    @property
    def InputType(self):
        return self.chain.InputType
    
    @property
    def OutputType(self):
        return self.chain.OutputType
    
    def get_verification_result(self):
        """Get the cached verification result, if any."""
        return self._cached_result
    
    def clear_cache(self):
        """Clear the cached verification result."""
        self._cached_result = None
        self._is_verified = False
