"""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, BaseModel
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


def _pydantic_to_bedrock_tool_schema(pydantic_model: BaseModel, tool_name: str = "structured_output") -> Dict:
    """
    Convert a Pydantic model to Bedrock tool schema format.
    
    Args:
        pydantic_model: Pydantic model class
        tool_name: Name for the tool
        
    Returns:
        Dictionary in Bedrock tool format
    """
    schema = pydantic_model.model_json_schema()
    
    # Extract properties and required fields from the Pydantic schema
    properties = schema.get("properties", {})
    required = schema.get("required", [])
    description = schema.get("description", f"Structured output for {pydantic_model.__name__}")
    
    return {
        "name": tool_name,
        "description": description,
        "input_schema": {
            "type": "object",
            "properties": properties,
            "required": required
        }
    }


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:
                            # Check if run_manager is async or sync
                            if hasattr(run_manager.on_llm_new_token, '__call__'):
                                try:
                                    import asyncio
                                    if asyncio.iscoroutinefunction(run_manager.on_llm_new_token):
                                        # Skip async callback in sync context for now
                                        pass
                                    else:
                                        run_manager.on_llm_new_token(text)
                                except:
                                    run_manager.on_llm_new_token(text)
                            else:
                                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]

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """
        Async generate chat response using AWS Bedrock.
        """
        # For now, run the sync version in a thread pool
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None, 
            lambda: self._generate(messages, stop, run_manager, **kwargs)
        )

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGeneration]:
        """
        Async stream chat response using AWS Bedrock.
        """
        # For now, run the sync version in a thread pool
        import asyncio
        loop = asyncio.get_event_loop()
        
        def sync_stream():
            return list(self._stream(messages, stop, run_manager, **kwargs))
        
        chunks = await loop.run_in_executor(None, sync_stream)
        for chunk in chunks:
            yield chunk

    def structured_generate(
        self,
        messages: List[BaseMessage], 
        pydantic_model: BaseModel,
        tool_name: str = "structured_output",
        **kwargs: Any
    ) -> BaseModel:
        """
        Generate structured output using Bedrock tool calling.
        
        Args:
            messages: List of chat messages
            pydantic_model: Pydantic model class to structure the output
            tool_name: Name for the tool
            **kwargs: Additional parameters
            
        Returns:
            Instance of the pydantic_model with parsed data
        """
        # Convert Pydantic model to Bedrock tool schema
        tool_schema = _pydantic_to_bedrock_tool_schema(pydantic_model, tool_name)
        
        # Convert messages to Bedrock format
        bedrock_messages = self._convert_messages_to_bedrock_format(messages)
        
        # Build request body
        request_body = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "messages": bedrock_messages,
            "tools": [tool_schema],
            "tool_choice": {"type": "tool", "name": tool_name}
        }
        
        # Filter out parameters that aren't supported by Bedrock tool calling
        bedrock_kwargs = {}
        supported_params = {'system', 'anthropic_version'}
        for key, value in kwargs.items():
            if key in supported_params:
                bedrock_kwargs[key] = value
        
        # Add supported kwargs to request body  
        request_body.update(bedrock_kwargs)
        
        try:
            # Call Bedrock
            response = self.client.invoke_model(
                modelId=self.model_id,
                body=json.dumps(request_body)
            )
            
            response_body = json.loads(response['body'].read())
            
            # Extract tool use from response
            content = response_body.get('content', [])
            for item in content:
                if item.get('type') == 'tool_use' and item.get('name') == tool_name:
                    tool_input = item.get('input', {})
                    # Parse the tool input into the Pydantic model
                    return pydantic_model(**tool_input)
            
            # If no tool use found, try to parse the text content as JSON
            for item in content:
                if item.get('type') == 'text':
                    try:
                        data = json.loads(item.get('text', '{}'))
                        return pydantic_model(**data)
                    except (json.JSONDecodeError, ValueError):
                        continue
            
            raise ValueError("No valid structured output found in response")
            
        except Exception as e:
            raise RuntimeError(f"Error in structured generation: {str(e)}")

    def with_structured_output(self, pydantic_model: BaseModel, **kwargs):
        """
        Return a wrapper that uses native structured output.
        This method makes ChatGovCloudBedrock compatible with LangChain's structured output interface.
        """
        class StructuredOutputWrapper:
            def __init__(self, chat_model, pydantic_model):
                self.chat_model = chat_model
                self.pydantic_model = pydantic_model
                
            def invoke(self, input_data, **invoke_kwargs):
                # Handle different input formats
                if hasattr(input_data, 'messages'):
                    messages = input_data.messages
                elif isinstance(input_data, list):
                    messages = input_data
                elif isinstance(input_data, str):
                    from langchain_core.messages import HumanMessage
                    messages = [HumanMessage(content=input_data)]
                else:
                    raise ValueError(f"Unsupported input type: {type(input_data)}")
                
                return self.chat_model.structured_generate(
                    messages=messages,
                    pydantic_model=self.pydantic_model,
                    **invoke_kwargs
                )
        
        return StructuredOutputWrapper(self, pydantic_model)

    @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,
        }

