"""
AWS Bedrock integration — wraps BedrockRuntime client to auto-trace.

Usage:
    import boto3
    from syntropy import Syntropy
    from syntropy.integrations.bedrock import wrap_bedrock

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    bedrock = boto3.client("bedrock-runtime")
    bedrock = wrap_bedrock(bedrock, syntropy=syn)

    # Auto-traced via Converse API:
    response = bedrock.converse(
        modelId="anthropic.claude-3-sonnet-20240229-v1:0",
        messages=[{"role": "user", "content": [{"text": "Hello!"}]}],
    )
"""

from __future__ import annotations

import json
import time
from typing import Any, Optional

from syntropy.client import Syntropy


class TracedBedrockRuntime:
    """Proxy around boto3 bedrock-runtime client."""

    def __init__(self, client: Any, syntropy: Syntropy, agent_id: str):
        self._client = client
        self._syntropy = syntropy
        self._agent_id = agent_id

    def converse(self, **kwargs: Any) -> Any:
        """Traced wrapper for Bedrock Converse API."""
        start = time.perf_counter()
        response = self._client.converse(**kwargs)
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        # Extract prompt
        messages = kwargs.get("messages", [])
        prompt = ""
        if messages:
            content = messages[-1].get("content", [])
            if content and isinstance(content, list):
                texts = [b.get("text", "") for b in content if isinstance(b, dict)]
                prompt = " ".join(texts)
            elif isinstance(content, str):
                prompt = content

        # Extract response
        output_msg = response.get("output", {}).get("message", {})
        output_content = output_msg.get("content", [])
        output = " ".join(
            b.get("text", "") for b in output_content if isinstance(b, dict)
        )

        # Token usage
        usage = response.get("usage", {})
        tokens_in = usage.get("inputTokens", 0)
        tokens_out = usage.get("outputTokens", 0)

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=kwargs.get("modelId", "bedrock"),
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "aws-bedrock",
                "command": "Converse",
                "stop_reason": response.get("stopReason"),
            },
        )

        return response

    def invoke_model(self, **kwargs: Any) -> Any:
        """Traced wrapper for Bedrock InvokeModel API."""
        start = time.perf_counter()
        response = self._client.invoke_model(**kwargs)
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        model_id = kwargs.get("modelId", "bedrock")

        # Parse request body
        try:
            body = json.loads(kwargs.get("body", "{}"))
        except (json.JSONDecodeError, TypeError):
            body = {}

        # Extract prompt — varies by model provider
        prompt = ""
        if "messages" in body:
            # Anthropic Messages API on Bedrock
            msgs = body["messages"]
            if msgs:
                content = msgs[-1].get("content", "")
                prompt = content if isinstance(content, str) else str(content)
        elif "prompt" in body:
            prompt = body["prompt"]
        elif "inputText" in body:
            prompt = body["inputText"]

        # Parse response body
        resp_body_bytes = response.get("body", b"")
        try:
            if hasattr(resp_body_bytes, "read"):
                resp_body_bytes = resp_body_bytes.read()
            resp_data = json.loads(resp_body_bytes)
        except (json.JSONDecodeError, TypeError):
            resp_data = {}

        # Extract output — varies by model
        output = ""
        if "content" in resp_data and isinstance(resp_data["content"], list):
            # Anthropic
            output = " ".join(
                b.get("text", "") for b in resp_data["content"] if isinstance(b, dict)
            )
        elif "results" in resp_data:
            # Titan
            output = resp_data["results"][0].get("outputText", "")
        elif "generation" in resp_data:
            # Llama
            output = resp_data["generation"]
        elif "completions" in resp_data:
            output = resp_data["completions"][0].get("data", {}).get("text", "")

        # Token usage — best-effort from various response formats
        tokens_in = resp_data.get("usage", {}).get("input_tokens", 0)
        tokens_out = resp_data.get("usage", {}).get("output_tokens", 0)

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=model_id,
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "aws-bedrock",
                "command": "InvokeModel",
            },
        )

        return response

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


def wrap_bedrock(
    client: Any,
    *,
    syntropy: Syntropy,
    agent_id: str = "",
) -> TracedBedrockRuntime:
    """
    Wrap a Bedrock Runtime client to auto-trace converse() and invoke_model() calls.

    Args:
        client: A boto3.client("bedrock-runtime") instance.
        syntropy: A Syntropy client instance.
        agent_id: Override agent_id.

    Returns:
        A drop-in replacement with tracing enabled.
    """
    return TracedBedrockRuntime(
        client, syntropy, agent_id or syntropy._config.agent_id
    )
