"""
AWS Bedrock provider integration for Amplitude AI SDK.

Wraps the ``boto3`` Bedrock Runtime client to automatically track LLM
usage in Amplitude.  Uses ``invoke_model`` / ``invoke_model_with_response_stream``
under the hood.

Usage::

    import boto3
    from amplitude import Amplitude
    from amplitude_ai.providers.bedrock import Bedrock

    amplitude = Amplitude("YOUR_API_KEY")
    bedrock = Bedrock(
        amplitude=amplitude,
        region_name="us-east-1",
    )

    response = bedrock.converse(
        modelId="anthropic.claude-3-sonnet-20240229-v1:0",
        messages=[{"role": "user", "content": [{"text": "Hello!"}]}],
        amplitude_user_id="user-1",
    )
"""
from __future__ import annotations

import json
import time
import uuid
from typing import Any, Dict, Iterator, List, Optional

try:
    import boto3
    from botocore.config import Config as BotoConfig

    BEDROCK_AVAILABLE = True
except ImportError:
    BEDROCK_AVAILABLE = False

from amplitude import Amplitude

from ..core.privacy import PrivacyConfig
from ..core.tracking import track_ai_message, track_user_message
from ..exceptions import ProviderError
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger
from .base import _apply_session_context, AmplitudeOrAI, resolve_amplitude


