"""Bedrock Runtime client wrapper with automatic retry logic."""

from __future__ import annotations

from typing import Any

import boto3

from alice_cli.errors import AliceCLIError, GatewayError, QuotaExceededError
from alice_cli.logging import get_logger
from alice_cli.retry import with_bedrock_retry

log = get_logger(__name__)


class BedrockClient:
    """Wrapper for the boto3 bedrock-runtime client.

    Adds automatic retry logic for transient errors (throttling, service
    unavailability) using the existing ``with_bedrock_retry`` decorator from
    ``alice_cli.retry``.

    Implements the ``BedrockRuntimeClient`` protocol so it can be used
    anywhere that protocol is expected.

    Example::

        session = boto3.Session(profile_name="drcc-ai", region_name="us-east-1")
        client = BedrockClient.from_session(session)
        response = client.converse(modelId="...", messages=[...])
    """

    def __init__(self, client: Any) -> None:
        self._client = client

    @classmethod
    def from_session(cls, session: boto3.Session) -> BedrockClient:
        """Create a BedrockClient from an existing boto3 Session."""
        return cls(session.client("bedrock-runtime"))

    @classmethod
    def from_session_via_gateway(
        cls,
        session: boto3.Session,
        gateway_url: str,
    ) -> BedrockClient:
        """Create a BedrockClient that routes through the Cheshire Gate gateway.

        The gateway uses ``authorizerType=NONE`` and relies on AWS SigV4 request
        signing for transport-level authentication. No extra headers are required.
        The boto3 client is pointed at *gateway_url* (the MCP endpoint) instead of
        the default Bedrock runtime endpoint.
        """
        client = session.client("bedrock-runtime", endpoint_url=gateway_url)
        return cls(client)

    @classmethod
    def from_credentials(
        cls,
        access_key: str,
        secret_key: str,
        region: str,
    ) -> BedrockClient:
        """Create a BedrockClient from API key credentials."""
        session = boto3.Session(
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name=region,
        )
        return cls.from_session(session)

    @with_bedrock_retry
    def converse(
        self,
        *,
        modelId: str,
        messages: list[dict[str, Any]],
        inferenceConfig: dict[str, Any] | None = None,
        system: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Invoke the Converse API with automatic retry on transient errors.

        Args:
            modelId: Model ID or inference profile ARN.
            messages: List of message dicts with ``role`` and ``content``.
            inferenceConfig: Optional inference configuration (e.g. maxTokens).
            system: Optional system prompt list.
            **kwargs: Additional parameters forwarded to boto3.

        Returns:
            Response dict with ``output`` and ``usage`` keys.

        Raises:
            ClientError: If the API call fails after exhausting retries.
            BotoCoreError: On low-level connection errors after retries.
        """
        call_kwargs: dict[str, Any] = {"modelId": modelId, "messages": messages}

        if inferenceConfig is not None:
            call_kwargs["inferenceConfig"] = inferenceConfig
        if system is not None:
            call_kwargs["system"] = system
        call_kwargs.update(kwargs)

        log.debug(
            "bedrock_converse",
            model_id=modelId,
            message_count=len(messages),
            has_system=system is not None,
            max_tokens=(inferenceConfig or {}).get("maxTokens"),
        )

        from botocore.exceptions import ClientError as _ClientError

        from alice_cli import personality as _personality
        from alice_cli.errors import GatewayError, QuotaExceededError

        try:
            response = self._client.converse(**call_kwargs)
        except _ClientError as exc:
            status = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode", 0)
            msg = exc.response.get("Error", {}).get("Message", str(exc))
            if status == 429:
                raise QuotaExceededError(_personality.error_quota_exceeded(msg)) from exc
            if status == 403:
                raise GatewayError(msg) from exc
            raise

        usage = response.get("usage", {})
        log.debug(
            "bedrock_converse_success",
            model_id=modelId,
            input_tokens=usage.get("inputTokens", 0),
            output_tokens=usage.get("outputTokens", 0),
        )

        return response  # type: ignore[no-any-return]


def _translate_mcp_response(mcp_result: dict[str, Any]) -> dict[str, Any]:
    """Convert an MCP JSON-RPC 2.0 result to a Bedrock Converse response shape.

    Mapping rules:
    - ``result.content[].text`` → ``output.message.content[].text``
    - ``result.usage.inputTokens`` → ``usage.inputTokens``
    - ``result.usage.outputTokens`` → ``usage.outputTokens``
    - Missing ``usage`` → ``{inputTokens: 0, outputTokens: 0}``
    - Missing ``content`` → empty list
    - ``output.message.role`` is always ``"assistant"``
    """
    result = mcp_result.get("result", {})
    mcp_content = result.get("content", [])
    mcp_usage = result.get("usage", {})

    content_blocks: list[dict[str, Any]] = [
        {"text": item["text"]} for item in mcp_content if item.get("type") == "text"
    ]

    return {
        "output": {
            "message": {
                "role": "assistant",
                "content": content_blocks,
            }
        },
        "usage": {
            "inputTokens": mcp_usage.get("inputTokens", 0),
            "outputTokens": mcp_usage.get("outputTokens", 0),
        },
    }


class GatewayBedrockClient:
    """MCP JSON-RPC 2.0 client for the Cheshire Gate AgentCore Gateway.

    Implements the ``BedrockRuntimeClient`` protocol so it can be used as a
    drop-in replacement for ``BedrockClient`` anywhere in alice-cli.

    The client constructs JSON-RPC 2.0 ``tools/call`` requests, signs them
    with SigV4 (``service=bedrock-agentcore``), POSTs to the gateway URL,
    and translates MCP responses back into the Bedrock Converse response shape.
    """

    def __init__(
        self,
        gateway_url: str,
        session: boto3.Session,
        model_to_tool: dict[str, str] | None = None,
    ) -> None:
        self._gateway_url = gateway_url
        self._session = session
        self._region = session.region_name or "us-east-1"
        # Reverse map: base_model_id → MCP tool name.
        # Populated from the /cheshire-gate/gateway-targets SSM parameter.
        # None means "not yet loaded" (lazy); empty dict means "no targets".
        self._model_to_tool: dict[str, str] = model_to_tool or {}

    @staticmethod
    def _base_model_id(model_id: str) -> str:
        """Strip Bedrock list-models context-window suffix (e.g. ``:300k``)."""
        import re

        return re.sub(r":\d+[kKmMgG]$", "", model_id)

    def _tool_name_for(self, model_id: str) -> str | None:
        """Return the MCP tool name for *model_id*, or ``None`` if no target exists."""
        base = self._base_model_id(model_id)
        return self._model_to_tool.get(base) or self._model_to_tool.get(model_id)

    def _build_jsonrpc_request(self, tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        """Construct a JSON-RPC 2.0 ``tools/call`` request body."""
        import uuid

        return {
            "jsonrpc": "2.0",
            "id": str(uuid.uuid4()),
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments,
            },
        }

    def _sign_request(self, url: str, body: bytes) -> dict[str, str]:
        """Sign an HTTPS POST with SigV4 for ``bedrock-agentcore``."""
        from botocore.auth import SigV4Auth
        from botocore.awsrequest import AWSRequest

        credentials = self._session.get_credentials()
        if credentials is None:
            msg = "No AWS credentials available for SigV4 signing"
            raise AliceCLIError(msg)
        frozen = credentials.get_frozen_credentials()

        aws_request = AWSRequest(
            method="POST",
            url=url,
            data=body,
            headers={"Content-Type": "application/json"},
        )
        SigV4Auth(frozen, "bedrock-agentcore", self._region).add_auth(aws_request)
        return dict(aws_request.headers)

    @with_bedrock_retry
    def converse(
        self,
        *,
        modelId: str,
        messages: list[dict[str, Any]],
        inferenceConfig: dict[str, Any] | None = None,
        system: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Send a Converse request through the Cheshire Gate MCP gateway.

        Constructs a JSON-RPC 2.0 ``tools/call`` request, signs it with SigV4,
        POSTs to the gateway, and translates the MCP response back into the
        Bedrock Converse response shape.
        """
        import json as _json

        import urllib3

        arguments: dict[str, Any] = {"modelId": modelId, "messages": messages}
        if inferenceConfig is not None:
            arguments["inferenceConfig"] = inferenceConfig
        if system is not None:
            arguments["system"] = system
        arguments.update(kwargs)

        tool_name = self._tool_name_for(modelId)
        if tool_name is None:
            # No gateway target for this model — fall back to direct Bedrock.
            log.debug("gateway_no_target_fallback", model_id=modelId)
            from botocore.exceptions import BotoCoreError, ClientError

            direct = BedrockClient.from_session(self._session)
            try:
                return direct.converse(
                    modelId=modelId,
                    messages=messages,
                    inferenceConfig=inferenceConfig,
                    system=system,
                    **kwargs,
                )
            except (ClientError, BotoCoreError):
                raise

        jsonrpc_body = self._build_jsonrpc_request(tool_name, arguments)
        body_bytes = _json.dumps(jsonrpc_body).encode()

        log.debug(
            "gateway_converse",
            gateway_url=self._gateway_url,
            tool_name=tool_name,
            model_id=modelId,
        )

        signed_headers = self._sign_request(self._gateway_url, body_bytes)

        http = urllib3.PoolManager()
        resp = http.request(
            "POST",
            self._gateway_url,
            body=body_bytes,
            headers=signed_headers,
        )

        # --- HTTP error handling ---
        if resp.status == 429:
            raise QuotaExceededError(
                f"Gateway quota exceeded: {resp.data.decode(errors='replace')}"
            )
        if resp.status == 403:
            raise GatewayError(f"Gateway access denied: {resp.data.decode(errors='replace')}")
        if resp.status < 200 or resp.status >= 300:
            raise AliceCLIError(
                f"Gateway request failed (HTTP {resp.status}): "
                f"{resp.data.decode(errors='replace')}"
            )

        # --- Parse JSON-RPC response ---
        result = _json.loads(resp.data.decode())

        if "error" in result:
            error_msg = result["error"].get("message", "Unknown gateway error")
            raise GatewayError(f"Gateway JSON-RPC error: {error_msg}")

        return _translate_mcp_response(result)
