import asyncio
import base64
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import (
    Any,
    AsyncGenerator,
    Dict,
    Iterable,
    Literal,
    Mapping,
    Optional,
    Sequence,
    Union,
    overload,
)

import aiohttp
import boto3
import logsim
import orjson
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.eventstream import EventStreamBuffer


class BedrockStreamError(Exception):
    """Raised when the Bedrock event stream surfaces an error payload."""

    def __init__(self, message_type: str, exception_type: str, detail: str):
        self.message_type = message_type
        self.exception_type = exception_type or "UnknownException"
        self.detail = detail
        super().__init__(
            f"Bedrock Error ({self.message_type}/{self.exception_type}): {self.detail}"  # noqa: E501
        )


class BedrockClientError(Exception):
    """Raised when the Bedrock Invoke API returns a non-success HTTP."""

    def __init__(self, status: int, error_type: str, detail: str):
        self.status = status
        self.error_type = error_type
        self.detail = detail
        super().__init__(f"{status} {error_type}: {detail}")


log = logsim.CustomLogger()


class Client:
    def __init__(
        self,
        region_name: str,
        assume_role_arn: Optional[str] = None,
        *,
        max_connections: int = 10000,
        request_timeout: Optional[float] = None,
        max_concurrency: Optional[int] = None,
        max_retries: int = 2,
        retry_backoff: float = 0.5,
        max_backoff: float = 6.0,
        retry_statuses: Optional[Sequence[int]] = None,
    ):
        self.region_name = region_name
        self.assume_role_arn = assume_role_arn
        self.connector = aiohttp.TCPConnector(
            limit=max_connections,
            ttl_dns_cache=3600,
            use_dns_cache=True,
            enable_cleanup_closed=True,
        )
        self.session = None
        self.expiration = None
        self.access_key = None
        self.secret_key = None
        self.session_token = None
        self.credentials = None

        self._client_timeout = (
            aiohttp.ClientTimeout(total=request_timeout) if request_timeout else None  # noqa: E501
        )
        self._max_concurrency = (
            max_concurrency if max_concurrency and max_concurrency > 0 else None  # noqa: E501
        )
        self._max_retries = max(0, max_retries)
        self._retry_backoff = max(0.0, retry_backoff)
        self._retry_backoff_cap = (
            max(max_backoff, self._retry_backoff)
            if max_backoff > 0
            else max(self._retry_backoff, 0.0)
        )
        self._retry_statuses = (
            tuple(retry_statuses)
            if retry_statuses
            else (
                408,
                424,
                429,
                500,
                502,
                503,
                504,
            )
        )
        self._request_semaphore: Optional[asyncio.Semaphore] = None
        self._credential_lock: Optional[asyncio.Lock] = None

        # Initialize credentials
        self._refresh_credentials_sync()

    def _refresh_credentials_sync(self):
        """Refresh AWS credentials, handling role assumption if needed"""
        if self.assume_role_arn:
            sts_client = boto3.client("sts")
            response = sts_client.assume_role(
                RoleArn=self.assume_role_arn,
                RoleSessionName="aiobedrock",
            )

            # Extract temporary credentials
            credentials = response["Credentials"]
            self.access_key = credentials["AccessKeyId"]
            self.secret_key = credentials["SecretAccessKey"]
            self.session_token = credentials["SessionToken"]
            self.expiration = credentials["Expiration"]

            log.info(f"Refreshed credentials, expires at: {self.expiration}")

            # Create session with temporary credentials
            boto3_session = boto3.Session(
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                aws_session_token=self.session_token,
                region_name=self.region_name,
            )
        else:
            # Use default credentials
            boto3_session = boto3.Session(region_name=self.region_name)

        self.credentials = boto3_session.get_credentials()

    async def _refresh_credentials(self):
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, self._refresh_credentials_sync)

    def _are_credentials_expired(self) -> bool:
        """Check if the current credentials are expired or about to expire"""
        if not self.assume_role_arn or not self.expiration:
            return False

        # Check if credentials expire within the next 5 minutes
        current_time = datetime.now(timezone.utc)
        expiration_time = self.expiration

        # Handle timezone-aware expiration time
        if expiration_time.tzinfo is None:
            expiration_time = expiration_time.replace(tzinfo=timezone.utc)

        time_until_expiration = expiration_time - current_time
        return time_until_expiration.total_seconds() < 300  # 5 minutes

    async def _ensure_valid_credentials(self):
        """Ensure credentials are valid, refreshing if necessary"""
        if self.credentials is None:
            await self._refresh_credentials()
            return

        if not self.assume_role_arn:
            return

        if not self._are_credentials_expired():
            return

        if self._credential_lock is None:
            self._credential_lock = asyncio.Lock()

        async with self._credential_lock:
            if self._are_credentials_expired():
                log.info("Credentials expired or expiring soon, refreshing...")
                await self._refresh_credentials()

    def _ensure_session(self) -> aiohttp.ClientSession:
        """Create the shared aiohttp session on first use."""

        if self.session is None or getattr(self.session, "closed", False):
            session_kwargs: Dict[str, Any] = {"connector": self.connector}
            if self._client_timeout is not None:
                session_kwargs["timeout"] = self._client_timeout
            self.session = aiohttp.ClientSession(**session_kwargs)
        return self.session

    async def __aenter__(self):
        self._ensure_session()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def close(self):
        if self.session is not None:
            await self.session.close()
            self.session = None

    async def _acquire_request_slot(self) -> Optional[asyncio.Semaphore]:
        if self._max_concurrency is None:
            return None

        if self._request_semaphore is None:
            self._request_semaphore = asyncio.Semaphore(self._max_concurrency)

        await self._request_semaphore.acquire()
        return self._request_semaphore

    @asynccontextmanager
    async def _request(
        self,
        url: str,
        headers: Dict[str, str],
        data: Union[str, bytes],
    ):
        limiter = await self._acquire_request_slot()
        session = self._ensure_session()
        post_kwargs: Dict[str, Any] = {
            "url": url,
            "headers": headers,
            "data": data,
        }
        if self._client_timeout is not None:
            post_kwargs["timeout"] = self._client_timeout

        try:
            async with session.post(**post_kwargs) as response:
                yield response
        finally:
            if limiter is not None:
                limiter.release()

    def _should_retry(self, exc: Exception, attempt: int) -> bool:
        if attempt >= self._max_retries:
            return False

        if isinstance(exc, BedrockClientError):
            return exc.status in self._retry_statuses

        if isinstance(exc, (aiohttp.ClientError, asyncio.TimeoutError)):
            return True

        return False

    async def _sleep_backoff(self, attempt: int) -> None:
        if self._retry_backoff <= 0:
            if attempt > 0:
                await asyncio.sleep(0)
            return

        delay = min(
            self._retry_backoff * (2**attempt),
            self._retry_backoff_cap,
        )
        if delay > 0:
            await asyncio.sleep(delay)

    async def invoke_model(self, body: str, modelId: str, **kwargs) -> bytes:
        """Invoke a model and return the response body as bytes"""
        url = f"https://bedrock-runtime.{self.region_name}.amazonaws.com/model/{modelId}/invoke"  # noqa: E501

        attempt = 0
        while True:
            await self._ensure_valid_credentials()
            headers = self._signed_request(
                body=body,
                url=url,
                method="POST",
                credentials=self.credentials,
                region_name=self.region_name,
                **kwargs,
            )

            try:
                async with self._request(url=url, headers=headers, data=body) as res:  # noqa: E501
                    await self._handle_error_response(res)
                    return await res.read()
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                if not self._should_retry(exc, attempt):
                    raise
                await self._sleep_backoff(attempt)
                attempt += 1

    async def invoke_model_with_response_stream(
        self, body: str, modelId: str, **kwargs
    ) -> AsyncGenerator[Union[Dict[str, Any], bytes], None]:
        """
        Invoke a model with streaming response
        """
        url = f"https://bedrock-runtime.{self.region_name}.amazonaws.com/model/{modelId}/invoke-with-response-stream"  # noqa: E501
        attempt = 0
        emitted = False

        while True:
            await self._ensure_valid_credentials()
            headers = self._signed_request(
                body=body,
                url=url,
                method="POST",
                credentials=self.credentials,
                region_name=self.region_name,
                **kwargs,
            )

            try:
                async with self._request(url=url, headers=headers, data=body) as res:  # noqa: E501
                    await self._handle_error_response(res)

                    event_stream_buffer = EventStreamBuffer()

                    async for chunk in res.content.iter_chunked(8192):
                        event_stream_buffer.add_data(chunk)

                        while True:
                            try:
                                message = event_stream_buffer.next()
                                processed_content = self._process_event_message(  # noqa: E501
                                    message=message,
                                )
                                if processed_content is not None:
                                    emitted = True
                                    yield processed_content
                            except StopIteration:
                                break
                            except Exception as e:
                                log.error(f"Error processing event: {e}")
                                break

                    while True:
                        try:
                            message = event_stream_buffer.next()
                            processed_content = self._process_event_message(
                                message,
                            )
                            if processed_content is not None:
                                emitted = True
                                yield processed_content
                        except StopIteration:
                            return
                        except Exception as e:
                            log.error(f"Error processing buffer: {e}")
                            return
            except asyncio.CancelledError:
                raise
            except BedrockStreamError:
                raise
            except Exception as exc:
                if emitted or not self._should_retry(exc, attempt):
                    raise
                await self._sleep_backoff(attempt)
                attempt += 1

    @overload
    async def invoke_many(
        self,
        requests: Iterable[Mapping[str, Any]],
        *,
        concurrency: Optional[int] = None,
        return_exceptions: Literal[False] = False,
    ) -> Sequence[bytes]: ...

    @overload
    async def invoke_many(
        self,
        requests: Iterable[Mapping[str, Any]],
        *,
        concurrency: Optional[int] = None,
        return_exceptions: Literal[True],
    ) -> Sequence[Union[bytes, BaseException]]: ...

    async def invoke_many(
        self,
        requests: Iterable[Mapping[str, Any]],
        *,
        concurrency: Optional[int] = None,
        return_exceptions: bool = False,
    ) -> Sequence[Union[bytes, BaseException]]:
        """Invoke multiple requests concurrently."""

        request_list = list(requests)
        if not request_list:
            return []

        limiter: Optional[asyncio.Semaphore] = None
        if concurrency and concurrency > 0:
            limiter = asyncio.Semaphore(concurrency)

        async def _invoke(entry: Mapping[str, Any]) -> bytes:
            if "body" not in entry or "modelId" not in entry:
                raise ValueError("Each request must include 'body' and 'modelId' keys")  # noqa: E501

            body_value = entry["body"]
            model_id = entry["modelId"]
            extra = {
                key: value
                for key, value in entry.items()
                if key not in {"body", "modelId"}
            }

            if limiter is not None:
                async with limiter:
                    return await self.invoke_model(
                        body=body_value, modelId=model_id, **extra
                    )

            return await self.invoke_model(
                body=body_value,
                modelId=model_id,
                **extra,
            )

        tasks = [asyncio.create_task(_invoke(item)) for item in request_list]
        try:
            if return_exceptions:
                gathered = await asyncio.gather(*tasks, return_exceptions=True)
                results: Sequence[Union[bytes, BaseException]] = tuple(gathered)  # noqa: E501
            else:
                gathered = await asyncio.gather(*tasks, return_exceptions=False)  # noqa: E501
                results = tuple(gathered)
        finally:
            for task in tasks:
                if not task.done():
                    task.cancel()

        return results

    def _process_event_message(
        self,
        message,
    ) -> Union[Dict[str, Any], bytes, None]:
        """Process an individual event message"""
        try:
            # Handle EventStreamMessage objects
            if hasattr(message, "headers") and hasattr(message, "payload"):
                headers = getattr(message, "headers", {})
                payload = getattr(message, "payload", b"")
            # Handle dict-like objects from manual parsing
            elif isinstance(message, dict):
                headers = message.get("headers", {})
                payload = message.get("payload", b"")
            else:
                log.warning(f"Unexpected message type: {type(message)}")
                return None

            headers_dict = self._normalize_headers(headers)

            message_type = headers_dict.get(":message-type")
            exception_type = headers_dict.get(":exception-type")
            event_type = headers_dict.get(":event-type")

            if isinstance(payload, memoryview):
                payload = payload.tobytes()
            elif isinstance(payload, bytearray):
                payload = bytes(payload)

            # Get content type from headers
            content_type = "application/json"  # Default
            content_type = headers_dict.get(":content-type", content_type)

            if message_type in {"exception", "error"} or (
                message_type == "event"
                and event_type
                in {
                    "error",
                    "exception",
                    "modelInvocationError",
                }
            ):
                detail = self._decode_payload_text(payload)
                raise BedrockStreamError(
                    message_type=message_type or "event",
                    exception_type=exception_type or event_type or "UnknownException",  # noqa: E501
                    detail=detail,
                )

            if not payload:
                return {}

            # Handle JSON content
            if "application/json" in content_type:
                try:
                    payload_data = orjson.loads(
                        self._ensure_text_payload(
                            payload=payload,
                        )
                    )

                    # Extract base64-encoded bytes if present
                    if "bytes" in payload_data:
                        return base64.b64decode(payload_data["bytes"])
                    elif (
                        "chunk" in payload_data
                        and isinstance(payload_data["chunk"], dict)
                        and "bytes" in payload_data["chunk"]
                    ):
                        return base64.b64decode(payload_data["chunk"]["bytes"])
                    else:
                        return payload_data

                except (UnicodeDecodeError, orjson.JSONDecodeError) as e:
                    log.error(f"Failed to parse JSON payload: {e}")
                    return payload
            else:
                # Return raw bytes for non-JSON content
                return payload

        except Exception as e:
            log.error(f"Error processing event message: {e}")
            return None

    def _normalize_headers(self, headers) -> Dict[str, Any]:
        """Convert botocore header objects into plain Python values."""

        if headers is None:
            return {}

        if isinstance(headers, dict):
            items = headers.items()
        elif hasattr(headers, "items"):
            items = headers.items()
        else:
            try:
                items = list(headers)
            except TypeError:
                return {}

        normalized: Dict[str, Any] = {}
        for key, value in items:
            if hasattr(value, "value"):
                normalized[key] = value.value
            else:
                normalized[key] = value

        return normalized

    def _decode_payload_text(self, payload: Union[bytes, str]) -> str:
        if isinstance(payload, bytes):
            return payload.decode("utf-8", errors="replace")
        if isinstance(payload, memoryview):
            return payload.tobytes().decode("utf-8", errors="replace")
        if isinstance(payload, bytearray):
            return bytes(payload).decode("utf-8", errors="replace")
        return str(payload)

    def _ensure_text_payload(self, payload: Union[str, bytes]) -> str:
        if isinstance(payload, str):
            return payload
        if isinstance(payload, bytes):
            return payload.decode("utf-8")
        if isinstance(payload, memoryview):
            return payload.tobytes().decode("utf-8")
        if isinstance(payload, bytearray):
            return bytes(payload).decode("utf-8")
        if hasattr(payload, "read"):
            return self._ensure_text_payload(payload.read())
        return str(payload)

    async def _handle_error_response(self, response: aiohttp.ClientResponse):
        """Handle HTTP error responses"""
        if response.status == 200:
            return

        error_text = await response.text()
        error_map = {
            403: "AccessDeniedException",
            408: "ModelTimeoutException",
            424: "ModelErrorException",
            429: "ThrottlingException",
            500: "InternalServerException",
        }

        error_type = error_map.get(response.status, "UnknownException")
        raise BedrockClientError(response.status, error_type, error_text)

    def _signed_request(
        self,
        credentials,
        url: str,
        method: str,
        body: str,
        region_name: str,
        **kwargs,
    ) -> Dict[str, str]:
        """Create a signed AWS request"""
        if credentials is None:
            raise RuntimeError("AWS credentials are not initialized")

        if hasattr(credentials, "get_frozen_credentials"):
            credentials = credentials.get_frozen_credentials()

        request = AWSRequest(method=method, url=url, data=body)
        request.headers.add_header("Host", url.split("/")[2])

        # Set appropriate headers based on the endpoint
        if "invoke-with-response-stream" in url:
            request.headers.add_header(
                "Accept",
                "application/vnd.amazon.eventstream",
            )
        else:
            request.headers.add_header(
                "Accept", kwargs.get("accept", "application/json")
            )

        request.headers.add_header(
            "Content-Type", kwargs.get("contentType", "application/json")
        )
        request.headers.add_header(
            "X-Amzn-Bedrock-Trace", kwargs.get("trace", "DISABLED")
        )

        # Optional headers
        optional_headers = [
            (
                "guardrailIdentifier",
                "X-Amzn-Bedrock-GuardrailIdentifier",
            ),
            (
                "guardrailVersion",
                "X-Amzn-Bedrock-GuardrailVersion",
            ),
            (
                "performanceConfigLatency",
                "X-Amzn-Bedrock-PerformanceConfig-Latency",
            ),
        ]

        for kwarg_key, header_name in optional_headers:
            if kwargs.get(kwarg_key):
                request.headers.add_header(
                    header_name,
                    kwargs.get(kwarg_key),  # type: ignore
                )

        # Sign the request
        SigV4Auth(credentials, "bedrock", region_name).add_auth(request)
        return dict(request.headers)
