"""Main Bedrock LLM client wrapper"""

import json
import asyncio
import time
from datetime import datetime, timezone
from typing import Optional, AsyncIterator, List, Dict, Any, Type
from pydantic import BaseModel
from logorator import Logger
from .config import BedrockConfig
from ..models import (
    TextRequest, 
    MessageRequest,
    TextResponse, 
    StreamChunk,
)
from ..utils import pydantic_to_tool_schema, TwoLevelCache, retry_on_error

# Default Bedrock model quotas for concurrency limiting
DEFAULT_MODEL_QUOTAS = {
    'claude-3-5-sonnet-v2': {'rpm': 10, 'tpm': 200000, 'concurrent': 1},
    'claude-3-5-sonnet': {'rpm': 200, 'tpm': 400000, 'concurrent': 2},
    'claude-3-sonnet': {'rpm': 200, 'tpm': 400000, 'concurrent': 2},
    'claude-3-haiku': {'rpm': 400, 'tpm': 400000, 'concurrent': 5},
    'claude-3-opus': {'rpm': 50, 'tpm': 200000, 'concurrent': 1},
    'llama': {'rpm': 500, 'tpm': 500000, 'concurrent': 5},
    'mistral': {'rpm': 300, 'tpm': 300000, 'concurrent': 3},
    'titan': {'rpm': 400, 'tpm': 400000, 'concurrent': 5},
}


