"""AWS Bedrock auto-instrumentation.

Patches ``botocore.client.BedrockRuntime.invoke_model`` and
``converse`` to capture model_call spans.
"""
from __future__ import annotations

import functools
import json
import logging
from datetime import datetime, timezone
from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, _make_error_span

logger = logging.getLogger("fibinaci.instrumentation")

_CONFIG = ProviderConfig(
    name="bedrock",
    input_tokens_field="inputTokens",
    output_tokens_field="outputTokens",
    finish_reason_getter=lambda r: r.get("stopReason") if isinstance(r, dict) else None,
)


def _extract_converse_usage(response: dict) -> tuple[int | None, int | None, str | None]:
    """Extract token counts from a Bedrock Converse response."""
    usage = response.get("usage", {})
    inp = usage.get("inputTokens")
    out = usage.get("outputTokens")
    stop = response.get("stopReason")
    return inp, out, stop


def _extract_invoke_usage(response: dict, model_id: str) -> tuple[int | None, int | None, str | None]:
    """Extract token counts from a raw invoke_model response body."""
    body = response.get("body")
    if body is None:
        return None, None, None

    try:
        if hasattr(body, "read"):
            raw = body.read()
            # Re-wrap so downstream code can still read
            import io
            response["body"] = io.BytesIO(raw)
            data = json.loads(raw)
        else:
            data = json.loads(body) if isinstance(body, (str, bytes)) else body
    except (json.JSONDecodeError, TypeError):
        return None, None, None

    # Anthropic via Bedrock
    if "anthropic" in model_id:
        usage = data.get("usage", {})
        return usage.get("input_tokens"), usage.get("output_tokens"), data.get("stop_reason")

    # Amazon Titan
    if "titan" in model_id:
        results = data.get("results", [{}])
        token_count = results[0].get("tokenCount") if results else None
        inp = data.get("inputTextTokenCount")
        return inp, token_count, data.get("completionReason")

    # Generic fallback
    usage = data.get("usage", {})
    return usage.get("input_tokens") or usage.get("prompt_tokens"), \
           usage.get("output_tokens") or usage.get("completion_tokens"), \
           data.get("stop_reason")


def _wrap_converse(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                inp, out, stop = _extract_converse_usage(response)
                model_id = kwargs.get("modelId", "")
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="bedrock",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model_id, inp, out),
                )
                sd.set_model_call(model_name=model_id, finish_reason=stop)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture bedrock converse span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, {"model": kwargs.get("modelId")}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture bedrock converse error span", exc_info=True)
            raise

    return wrapper


def _wrap_invoke_model(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                model_id = kwargs.get("modelId", "")
                inp, out, stop = _extract_invoke_usage(response, model_id)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="bedrock",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model_id, inp, out),
                )
                sd.set_model_call(model_name=model_id, finish_reason=stop)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture bedrock invoke_model span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, {"model": kwargs.get("modelId")}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture bedrock invoke_model error span", exc_info=True)
            raise

    return wrapper


class BedrockInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._originals: dict[str, Any] = {}
        self._patched_client_cls: Any = None

    def instrument(self) -> None:
        import botocore.client

        # We patch the general client _make_api_call indirectly by
        # targeting the specific service methods after they're created.
        # Instead, we hook via botocore event system for cleaner patching.
        original_init = botocore.client.ClientCreator.create_client

        originals = self._originals
        originals["create_client"] = original_init

        @functools.wraps(original_init)
        def patched_create_client(creator_self: Any, *args: Any, **kwargs: Any) -> Any:
            client = original_init(creator_self, *args, **kwargs)
            service_name = getattr(client, "_service_model", None)
            svc = getattr(service_name, "service_name", "") if service_name else ""

            if svc == "bedrock-runtime":
                if hasattr(client, "converse") and "converse" not in originals:
                    originals["converse"] = client.converse
                    client.converse = _wrap_converse(client.converse).__get__(client)

                if hasattr(client, "invoke_model") and "invoke_model" not in originals:
                    originals["invoke_model"] = client.invoke_model
                    client.invoke_model = _wrap_invoke_model(client.invoke_model).__get__(client)

            return client

        botocore.client.ClientCreator.create_client = patched_create_client

    def uninstrument(self) -> None:
        if not self._originals:
            return
        import botocore.client

        if "create_client" in self._originals:
            botocore.client.ClientCreator.create_client = self._originals["create_client"]
        self._originals.clear()
