"""
AWS Bedrock client wrapper for Raindrop observability.

Wraps a boto3 bedrock-runtime client to automatically capture
converse() and invoke_model() calls and ship them to Raindrop.

Usage:
    from raindrop_bedrock import create_raindrop_bedrock

    raindrop = create_raindrop_bedrock(api_key="rk_...")
    bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
    wrapped = raindrop["wrap"](bedrock)

    response = wrapped.converse(modelId="anthropic.claude-3-5-sonnet-20241022-v2:0", ...)
    raindrop["flush"]()
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Dict, Optional

import raindrop.analytics as raindrop

def create_raindrop_bedrock(
    api_key: str,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> dict:
    """
    Create a Raindrop-instrumented Bedrock wrapper.

    Args:
        api_key: Raindrop API key (rk_...).
        user_id: Associate all events with this user.
        convo_id: Conversation ID to group related events.

    Returns:
        Dict with 'wrap', 'flush', and 'shutdown'.

    Usage:
        raindrop = create_raindrop_bedrock(api_key="rk_...")
        bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
        wrapped = raindrop["wrap"](bedrock)
        response = wrapped.converse(modelId="...", messages=[...])
        raindrop["flush"]()
    """
    raindrop.init(api_key=api_key)

    def wrap(client: Any) -> Any:
        return wrap_bedrock_client(client, user_id=user_id, convo_id=convo_id)

    return {
        "wrap": wrap,
        "flush": raindrop.flush,
        "shutdown": raindrop.shutdown,
    }


def wrap_bedrock_client(
    client: Any,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> Any:
    """Wrap a boto3 bedrock-runtime client to capture LLM calls."""

    original_converse = getattr(client, "converse", None)
    original_invoke_model = getattr(client, "invoke_model", None)

    if original_converse:
        def wrapped_converse(**kwargs: Any) -> Any:
            return _instrument_converse(original_converse, kwargs, user_id, convo_id)
        client.converse = wrapped_converse

    if original_invoke_model:
        def wrapped_invoke_model(**kwargs: Any) -> Any:
            return _instrument_invoke_model(original_invoke_model, kwargs, user_id, convo_id)
        client.invoke_model = wrapped_invoke_model

    return client


def _instrument_converse(
    original_fn: Any,
    kwargs: Dict[str, Any],
    user_id: Optional[str],
    convo_id: Optional[str],
) -> Any:
    try:
        model_id = kwargs.get("modelId")
        input_text = _format_converse_input(kwargs.get("messages"))
        event_id = str(uuid.uuid4())
    except Exception:
        return original_fn(**kwargs)

    try:
        response = original_fn(**kwargs)
    except Exception:
        try:
            raindrop.track_ai(
                user_id=user_id or "unknown",
                event="ai_generation",
                event_id=event_id,
                input=input_text,
                model=model_id,
                convo_id=convo_id,
            )
        except Exception:
            pass
        raise

    try:
        output_text = _extract_converse_output(response)
        usage = response.get("usage", {})
        prompt_tokens = usage.get("inputTokens")
        completion_tokens = usage.get("outputTokens")

        properties: Dict[str, Any] = {}
        if prompt_tokens is not None:
            properties["ai.usage.prompt_tokens"] = prompt_tokens
        if completion_tokens is not None:
            properties["ai.usage.completion_tokens"] = completion_tokens

        raindrop.track_ai(
            user_id=user_id or "unknown",
            event="ai_generation",
            event_id=event_id,
            input=input_text,
            output=output_text,
            model=model_id,
            properties=properties or None,
            convo_id=convo_id,
        )
    except Exception:
        pass

    return response


def _instrument_invoke_model(
    original_fn: Any,
    kwargs: Dict[str, Any],
    user_id: Optional[str],
    convo_id: Optional[str],
) -> Any:
    try:
        model_id = kwargs.get("modelId")
        input_text = _parse_body(kwargs.get("body"))
        event_id = str(uuid.uuid4())
    except Exception:
        return original_fn(**kwargs)

    try:
        response = original_fn(**kwargs)
    except Exception:
        try:
            raindrop.track_ai(
                user_id=user_id or "unknown",
                event="ai_generation",
                event_id=event_id,
                input=input_text,
                model=model_id,
                convo_id=convo_id,
            )
        except Exception:
            pass
        raise

    try:
        output_text, prompt_tokens, completion_tokens = _parse_invoke_response(response)

        properties: Dict[str, Any] = {}
        if prompt_tokens is not None:
            properties["ai.usage.prompt_tokens"] = prompt_tokens
        if completion_tokens is not None:
            properties["ai.usage.completion_tokens"] = completion_tokens

        raindrop.track_ai(
            user_id=user_id or "unknown",
            event="ai_generation",
            event_id=event_id,
            input=input_text,
            output=output_text,
            model=model_id,
            properties=properties or None,
            convo_id=convo_id,
        )
    except Exception:
        pass

    return response


def _format_converse_input(messages: Any) -> Optional[str]:
    if not messages:
        return None
    user_msgs = [m for m in messages if isinstance(m, dict) and m.get("role") == "user"]
    to_extract = user_msgs if user_msgs else messages[-1:]
    parts = []
    for msg in to_extract:
        content = msg.get("content", []) if isinstance(msg, dict) else []
        text_parts = []
        for block in (content if isinstance(content, list) else []):
            if isinstance(block, dict) and block.get("text"):
                text_parts.append(block["text"])
        if text_parts:
            parts.append(" ".join(text_parts))
    return "\n".join(parts) if parts else None


def _extract_converse_output(response: Dict[str, Any]) -> Optional[str]:
    output = response.get("output", {})
    message = output.get("message", {}) if isinstance(output, dict) else {}
    content = message.get("content", []) if isinstance(message, dict) else []
    texts = []
    for block in (content if isinstance(content, list) else []):
        if isinstance(block, dict) and "text" in block:
            texts.append(block["text"])
    return "\n".join(texts) if texts else None


def _parse_body(body: Any) -> Optional[str]:
    if body is None:
        return None
    try:
        if isinstance(body, (bytes, bytearray)):
            return body.decode("utf-8")
        if isinstance(body, str):
            return body
        if hasattr(body, "read"):
            data = body.read()
            if hasattr(body, "seek"):
                body.seek(0)
            return data.decode("utf-8") if isinstance(data, bytes) else str(data)
        return str(body)
    except Exception:
        return None


def _parse_invoke_response(
    response: Dict[str, Any],
) -> tuple:
    """Returns (output_text, prompt_tokens, completion_tokens).

    If the body is a streaming object (like botocore.response.StreamingBody),
    it is read and replaced with a BytesIO so the caller can still call .read().
    """
    body = response.get("body")
    if not body:
        return None, None, None

    try:
        if hasattr(body, "read"):
            raw = body.read()
            import io
            response["body"] = io.BytesIO(raw)
        elif isinstance(body, (bytes, bytearray)):
            raw = body
        else:
            raw = str(body).encode("utf-8")

        parsed = json.loads(raw)

        # Anthropic Claude format
        content = parsed.get("content")
        if isinstance(content, list):
            output = "\n".join(
                c.get("text", "") for c in content if isinstance(c, dict) and c.get("type") == "text"
            ) or None
            usage = parsed.get("usage", {})
            return output, usage.get("input_tokens"), usage.get("output_tokens")

        # Amazon Titan format
        results = parsed.get("results")
        if isinstance(results, list) and results:
            output = "\n".join(
                r.get("outputText", "") for r in results if isinstance(r, dict)
            ) or None
            return output, parsed.get("inputTextTokenCount"), results[0].get("tokenCount") if isinstance(results[0], dict) else None

        # Llama format
        generation = parsed.get("generation")
        if generation:
            return generation, parsed.get("prompt_token_count"), parsed.get("generation_token_count")

        # Fallback
        output = parsed.get("output") or parsed.get("completion") or parsed.get("generated_text")
        return output, None, None

    except Exception:
        return None, None, None