class BedrockLLMClient:
    """Async client for text generation with AWS Bedrock LLMs"""

    def __init__(self, config: Optional[BedrockConfig] = None, max_concurrent: Optional[int] = None, dynamo_table_name: Optional[str] = None, cache_ttl_days: Optional[float] = None):
        """Initialize the Bedrock client
        
        Args:
            config: BedrockConfig instance. If None, creates default config.
            max_concurrent: Max concurrent requests. Overrides config.max_concurrent if provided.
            dynamo_table_name: DynamoDB table name for shared cache. If None, only local cache is used.
            cache_ttl_days: TTL for DynamoDB cache entries in days. Defaults to 365.
        """
        self.config = config or BedrockConfig()
        self.config.validate()
        self.client = None
        self.models_client = None
        cache_kwargs = {"dynamo_table_name": dynamo_table_name}
        if cache_ttl_days is not None:
            cache_kwargs["ttl_days"] = cache_ttl_days
        self.cache = TwoLevelCache(**cache_kwargs)
        self._semaphores: Dict[str, asyncio.Semaphore] = {}
        self._max_concurrent = max_concurrent if max_concurrent is not None else self.config.max_concurrent

    async def _init_client(self):
        """Initialize aioboto3 Bedrock client"""
        try:
            import aioboto3
            creds = self.config.get_credentials()
            session = aioboto3.Session()
            self.client = await session.client("bedrock-runtime", **creds).__aenter__()
            self.models_client = await session.client("bedrock", **creds).__aenter__()
        except ImportError:
            raise ImportError("aioboto3 is required. Install with: pip install aioboto3")
        except Exception:
            raise

    async def close(self):
        """Close the client connections"""
        if self.client:
            await self.client.__aexit__(None, None, None)
        if self.models_client:
            await self.models_client.__aexit__(None, None, None)

    async def __aenter__(self):
        """Async context manager entry"""
        await self._init_client()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        await self.close()

    def _get_semaphore(self, model: str) -> asyncio.Semaphore:
        """Get or create semaphore for model to limit concurrent requests"""
        if model not in self._semaphores:
            # Use explicit max_concurrent or infer from model defaults
            if self._max_concurrent:
                limit = self._max_concurrent
            else:
                # Get default concurrent limit for this model
                limit = 2  # Safe default
                for pattern, quotas in DEFAULT_MODEL_QUOTAS.items():
                    if pattern in model.lower():
                        limit = quotas['concurrent']
                        break
            
            self._semaphores[model] = asyncio.Semaphore(limit)
        
        return self._semaphores[model]

    async def _invoke_model_with_retry(self, **kwargs):
        """Invoke model with retry logic"""
        @retry_on_error(
            max_retries=self.config.max_retries,
            base_delay=self.config.retry_delay,
            max_delay=self.config.max_retry_delay,
        )
        async def _invoke():
            return await self.client.invoke_model(**kwargs)
        
        return await _invoke()

    async def list_available_models(self) -> List[Dict[str, Any]]:
        """List all available models in Bedrock"""
        if not self.models_client:
            await self._init_client()
        try:
            response = await self.models_client.list_foundation_models()
            return response.get("modelSummaries", [])
        except Exception:
            return []
    
    async def list_available_model_ids(self) -> List[str]:
        """List all available model IDs in Bedrock
        
        Returns:
            List of model ID strings
        """
        models = await self.list_available_models()
        return [m.get("modelId") for m in models if isinstance(m, dict) and m.get("modelId")]

    def __str__(self):
        return f"BedrockLLMClient({self.config.default_model})"

    @Logger(exclude_args=[])
    async def generate_text(self, request: TextRequest) -> TextResponse:
        """Generate text from a prompt"""
        if not self.client:
            await self._init_client()
            
        model = request.model or self.config.default_model
        # If no temperature specified, use 0 (deterministic + cacheable)
        temperature = request.temperature if request.temperature is not None else 0
        
        # Generate cache key for this specific request
        cache_key = None
        if temperature == 0 and not request.stream:
            cache_key = self._generate_cache_key(
                model=model,
                prompt=request.prompt,
                max_tokens=request.max_tokens or self.config.max_tokens,
                top_p=request.top_p,
                top_k=request.top_k,
                system_prompt=request.system_prompt,
                response_format=request.response_format.__name__ if request.response_format else None
            )
        
        if request.clear_cache and cache_key:
            self.cache.clear(cache_key)
        
        if request.use_cache and cache_key:
            cached, cache_source = self.cache.get(cache_key)
            if cached:
                Logger.note(f"Cache hit [{cache_key[:8]}] - {model}")
                result = self._deserialize_response(cached["data"], request.response_format)
                result.cache_source = cache_source
                result.cache_key = cache_key
                return result
        
        if request.reasoning_effort:
            Logger.note(f"reasoning_effort='{request.reasoning_effort}' ignored by Bedrock")

        prompt_preview = request.prompt[:60] + "..." if len(request.prompt) > 60 else request.prompt
        Logger.note(f"{model} | temp={temperature} | {prompt_preview}")
        
        body = self._build_request_body(
            model=model,
            prompt=request.prompt,
            temperature=temperature,
            max_tokens=request.max_tokens or self.config.max_tokens,
            top_p=request.top_p or self.config.top_p,
            top_k=request.top_k or self.config.top_k,
            system_prompt=request.system_prompt,
            response_format=request.response_format,
        )

        try:
            started_at = datetime.now(timezone.utc).isoformat()
            t0 = time.monotonic()
            semaphore = self._get_semaphore(model)
            async with semaphore:
                response = await self._invoke_model_with_retry(
                    modelId=model,
                    body=json.dumps(body),
                    contentType="application/json",
                )
            elapsed = round(time.monotonic() - t0, 3)
            
            response_body = json.loads(await response["body"].read())
            result = self._parse_response(response_body, model, request.response_format)
            result.timestamp = started_at
            result.elapsed_seconds = elapsed
            result.metadata["prompt"] = request.prompt
            result.metadata["response_format"] = request.response_format.model_json_schema() if request.response_format else None
            Logger.note(f"{result.input_tokens} in / {result.output_tokens} out | {result.text[:50]}")
            
            if cache_key:
                cache_metadata = {
                    "prompt": request.prompt,
                    "model": model,
                    "temperature": temperature,
                    "max_tokens": request.max_tokens or self.config.max_tokens,
                    "system_prompt": request.system_prompt,
                    "response_format": request.response_format.__name__ if request.response_format else None,
                    "top_p": request.top_p or self.config.top_p,
                    "top_k": request.top_k or self.config.top_k,
                }
                self.cache.set(cache_key, self._serialize_response(result), cache_metadata)
            result.cache_key = cache_key
            return result
            
        except Exception:
            raise

    async def generate_text_stream(self, request: TextRequest) -> AsyncIterator[StreamChunk]:
        """Stream text generation
        
        Args:
            request: TextRequest with prompt and parameters
            
        Yields:
            StreamChunk objects with partial text
        """
        if not self.client:
            await self._init_client()
            
        model = request.model or self.config.default_model
        
        body = self._build_request_body(
            model=model,
            prompt=request.prompt,
            temperature=request.temperature or self.config.temperature,
            max_tokens=request.max_tokens or self.config.max_tokens,
            top_p=request.top_p or self.config.top_p,
            top_k=request.top_k or self.config.top_k,
        )

        try:
            response = await self.client.invoke_model_with_response_stream(
                modelId=model,
                body=json.dumps(body),
                contentType="application/json",
            )
            
            async for event in response["body"]:
                if "chunk" in event:
                    chunk_data = json.loads(event["chunk"]["bytes"])
                    text = self._extract_text_from_chunk(chunk_data, model)
                    if text:
                        yield StreamChunk(text=text, model=model)
                        
        except Exception:
            raise

    @Logger(exclude_args=[])
    async def send_message(self, request: MessageRequest) -> TextResponse:
        """Send a message in a conversation"""
        if not self.client:
            await self._init_client()
            
        model = request.model or self.config.default_model
        # If no temperature specified, use 0 (deterministic + cacheable)
        temperature = request.temperature if request.temperature is not None else 0
        
        # Generate cache key for this specific request
        cache_key = None
        if temperature == 0 and not request.stream:
            messages_str = json.dumps([{"role": m.role, "content": m.content} for m in request.messages])
            cache_key = self._generate_cache_key(
                model=model,
                messages=messages_str,
                max_tokens=request.max_tokens or self.config.max_tokens,
                system_prompt=request.system_prompt,
                response_format=request.response_format.__name__ if request.response_format else None
            )
        
        if request.clear_cache and cache_key:
            self.cache.clear(cache_key)
        
        if request.use_cache and cache_key:
            cached, cache_source = self.cache.get(cache_key)
            if cached:
                Logger.note(f"Cache hit [{cache_key[:8]}] - {model}")
                result = self._deserialize_response(cached["data"], request.response_format)
                result.cache_source = cache_source
                result.cache_key = cache_key
                return result
        
        last_msg = request.messages[-1].content[:60] if request.messages else ""
        Logger.note(f"{model} | {len(request.messages)} messages | {last_msg}")
        
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        body = {
            "anthropic_version": "bedrock-2023-05-31",
            "messages": messages,
            "max_tokens": request.max_tokens or self.config.max_tokens,
            "temperature": temperature,
        }
        
        if request.system_prompt:
            body["system"] = request.system_prompt
            
        if request.response_format and "claude" in model.lower():
            tool_schema = pydantic_to_tool_schema(request.response_format)
            body["tools"] = [tool_schema]
            body["tool_choice"] = {"type": "tool", "name": tool_schema["name"]}

        try:
            started_at = datetime.now(timezone.utc).isoformat()
            t0 = time.monotonic()
            semaphore = self._get_semaphore(model)
            async with semaphore:
                response = await self._invoke_model_with_retry(
                    modelId=model,
                    body=json.dumps(body),
                    contentType="application/json",
                )
            elapsed = round(time.monotonic() - t0, 3)
            
            response_body = json.loads(await response["body"].read())
            result = self._parse_response(response_body, model, request.response_format)
            result.timestamp = started_at
            result.elapsed_seconds = elapsed
            result.metadata["messages"] = [{"role": m.role, "content": m.content} for m in request.messages]
            result.metadata["response_format"] = request.response_format.model_json_schema() if request.response_format else None
            Logger.note(f"{result.input_tokens} in / {result.output_tokens} out | {result.text[:50]}")
            
            if cache_key:
                cache_metadata = {
                    "messages": [{"role": m.role, "content": m.content} for m in request.messages],
                    "model": model,
                    "temperature": temperature,
                    "max_tokens": request.max_tokens or self.config.max_tokens,
                    "system_prompt": request.system_prompt,
                    "response_format": request.response_format.__name__ if request.response_format else None,
                }
                self.cache.set(cache_key, self._serialize_response(result), cache_metadata)
            result.cache_key = cache_key
            return result
            
        except Exception:
            raise

    async def send_message_stream(self, request: MessageRequest) -> AsyncIterator[StreamChunk]:
        """Stream a conversation message
        
        Args:
            request: MessageRequest with message history
            
        Yields:
            StreamChunk objects with partial responses
        """
        if not self.client:
            await self._init_client()
            
        model = request.model or self.config.default_model
        
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        body = {
            "anthropic_version": "bedrock-2023-05-31",
            "messages": messages,
            "max_tokens": request.max_tokens or self.config.max_tokens,
            "temperature": request.temperature or self.config.temperature,
        }
        
        if request.system_prompt:
            body["system"] = request.system_prompt

        try:
            response = await self.client.invoke_model_with_response_stream(
                modelId=model,
                body=json.dumps(body),
                contentType="application/json",
            )
            
            async for event in response["body"]:
                if "chunk" in event:
                    chunk_data = json.loads(event["chunk"]["bytes"])
                    text = self._extract_text_from_chunk(chunk_data, model)
                    if text:
                        yield StreamChunk(text=text, model=model)
                        
        except Exception:
            raise

    def _build_request_body(
        self,
        model: str,
        prompt: str,
        temperature: float,
        max_tokens: int,
        top_p: float,
        top_k: int,
        system_prompt: Optional[str] = None,
        response_format: Optional[Type[BaseModel]] = None,
    ) -> Dict[str, Any]:
        """Build request body for text generation based on model type"""
        
        if "claude" in model.lower():
            # Claude 3+ models use Messages API
            body = {
                "anthropic_version": "bedrock-2023-05-31",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": max_tokens,
                "temperature": temperature,
            }
            if system_prompt:
                body["system"] = system_prompt
            if response_format:
                tool_schema = pydantic_to_tool_schema(response_format)
                body["tools"] = [tool_schema]
                body["tool_choice"] = {"type": "tool", "name": tool_schema["name"]}
        elif "llama" in model.lower():
            # Llama models
            body = {
                "prompt": prompt,
                "max_gen_len": max_tokens,
                "temperature": temperature,
                "top_p": top_p,
            }
        elif "mistral" in model.lower():
            # Mistral models
            body = {
                "prompt": prompt,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p,
            }
        else:
            # Default/generic format
            body = {
                "prompt": prompt,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p,
            }
        
        return body

    def _parse_response(self, response_body: Dict[str, Any], model: str, response_format: Optional[Type[BaseModel]] = None) -> TextResponse:
        """Parse response based on model type"""
        if "claude" in model.lower():
            # Check for tool use (structured output)
            content = response_body.get("content", [])
            if content and content[0].get("type") == "tool_use" and response_format:
                tool_input = content[0].get("input", {})
                structured_data = response_format(**tool_input)
                text = json.dumps(tool_input, indent=2)
            else:
                # Regular text response
                text = response_body["content"][0]["text"]
                structured_data = None
                
            stop_reason = response_body.get("stop_reason", "")
            input_tokens = response_body.get("usage", {}).get("input_tokens", 0)
            output_tokens = response_body.get("usage", {}).get("output_tokens", 0)
        elif "llama" in model.lower():
            # Llama response format
            text = response_body.get("generation", "")
            stop_reason = response_body.get("stop_reason", "")
            input_tokens = 0
            output_tokens = 0
        elif "mistral" in model.lower():
            # Mistral response format
            outputs = response_body.get("outputs", [])
            text = outputs[0].get("text", "") if outputs else ""
            stop_reason = outputs[0].get("stop_reason", "") if outputs else ""
            input_tokens = 0
            output_tokens = 0
        else:
            # Generic handling
            text = response_body.get("generated_text", response_body.get("generation", ""))
            stop_reason = response_body.get("stop_reason", "")
            input_tokens = 0
            output_tokens = 0

        return TextResponse(
            text=text,
            model=model,
            stop_reason=stop_reason,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            structured_data=structured_data,
        )

    def _extract_text_from_chunk(self, chunk_data: Dict[str, Any], model: str) -> str:
        """Extract text from streaming chunk based on model type"""
        if "claude" in model.lower():
            if chunk_data.get("type") == "content_block_delta":
                return chunk_data.get("delta", {}).get("text", "")
        elif "llama" in model.lower():
            return chunk_data.get("generation", "")
        
        return ""
    
    def _generate_cache_key(self, **kwargs) -> str:
        """Generate cache key from request parameters"""
        return self.cache._generate_key(**kwargs)
    
    def _serialize_response(self, response: TextResponse) -> Dict[str, Any]:
        """Serialize TextResponse for caching"""
        return {
            "text": response.text,
            "model": response.model,
            "stop_reason": response.stop_reason,
            "input_tokens": response.input_tokens,
            "output_tokens": response.output_tokens,
            "reasoning_tokens": response.reasoning_tokens,
            "cached_tokens": response.cached_tokens,
            "timestamp": response.timestamp,
            "elapsed_seconds": response.elapsed_seconds,
            "metadata": response.metadata,
            "structured_data": response.structured_data.model_dump() if response.structured_data else None,
        }
    
    def _deserialize_response(self, data: Dict[str, Any], response_format: Optional[Type[BaseModel]] = None) -> TextResponse:
        """Deserialize cached data back to TextResponse"""
        structured_data = None
        if data.get("structured_data") and response_format:
            structured_data = response_format(**data["structured_data"])
        
        return TextResponse(
            text=data["text"],
            model=data["model"],
            stop_reason=data["stop_reason"],
            input_tokens=data["input_tokens"],
            output_tokens=data["output_tokens"],
            reasoning_tokens=data.get("reasoning_tokens", 0),
            cached_tokens=data.get("cached_tokens", 0),
            timestamp=data.get("timestamp"),
            elapsed_seconds=data.get("elapsed_seconds"),
            metadata=data.get("metadata", {}),
            structured_data=structured_data,
        )
