"""Vercel AI SDK (Python) auto-instrumentation.

Patches ``ai.generate_text`` and ``ai.stream_text`` (when available in
the Python edge-case wrapper) to capture model_call spans.

The Vercel AI SDK is primarily TypeScript. This instrumentor targets the
thin Python wrappers that some projects use for server-side calls.
"""
from __future__ import annotations

import functools
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="vercel_ai",
    input_tokens_field="prompt_tokens",
    output_tokens_field="completion_tokens",
    finish_reason_getter=lambda r: getattr(r, "finish_reason", None) if not isinstance(r, dict) else r.get("finishReason"),
)


def _extract_vercel_usage(response: Any) -> tuple[int | None, int | None, str | None]:
    """Extract usage from a Vercel AI response object."""
    usage = getattr(response, "usage", None)
    if usage is None and isinstance(response, dict):
        usage = response.get("usage", {})

    if usage is None:
        return None, None, None

    if isinstance(usage, dict):
        inp = usage.get("promptTokens") or usage.get("prompt_tokens")
        out = usage.get("completionTokens") or usage.get("completion_tokens")
    else:
        inp = getattr(usage, "promptTokens", None) or getattr(usage, "prompt_tokens", None)
        out = getattr(usage, "completionTokens", None) or getattr(usage, "completion_tokens", None)

    finish = None
    if isinstance(response, dict):
        finish = response.get("finishReason")
    else:
        finish = getattr(response, "finishReason", None) or getattr(response, "finish_reason", None)

    return inp, out, finish


def _wrap_generate_text(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(*args, **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, finish = _extract_vercel_usage(response)
                model = kwargs.get("model") or (args[0] if args else None)
                model_name = model if isinstance(model, str) else getattr(model, "modelId", None)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="vercel_ai",
                    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_name, inp, out),
                )
                sd.set_model_call(model_name=model_name, finish_reason=finish)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture vercel_ai 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, kwargs, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture vercel_ai error span", exc_info=True)
            raise

    return wrapper


def _wrap_agenerate_text(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(*args, **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, finish = _extract_vercel_usage(response)
                model = kwargs.get("model") or (args[0] if args else None)
                model_name = model if isinstance(model, str) else getattr(model, "modelId", None)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="vercel_ai",
                    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_name, inp, out),
                )
                sd.set_model_call(model_name=model_name, finish_reason=finish)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async vercel_ai 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, kwargs, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async vercel_ai error span", exc_info=True)
            raise

    return wrapper


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

    def instrument(self) -> None:
        import ai

        if hasattr(ai, "generate_text"):
            self._originals["generate_text"] = ai.generate_text
            ai.generate_text = _wrap_generate_text(ai.generate_text)

        if hasattr(ai, "agenerate_text"):
            self._originals["agenerate_text"] = ai.agenerate_text
            ai.agenerate_text = _wrap_agenerate_text(ai.agenerate_text)

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

        if "generate_text" in self._originals:
            ai.generate_text = self._originals["generate_text"]
        if "agenerate_text" in self._originals:
            ai.agenerate_text = self._originals["agenerate_text"]
        self._originals.clear()
