"""Support classes for different LLM backends (e.g., AWS GovCloud LLMs)"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_llm.backends.ipynb.

# %% auto 0
__all__ = ['ChatGovCloudBedrock']

# %% ../../nbs/00_llm.backends.ipynb 3
"""
Custom LangChain Chat classes for various frameworks and cloud providers.
"""

import json
import os
from typing import Any, Dict, Iterator, List, Optional, Union

from pydantic import Field
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_core.outputs import ChatGeneration, ChatResult


class ChatGovCloudBedrock(BaseChatModel):
    """
    Custom LangChain Chat model for AWS GovCloud Bedrock.
    
    This class provides integration with Amazon Bedrock running in AWS GovCloud regions,
    supporting custom VPC endpoints and GovCloud-specific configurations.
    """
    
    model_id: str = Field(description="The Bedrock model ID or inference profile ARN")
    region_name: str = Field(default="us-gov-east-1", description="AWS GovCloud region")
    endpoint_url: Optional[str] = Field(default=None, description="Custom endpoint URL for VPC endpoints")
    aws_access_key_id: Optional[str] = Field(default=None, description="AWS access key ID")
    aws_secret_access_key: Optional[str] = Field(default=None, description="AWS secret access key")
    max_tokens: int = Field(default=512, description="Maximum tokens to generate")
    temperature: float = Field(default=0.7, description="Sampling temperature")
    streaming: bool = Field(default=False, description="Enable streaming responses")
    client: Any = Field(default=None, exclude=True, description="Boto3 client instance")

    def __init__(
        self,
        model_id: str,
        region_name: str = "us-gov-east-1",
        endpoint_url: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_secret_access_key: Optional[str] = None,
        max_tokens: int = 512,
        temperature: float = 0.7,
        streaming: bool = False,
        callbacks: Optional[List] = None,
        **kwargs
    ):
        """
        Initialize ChatGovCloudBedrock.
        
        Args:
            model_id: The Bedrock model ID or inference profile ARN
            region_name: AWS GovCloud region (us-gov-east-1 or us-gov-west-1)
            endpoint_url: Custom endpoint URL for VPC endpoints
            aws_access_key_id: AWS access key ID (or use environment variables)
            aws_secret_access_key: AWS secret access key (or use environment variables)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            streaming: Enable streaming responses
            callbacks: LangChain callbacks
            **kwargs: Additional parameters
        """
        super().__init__(
            model_id=model_id,
            region_name=region_name,
            endpoint_url=endpoint_url,
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
            max_tokens=max_tokens,
            temperature=temperature,
            streaming=streaming,
            callbacks=callbacks,
            **kwargs
        )

    def model_post_init(self, __context: Any) -> None:
        """Initialize boto3 client after model validation."""
        # Initialize boto3 client
        client_kwargs = {
            "service_name": "bedrock-runtime",
            "region_name": self.region_name,
        }
        
        if self.endpoint_url:
            client_kwargs["endpoint_url"] = self.endpoint_url
            
        if self.aws_access_key_id:
            client_kwargs["aws_access_key_id"] = self.aws_access_key_id
        else:
            client_kwargs["aws_access_key_id"] = os.environ.get('AWS_ACCESS_KEY_ID')
            
        if self.aws_secret_access_key:
            client_kwargs["aws_secret_access_key"] = self.aws_secret_access_key
        else:
            client_kwargs["aws_secret_access_key"] = os.environ.get('AWS_SECRET_ACCESS_KEY')
        
        try:
            import boto3
        except ImportError:
            raise ImportError('Please install boto3: pip install boto3')

        self.client = boto3.client(**client_kwargs)

    def _convert_messages_to_bedrock_format(self, messages: List[BaseMessage]) -> List[Dict[str, str]]:
        """
        Convert LangChain messages to Bedrock API format.
        
        Args:
            messages: List of LangChain BaseMessage objects
            
        Returns:
            List of dictionaries in Bedrock API format
        """
        bedrock_messages = []
        
        for message in messages:
            if isinstance(message, HumanMessage):
                bedrock_messages.append({
                    "role": "user", 
                    "content": message.content
                })
            elif isinstance(message, AIMessage):
                bedrock_messages.append({
                    "role": "assistant", 
                    "content": message.content
                })
            elif isinstance(message, SystemMessage):
                # For Claude models, system messages can be handled as user messages
                # or through the system parameter in the request body
                bedrock_messages.append({
                    "role": "user", 
                    "content": f"System: {message.content}"
                })
            else:
                # Default to user role for unknown message types
                bedrock_messages.append({
                    "role": "user", 
                    "content": str(message.content)
                })
        
        return bedrock_messages

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """
        Generate chat response using AWS Bedrock.
        
        Args:
            messages: List of chat messages
            stop: List of stop sequences
            run_manager: Callback manager
            **kwargs: Additional parameters
            
        Returns:
            ChatResult containing the response
        """
        # If streaming is enabled, use streaming and collect all chunks
        if self.streaming:
            chunks = []
            full_content = ""
            for chunk in self._stream(messages, stop=stop, run_manager=run_manager, **kwargs):
                chunks.append(chunk)
                if chunk.message.content:
                    full_content += chunk.message.content
            
            # Return final result with complete content
            final_message = AIMessage(content=full_content)
            return ChatResult(generations=[ChatGeneration(message=final_message)])
        
        # Non-streaming path
        # Convert messages to Bedrock format
        bedrock_messages = self._convert_messages_to_bedrock_format(messages)
        
        # Prepare request body
        body = {
            "messages": bedrock_messages,
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "temperature": kwargs.get("temperature", self.temperature),
            "anthropic_version": "bedrock-2023-05-31"
        }
        
        # Add stop sequences if provided
        if stop:
            body["stop_sequences"] = stop
            
        try:
            # Make the API call
            response = self.client.invoke_model(
                modelId=self.model_id,
                contentType="application/json",
                accept="application/json",
                body=json.dumps(body).encode("utf-8")
            )
            
            # Parse response
            response_body = json.loads(response["body"].read().decode("utf-8"))
            
            # Extract the generated text
            if "content" in response_body and len(response_body["content"]) > 0:
                generated_text = response_body["content"][0]["text"]
            else:
                generated_text = ""
            
            # Create AIMessage
            message = AIMessage(content=generated_text)
            
            # Create ChatGeneration
            generation = ChatGeneration(message=message)
            
            return ChatResult(generations=[generation])
            
        except Exception as e:
            raise RuntimeError(f"Error calling AWS Bedrock: {str(e)}")

    def _stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGeneration]:
        """
        Stream chat response using AWS Bedrock.
        """
        # Convert messages to Bedrock format
        bedrock_messages = self._convert_messages_to_bedrock_format(messages)
        
        # Prepare request body
        body = {
            "messages": bedrock_messages,
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "temperature": kwargs.get("temperature", self.temperature),
            "anthropic_version": "bedrock-2023-05-31"
        }
        
        # Add stop sequences if provided
        if stop:
            body["stop_sequences"] = stop
            
        try:
            # Use streaming API
            response = self.client.invoke_model_with_response_stream(
                modelId=self.model_id,
                contentType="application/json",
                accept="application/json",
                body=json.dumps(body).encode("utf-8")
            )
            
            # Process streaming response
            for event in response["body"]:
                chunk = json.loads(event["chunk"]["bytes"].decode("utf-8"))
                
                if chunk.get("type") == "content_block_delta":
                    if "delta" in chunk and "text" in chunk["delta"]:
                        text = chunk["delta"]["text"]
                        if run_manager:
                            run_manager.on_llm_new_token(text)
                        yield ChatGeneration(message=AIMessage(content=text))
                        
        except Exception as e:
            # Fallback to non-streaming if streaming fails
            result = self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
            yield result.generations[0]

    @property
    def _llm_type(self) -> str:
        """Return type of chat model."""
        return "govcloud-bedrock"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Get identifying parameters."""
        return {
            "model_id": self.model_id,
            "region_name": self.region_name,
            "endpoint_url": self.endpoint_url,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
        }

