"""
AWS Bedrock LLM Provider for MCAL.
"""

import json
import asyncio
import logging
from typing import Dict, List, Any, Optional
import boto3
from botocore.exceptions import ClientError

from ..core.retry import (
    llm_retry,
    classify_boto_error,
    LLMThrottlingError,
    LLMServerError,
    LLMTimeoutError,
)
from ..core.rate_limiter import (
    TokenBucketRateLimiter,
    RateLimitConfig,
    create_rate_limiter,
)

logger = logging.getLogger(__name__)


class BedrockProvider:
    """AWS Bedrock LLM provider for MCAL."""
    
    MODELS = {
        # Fast models (simple tasks: intent detection, graph updates)
        # Use cross-region inference profile IDs for on-demand throughput
        "llama-3.1-8b": {"id": "us.meta.llama3-1-8b-instruct-v1:0", "max_tokens": 4096, "type": "llama", "tier": "fast"},
        "llama-3.2-3b": {"id": "us.meta.llama3-2-3b-instruct-v1:0", "max_tokens": 4096, "type": "llama", "tier": "fast"},
        # Smart models (complex tasks: decision extraction, reasoning)
        "llama-3.3-70b": {"id": "us.meta.llama3-3-70b-instruct-v1:0", "max_tokens": 4096, "type": "llama", "tier": "smart"},
        "llama-4-maverick": {"id": "us.meta.llama4-maverick-17b-instruct-v1:0", "max_tokens": 4096, "type": "llama", "tier": "smart"},
        "llama-3.1-70b": {"id": "us.meta.llama3-1-70b-instruct-v1:0", "max_tokens": 4096, "type": "llama", "tier": "smart"},
    }
    
    # Default models for tiered selection
    DEFAULT_FAST_MODEL = "llama-3.1-8b"
    DEFAULT_SMART_MODEL = "llama-3.3-70b"
    
    def __init__(
        self,
        model: str = "llama-3.3-70b",
        region: str = "us-east-1",
        temperature: float = 0.7,
        max_tokens: int = 2048,
        # Rate limiting options (Issue #39)
        enable_rate_limiting: bool = False,
        rate_limit_rpm: int = 60,
        rate_limit_tpm: int = 100_000,
        max_cost_per_hour: Optional[float] = None,
        enable_cost_logging: bool = True,
    ):
        """
        Initialize Bedrock provider.
        
        Args:
            model: Model name (e.g., "llama-3.3-70b")
            region: AWS region
            temperature: Sampling temperature
            max_tokens: Maximum tokens per response
            enable_rate_limiting: Enable rate limiting (default: False)
            rate_limit_rpm: Requests per minute limit (default: 60)
            rate_limit_tpm: Tokens per minute limit (default: 100,000)
            max_cost_per_hour: Maximum cost per hour in USD (None = unlimited)
            enable_cost_logging: Log cost estimates for each call
        """
        if model not in self.MODELS:
            raise ValueError(f"Unknown model: {model}")
        self.model_name = model
        self.model_config = self.MODELS[model]
        self.model_id = self.model_config["id"]
        self.temperature = temperature
        self.max_tokens = min(max_tokens, self.model_config["max_tokens"])
        self.bedrock = boto3.client('bedrock-runtime', region_name=region)
        
        # Token tracking
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        
        # Rate limiting (Issue #39)
        self._rate_limiter: Optional[TokenBucketRateLimiter] = None
        if enable_rate_limiting:
            self._rate_limiter = create_rate_limiter(
                model=model,
                provider="bedrock",
                rpm=rate_limit_rpm,
                tpm=rate_limit_tpm,
                max_cost_per_hour=max_cost_per_hour,
                enable_cost_logging=enable_cost_logging,
            )
            logger.info(
                f"Rate limiting enabled for {model}: "
                f"{rate_limit_rpm} RPM, {rate_limit_tpm} TPM"
                + (f", max ${max_cost_per_hour}/hr" if max_cost_per_hour else "")
            )
    
    def get_token_usage(self) -> dict:
        """Get total token usage."""
        return {
            "prompt_tokens": self.total_prompt_tokens,
            "completion_tokens": self.total_completion_tokens,
            "total_tokens": self.total_prompt_tokens + self.total_completion_tokens
        }
    
    def reset_token_usage(self):
        """Reset token counters."""
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
    
    @llm_retry(max_attempts=3, min_wait=1.0, max_wait=10.0)
    async def generate(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """
        Generate a response from Bedrock with automatic retry on transient failures.
        
        Retries on:
        - ThrottlingException (rate limits)
        - ServiceUnavailableException (5xx errors)
        - Connection timeouts
        
        Args:
            messages: List of message dicts with 'role' and 'content'
            **kwargs: Additional parameters (max_tokens, temperature)
            
        Returns:
            Generated text response
            
        Raises:
            LLMThrottlingError: After max retries on throttling
            LLMServerError: After max retries on server errors
            PermissionError: On access denied (not retried)
            RuntimeError: On other non-retryable errors or cost limit exceeded
        """
        # Rate limiting (Issue #39)
        if self._rate_limiter:
            # Estimate tokens based on message content
            total_chars = sum(len(m.get("content", "")) for m in messages)
            estimated_tokens = max(total_chars // 4, 500)  # Rough estimate
            await self._rate_limiter.acquire(estimated_tokens=estimated_tokens)
        
        try:
            # Convert message format for Converse API (content must be a list)
            formatted_messages = []
            for msg in messages:
                formatted_messages.append({
                    "role": msg["role"],
                    "content": [{"text": msg["content"]}]
                })
            
            response = await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: self.bedrock.converse(
                    modelId=self.model_id,
                    messages=formatted_messages,
                    inferenceConfig={
                        "maxTokens": kwargs.get("max_tokens", self.max_tokens),
                        "temperature": kwargs.get("temperature", self.temperature),
                        "topP": 0.9,
                    }
                )
            )
            
            # Track token usage from Bedrock response
            input_tokens = 0
            output_tokens = 0
            if 'usage' in response:
                input_tokens = response['usage'].get('inputTokens', 0)
                output_tokens = response['usage'].get('outputTokens', 0)
                self.total_prompt_tokens += input_tokens
                self.total_completion_tokens += output_tokens
            
            # Record usage for rate limiting (Issue #39)
            if self._rate_limiter:
                self._rate_limiter.record_usage(input_tokens, output_tokens)
            
            return response['output']['message']['content'][0]['text']
        except ClientError as e:
            error_code = e.response['Error']['Code']
            error_msg = e.response['Error']['Message']
            
            # Non-retryable: access denied
            if error_code == "AccessDeniedException":
                raise PermissionError(f"Bedrock access denied: {error_msg}")
            
            # Classify and potentially retry
            classified_error = classify_boto_error(error_code, error_msg)
            logger.warning(f"Bedrock error ({error_code}): {error_msg} - {'retrying' if isinstance(classified_error, (LLMThrottlingError, LLMServerError)) else 'not retrying'}")
            raise classified_error
    
    def get_rate_limit_stats(self) -> Optional[Dict[str, Any]]:
        """Get rate limiter statistics (if enabled)."""
        if self._rate_limiter:
            return self._rate_limiter.get_stats()
        return None


async def test_bedrock_provider():
    print("Testing Bedrock Provider...")
    print("=" * 70)
    for model in ["llama-3.3-70b", "llama-4-maverick"]:
        try:
            print(f"\nTesting {model}...")
            provider = BedrockProvider(model=model, max_tokens=100)
            messages = [{"role": "user", "content": "What is 2+2?"}]
            response = await provider.generate(messages)
            print(f"✅ Response: {response[:100]}")
        except Exception as e:
            print(f"❌ Error: {str(e)[:150]}")


if __name__ == "__main__":
    asyncio.run(test_bedrock_provider())
