"""AWS Bedrock model implementation.

This module provides the Bedrock model implementation using the
AWS Bedrock Converse API.
"""

import logging
from typing import Any, Dict, List, Optional

from agent_factory.models.base import (
    BaseModel,
    Message,
    ModelError,
    ModelNotAvailableError,
    ModelRateLimitError,
    ModelResponse,
    StopReason,
    ToolDefinition,
    ToolUse,
)

logger = logging.getLogger(__name__)


class BedrockModel(BaseModel):
    """AWS Bedrock model implementation.

    Uses the Bedrock Converse API for consistent interface across
    different foundation models.

    Example:
        model = BedrockModel(
            model_id="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
            region="us-east-1",
        )
        response = model.converse(messages, system_prompt)
    """

    def __init__(
        self,
        model_id: str = "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
        region: str = "us-east-1",
        client: Optional[Any] = None,
    ):
        """Initialize Bedrock model.

        Args:
            model_id: Bedrock model ID (cross-region inference profile)
            region: AWS region
            client: Optional pre-configured boto3 client
        """
        self._model_id = model_id
        self._region = region
        self._client = client

    def _get_client(self):
        """Get or create the Bedrock client."""
        if self._client is None:
            try:
                import boto3

                self._client = boto3.client(
                    "bedrock-runtime",
                    region_name=self._region,
                )
            except ImportError:
                raise ModelNotAvailableError(
                    "boto3 is required for Bedrock. "
                    "Install with: pip install agent-factory-core[bedrock]",
                    model_id=self._model_id,
                )
        return self._client

    def converse(
        self,
        messages: List[Message],
        system_prompt: str,
        tools: Optional[List[ToolDefinition]] = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> ModelResponse:
        """Send messages to Bedrock and get a response."""
        client = self._get_client()

        # Convert messages to Bedrock format
        bedrock_messages = [msg.to_bedrock_format() for msg in messages]

        # Build request
        request: Dict[str, Any] = {
            "modelId": self._model_id,
            "messages": bedrock_messages,
            "system": [{"text": system_prompt}],
            "inferenceConfig": {
                "maxTokens": max_tokens,
                "temperature": temperature,
            },
        }

        # Add tools if provided
        if tools:
            request["toolConfig"] = {"tools": [tool.to_bedrock_format() for tool in tools]}

        try:
            response = client.converse(**request)
            return self._parse_response(response)

        except client.exceptions.ThrottlingException as e:
            raise ModelRateLimitError(
                f"Bedrock rate limited: {e}",
                model_id=self._model_id,
                original_error=e,
            )
        except Exception as e:
            # Check for transient errors
            is_transient = "throttl" in str(e).lower() or "timeout" in str(e).lower()
            raise ModelError(
                f"Bedrock error: {e}",
                model_id=self._model_id,
                is_transient=is_transient,
                original_error=e,
            )

    def _parse_response(self, response: Dict[str, Any]) -> ModelResponse:
        """Parse Bedrock response into ModelResponse."""
        output = response.get("output", {})
        message = output.get("message", {})
        content_items = message.get("content", [])

        # Extract text content
        text_parts = []
        tool_uses = []

        for item in content_items:
            if "text" in item:
                text_parts.append(item["text"])
            elif "toolUse" in item:
                tu = item["toolUse"]
                tool_uses.append(
                    ToolUse(
                        tool_use_id=tu.get("toolUseId", ""),
                        name=tu.get("name", ""),
                        input=tu.get("input", {}),
                    )
                )

        # Map stop reason
        stop_reason_map = {
            "end_turn": StopReason.END_TURN,
            "tool_use": StopReason.TOOL_USE,
            "max_tokens": StopReason.MAX_TOKENS,
            "stop_sequence": StopReason.STOP_SEQUENCE,
        }
        raw_stop = response.get("stopReason", "end_turn")
        stop_reason = stop_reason_map.get(raw_stop, StopReason.END_TURN)

        # Extract usage
        usage = response.get("usage", {})

        return ModelResponse(
            content="\n".join(text_parts),
            stop_reason=stop_reason,
            tool_uses=tool_uses,
            usage={
                "input_tokens": usage.get("inputTokens", 0),
                "output_tokens": usage.get("outputTokens", 0),
            },
            model_id=self._model_id,
            raw_response=response,
        )

    def is_available(self) -> bool:
        """Check if Bedrock is available."""
        try:
            self._get_client()
            return True
        except Exception:
            return False

    @property
    def model_id(self) -> str:
        """Get the model identifier."""
        return self._model_id