class Bedrock:
    """
    AWS Bedrock wrapper that automatically tracks usage in Amplitude.

    Wraps the ``bedrock-runtime`` ``converse`` and ``converse_stream`` APIs.
    The wrapper is intentionally thin -- it delegates to boto3 for all AWS

    .. todo:: Refactor to inherit from :class:`BaseAIProvider` for shared
       tracking logic (privacy config, metadata helpers, etc.).
    communication and only adds Amplitude event tracking on top.
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        privacy_config: Optional[PrivacyConfig] = None,
        *,
        region_name: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_secret_access_key: Optional[str] = None,
        aws_session_token: Optional[str] = None,
        boto_config: Optional[Any] = None,
        client: Optional[Any] = None,
    ) -> None:
        """
        Initialise the Bedrock provider.

        You can either pass an existing ``bedrock-runtime`` client via
        ``client=`` or let the wrapper create one from the provided
        credentials / region.

        Args:
            amplitude: ``Amplitude`` client or ``AmplitudeAI`` instance.
            privacy_config: Privacy configuration for content sanitisation.
            region_name: AWS region (e.g., ``"us-east-1"``).
            aws_access_key_id: Explicit AWS access key (optional).
            aws_secret_access_key: Explicit AWS secret key (optional).
            aws_session_token: Explicit AWS session token (optional).
            boto_config: ``botocore.config.Config`` overrides.
            client: Pre-configured ``boto3`` bedrock-runtime client.
        """
        if not BEDROCK_AVAILABLE:
            raise ProviderError(
                "boto3 is not installed. Install with: pip install boto3"
            )

        self.amplitude = resolve_amplitude(amplitude)
        self.privacy_config = privacy_config or PrivacyConfig()

        if client is not None:
            self._client = client
        else:
            session_kwargs: Dict[str, Any] = {}
            if region_name:
                session_kwargs["region_name"] = region_name
            if aws_access_key_id:
                session_kwargs["aws_access_key_id"] = aws_access_key_id
            if aws_secret_access_key:
                session_kwargs["aws_secret_access_key"] = aws_secret_access_key
            if aws_session_token:
                session_kwargs["aws_session_token"] = aws_session_token

            client_kwargs: Dict[str, Any] = {}
            if boto_config is not None:
                client_kwargs["config"] = boto_config

            session = boto3.Session(**session_kwargs)
            self._client = session.client("bedrock-runtime", **client_kwargs)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def converse(
        self,
        *,
        modelId: str,
        messages: List[Dict[str, Any]],
        amplitude_user_id: Optional[str] = None,
        amplitude_session_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Call Bedrock ``Converse`` with automatic tracking.

        All ``amplitude_*`` parameters are stripped before forwarding to AWS.
        Any additional ``**kwargs`` (``system``, ``inferenceConfig``,
        ``toolConfig``, etc.) are passed through verbatim.

        Returns the raw Bedrock ``Converse`` response dict.
        """
        # Auto-inherit session/agent context when inside agent.session()
        (
            amplitude_user_id, amplitude_session_id,
            amplitude_trace_id, amplitude_turn_id,
            amplitude_event_properties, amplitude_groups,
        ) = _apply_session_context(
            user_id=amplitude_user_id,
            session_id=amplitude_session_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            groups=amplitude_groups,
        )

        trace_id = amplitude_trace_id or str(uuid.uuid4())
        session_id = amplitude_session_id or str(uuid.uuid4())

        start_time = time.time()
        is_error = False
        error_message: Optional[str] = None
        response: Optional[Dict[str, Any]] = None

        try:
            response = self._client.converse(
                modelId=modelId,
                messages=messages,
                **kwargs,
            )
        except Exception as exc:
            is_error = True
            error_message = str(exc)
            raise
        finally:
            latency_ms = (time.time() - start_time) * 1000
            if amplitude_user_id:
                self._track(
                    user_id=amplitude_user_id,
                    model_id=modelId,
                    messages=messages,
                    response=response,
                    latency_ms=latency_ms,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=is_error,
                    error_message=error_message,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    aws_kwargs=kwargs,
                )

        return response  # type: ignore[return-value]

    def converse_stream(
        self,
        *,
        modelId: str,
        messages: List[Dict[str, Any]],
        amplitude_user_id: Optional[str] = None,
        amplitude_session_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Iterator[Dict[str, Any]]:
        """Call Bedrock ``ConverseStream`` with automatic tracking.

        Yields the raw ``stream`` events.  Tracking is emitted once the
        stream is fully consumed (or errors).
        """
        # Auto-inherit session/agent context when inside agent.session()
        (
            amplitude_user_id, amplitude_session_id,
            amplitude_trace_id, amplitude_turn_id,
            amplitude_event_properties, amplitude_groups,
        ) = _apply_session_context(
            user_id=amplitude_user_id,
            session_id=amplitude_session_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            groups=amplitude_groups,
        )

        trace_id = amplitude_trace_id or str(uuid.uuid4())
        session_id = amplitude_session_id or str(uuid.uuid4())
        start_time = time.time()
        first_byte_time: Optional[float] = None

        content_chunks: List[str] = []
        reasoning_chunks: List[str] = []
        streaming_tool_calls: List[Dict[str, Any]] = []
        input_tokens: Optional[int] = None
        output_tokens: Optional[int] = None
        cache_read_tokens: Optional[int] = None
        cache_creation_tokens: Optional[int] = None
        finish_reason: Optional[str] = None
        is_error = False
        error_message: Optional[str] = None

        try:
            raw = self._client.converse_stream(
                modelId=modelId,
                messages=messages,
                **kwargs,
            )
        except Exception as exc:
            is_error = True
            error_message = str(exc)
            # Track the error and re-raise
            if amplitude_user_id:
                self._track(
                    user_id=amplitude_user_id,
                    model_id=modelId,
                    messages=messages,
                    response=None,
                    latency_ms=(time.time() - start_time) * 1000,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=error_message,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    aws_kwargs={**kwargs, "_is_streaming": True},
                )
            raise

        try:
            for event in raw.get("stream", []):
                try:
                    if "contentBlockStart" in event:
                        start_block = event["contentBlockStart"].get("start", {})
                        tool_use = start_block.get("toolUse")
                        if tool_use:
                            streaming_tool_calls.append({
                                "type": "function",
                                "id": tool_use.get("toolUseId", ""),
                                "function": {
                                    "name": tool_use.get("name", ""),
                                    "arguments": "",
                                },
                            })

                    elif "contentBlockDelta" in event:
                        if first_byte_time is None:
                            first_byte_time = time.time()
                        delta = event["contentBlockDelta"].get("delta", {})
                        text = delta.get("text", "")
                        if text:
                            content_chunks.append(text)
                        reasoning_text = delta.get("reasoningText", "")
                        if reasoning_text:
                            reasoning_chunks.append(reasoning_text)
                        tool_input = delta.get("toolUse", {}).get("input", "")
                        if tool_input and streaming_tool_calls:
                            streaming_tool_calls[-1]["function"]["arguments"] += str(tool_input)

                    elif "metadata" in event:
                        usage = event["metadata"].get("usage", {})
                        input_tokens = usage.get("inputTokens")
                        output_tokens = usage.get("outputTokens")
                        cache_read_tokens = usage.get("cacheReadInputTokenCount")
                        cache_creation_tokens = usage.get("cacheWriteInputTokenCount")

                    elif "messageStop" in event:
                        finish_reason = event["messageStop"].get("stopReason")
                except Exception as tracking_exc:
                    get_logger(self.amplitude).warning(
                        f"Tracking failed for chunk: {tracking_exc}"
                    )

                yield event

        except Exception as exc:
            is_error = True
            error_message = str(exc)
            raise
        finally:
            if amplitude_user_id:
                latency_ms = (time.time() - start_time) * 1000
                ttfb_ms = (
                    (first_byte_time - start_time) * 1000
                    if first_byte_time
                    else None
                )
                full_content = "".join(content_chunks)
                full_reasoning = "".join(reasoning_chunks) if reasoning_chunks else None

                total_cost_usd = None
                if input_tokens is not None and output_tokens is not None:
                    try:
                        total_cost_usd = calculate_cost(
                            model_name=modelId,
                            input_tokens=input_tokens,
                            output_tokens=output_tokens,
                            cache_read_input_tokens=cache_read_tokens or 0,
                            cache_creation_input_tokens=cache_creation_tokens or 0,
                            default_provider='bedrock',
                        )
                    except Exception:
                        pass

                # Synthesise a response-like dict for the tracker
                synth_content: List[Dict[str, Any]] = [{"text": full_content}]
                if full_reasoning:
                    synth_content.append(
                        {"reasoningContent": {"reasoningText": {"text": full_reasoning}}}
                    )
                for tc in streaming_tool_calls:
                    synth_content.append({
                        "toolUse": {
                            "toolUseId": tc.get("id", ""),
                            "name": tc["function"]["name"],
                            "input": tc["function"]["arguments"],
                        }
                    })
                synth_response: Dict[str, Any] = {
                    "output": {
                        "message": {
                            "role": "assistant",
                            "content": synth_content,
                        }
                    },
                    "stopReason": finish_reason,
                }
                if input_tokens is not None or output_tokens is not None:
                    synth_response["usage"] = {
                        "inputTokens": input_tokens,
                        "outputTokens": output_tokens,
                    }
                    if cache_read_tokens is not None:
                        synth_response["usage"]["cacheReadInputTokenCount"] = cache_read_tokens
                    if cache_creation_tokens is not None:
                        synth_response["usage"]["cacheWriteInputTokenCount"] = cache_creation_tokens

                self._track(
                    user_id=amplitude_user_id,
                    model_id=modelId,
                    messages=messages,
                    response=synth_response,
                    latency_ms=latency_ms,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=is_error,
                    error_message=error_message,
                    total_cost_usd=total_cost_usd,
                    provider_ttfb_ms=ttfb_ms,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    aws_kwargs={**kwargs, "_is_streaming": True},
                )

    # ------------------------------------------------------------------
    # Passthrough to underlying client
    # ------------------------------------------------------------------

    @property
    def client(self) -> Any:
        """Access the underlying ``boto3`` bedrock-runtime client."""
        return self._client

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_content_text(messages: List[Dict[str, Any]]) -> str:
        """Pull the last user message text out of Bedrock ``messages`` format."""
        for msg in reversed(messages):
            if msg.get("role") == "user":
                for block in msg.get("content", []):
                    if isinstance(block, dict) and "text" in block:
                        return block["text"]
        return ""

    @staticmethod
    def _extract_response_text(response: Optional[Dict[str, Any]]) -> str:
        """Pull assistant text from a Bedrock ``Converse`` response dict."""
        if response is None:
            return ""
        output = response.get("output", {})
        message = output.get("message", {})
        for block in message.get("content", []):
            if isinstance(block, dict) and "text" in block:
                return block["text"]
        return ""

    @staticmethod
    def _extract_reasoning_content(response: Optional[Dict[str, Any]]) -> Optional[str]:
        """Extract reasoning/thinking content from a Bedrock Converse response.

        Bedrock models that support reasoning include
        ``reasoningContent.reasoningText.text`` blocks in the response
        ``output.message.content`` array.
        """
        if response is None:
            return None
        message = response.get("output", {}).get("message", {})
        parts: List[str] = []
        for block in message.get("content", []):
            if not isinstance(block, dict):
                continue
            rc = block.get("reasoningContent")
            if rc is None:
                continue
            rt = rc.get("reasoningText") if isinstance(rc, dict) else None
            if isinstance(rt, dict) and "text" in rt:
                parts.append(rt["text"])
        return "\n".join(parts) if parts else None

    @staticmethod
    def _extract_usage(response: Optional[Dict[str, Any]]) -> Dict[str, Optional[int]]:
        """Extract token counts from a Bedrock response.

        Bedrock's ``converse`` API includes ``cacheReadInputTokenCount`` and
        ``cacheWriteInputTokenCount`` in the usage object when prompt caching
        is active.  We map these to the canonical SDK names so that
        ``calculate_cost`` can apply cache-aware pricing.
        """
        if response is None:
            return {
                "input_tokens": None,
                "output_tokens": None,
                "total_tokens": None,
                "cache_read_input_tokens": None,
                "cache_creation_input_tokens": None,
            }
        usage = response.get("usage", {})
        inp = usage.get("inputTokens")
        out = usage.get("outputTokens")
        total = inp + out if inp is not None and out is not None else None
        cache_read = usage.get("cacheReadInputTokenCount")
        cache_creation = usage.get("cacheWriteInputTokenCount")
        return {
            "input_tokens": inp,
            "output_tokens": out,
            "total_tokens": total,
            "cache_read_input_tokens": cache_read,
            "cache_creation_input_tokens": cache_creation,
        }

    @staticmethod
    def _extract_finish_reason(response: Optional[Dict[str, Any]]) -> Optional[str]:
        if response is None:
            return None
        reason = response.get("stopReason")
        if reason == "end_turn":
            return "stop"
        if reason == "max_tokens":
            return "length"
        return reason

    @staticmethod
    def _extract_tool_calls(response: Optional[Dict[str, Any]]) -> Optional[List[Dict[str, Any]]]:
        """Extract tool calls from a Bedrock Converse response.

        Bedrock includes tool use blocks in ``output.message.content``
        with a ``toolUse`` key containing ``toolUseId``, ``name``, and
        ``input``.
        """
        if response is None:
            return None
        message = response.get("output", {}).get("message", {})
        tool_calls: List[Dict[str, Any]] = []
        for block in message.get("content", []):
            if not isinstance(block, dict):
                continue
            tool_use = block.get("toolUse")
            if tool_use is None:
                continue
            tool_calls.append({
                "type": "function",
                "id": tool_use.get("toolUseId", ""),
                "function": {
                    "name": tool_use.get("name", ""),
                    "arguments": str(tool_use.get("input", {})),
                },
            })
        return tool_calls if tool_calls else None

    def _extract_system_prompt(self, aws_kwargs: Dict[str, Any]) -> Optional[str]:
        """Extract system prompt from Bedrock Converse ``system`` parameter."""
        system = aws_kwargs.get("system")
        if not system:
            return None
        if isinstance(system, list):
            parts = []
            for block in system:
                if isinstance(block, dict) and "text" in block:
                    parts.append(block["text"])
            return "".join(parts) if parts else None
        if isinstance(system, str):
            return system
        return str(system)

    @staticmethod
    def _extract_tool_definitions(aws_kwargs: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:
        """Extract tool definitions from Bedrock ``toolConfig``."""
        tool_config = aws_kwargs.get("toolConfig")
        if isinstance(tool_config, dict):
            tools = tool_config.get("tools")
            if tools and isinstance(tools, list):
                return tools
        return None

    @staticmethod
    def _extract_model_params(aws_kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Extract model parameters from Bedrock ``inferenceConfig``."""
        result: Dict[str, Any] = {
            "temperature": None,
            "top_p": None,
            "max_output_tokens": None,
        }
        ic = aws_kwargs.get("inferenceConfig")
        if isinstance(ic, dict):
            result["temperature"] = ic.get("temperature")
            result["top_p"] = ic.get("topP")
            result["max_output_tokens"] = ic.get("maxTokens")
        return result

    def _track(
        self,
        *,
        user_id: str,
        model_id: str,
        messages: List[Dict[str, Any]],
        response: Optional[Dict[str, Any]],
        latency_ms: float,
        session_id: Optional[str],
        trace_id: Optional[str],
        turn_id: Optional[int],
        is_error: bool = False,
        error_message: Optional[str] = None,
        total_cost_usd: Optional[float] = None,
        provider_ttfb_ms: Optional[float] = None,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
        aws_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Emit Amplitude tracking events."""
        try:
            current_turn = turn_id or 1

            # Extract agent-level fields that _apply_session_context may
            # have injected into event_properties.
            _ep = dict(event_properties) if event_properties else {}
            _ctx_agent_id = _ep.pop("agent_id", None)
            _ctx_parent_agent_id = _ep.pop("parent_agent_id", None)
            _ctx_env = _ep.pop("env", None)
            _ctx_customer_org_id = _ep.pop("customer_org_id", None)
            _ctx_agent_version = _ep.pop("agent_version", None)
            _ctx_context = _ep.pop("context", None)
            if not _ep and event_properties is not None:
                event_properties = None
            else:
                event_properties = _ep or event_properties

            # Track user messages
            user_text = self._extract_content_text(messages)
            if user_text:
                track_user_message(
                    amplitude=self.amplitude,
                    user_id=user_id,
                    message_content=user_text,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=current_turn,
                    agent_id=_ctx_agent_id,
                    parent_agent_id=_ctx_parent_agent_id,
                    env=_ctx_env,
                    customer_org_id=_ctx_customer_org_id,
                    agent_version=_ctx_agent_version,
                    context=_ctx_context,
                    message_source="agent" if _ctx_parent_agent_id else "user",
                    event_properties=event_properties,
                    user_properties=user_properties,
                    groups=groups,
                    privacy_config=self.privacy_config,
                )
                current_turn += 1

            # Extract response data
            response_content = self._extract_response_text(response) if not is_error else ""
            usage = self._extract_usage(response) if not is_error else {}
            finish_reason = self._extract_finish_reason(response) if not is_error else None
            tool_calls = self._extract_tool_calls(response) if not is_error else None
            reasoning_content = (
                self._extract_reasoning_content(response) if not is_error else None
            )

            # Extract model configuration from the original API kwargs
            _kw = aws_kwargs or {}
            system_prompt = self._extract_system_prompt(_kw)
            model_params = self._extract_model_params(_kw)
            is_streaming = _kw.get("_is_streaming")
            tool_definitions = self._extract_tool_definitions(_kw)

            if total_cost_usd is None and not is_error:
                inp = usage.get("input_tokens")
                out = usage.get("output_tokens")
                if inp is not None and out is not None:
                    try:
                        total_cost_usd = calculate_cost(
                            model_name=model_id,
                            input_tokens=inp,
                            output_tokens=out,
                            cache_read_input_tokens=usage.get("cache_read_input_tokens") or 0,
                            cache_creation_input_tokens=usage.get("cache_creation_input_tokens") or 0,
                            default_provider='bedrock',
                        )
                    except Exception:
                        pass

            track_ai_message(
                amplitude=self.amplitude,
                user_id=user_id,
                model_name=model_id,
                provider="bedrock",
                response_content=response_content,
                latency_ms=latency_ms,
                session_id=session_id,
                trace_id=trace_id,
                turn_id=current_turn,
                agent_id=_ctx_agent_id,
                parent_agent_id=_ctx_parent_agent_id,
                env=_ctx_env,
                customer_org_id=_ctx_customer_org_id,
                agent_version=_ctx_agent_version,
                context=_ctx_context,
                provider_ttfb_ms=provider_ttfb_ms,
                is_error=is_error,
                error_message=error_message,
                finish_reason=finish_reason,
                tool_calls=tool_calls,
                total_cost_usd=total_cost_usd,
                reasoning_content=reasoning_content,
                tool_definitions=tool_definitions,
                system_prompt=system_prompt,
                temperature=model_params.get("temperature"),
                max_output_tokens=model_params.get("max_output_tokens"),
                top_p=model_params.get("top_p"),
                is_streaming=is_streaming,
                event_properties=event_properties,
                user_properties=user_properties,
                groups=groups,
                privacy_config=self.privacy_config,
                input_tokens=usage.get("input_tokens"),
                output_tokens=usage.get("output_tokens"),
                total_tokens=usage.get("total_tokens"),
                cache_read_input_tokens=usage.get("cache_read_input_tokens"),
                cache_creation_input_tokens=usage.get("cache_creation_input_tokens"),
            )
        except Exception as exc:
            get_logger(self.amplitude).error(f"Failed to track Bedrock completion: {exc}")
