"""Google Generative AI auto-instrumentation.

Patches ``google.generativeai.GenerativeModel.generate_content`` and
``generate_content_async`` to capture model_call spans.
"""
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="google",
    input_tokens_field="prompt_token_count",
    output_tokens_field="candidates_token_count",
    finish_reason_getter=lambda r: None,
)


def _extract_google_usage(response: Any) -> tuple[int | None, int | None]:
    """Extract token counts from a Google GenerateContentResponse."""
    usage = getattr(response, "usage_metadata", None)
    if usage is None:
        return None, None
    inp = getattr(usage, "prompt_token_count", None)
    out = getattr(usage, "candidates_token_count", None)
    return inp, out


def _wrap_generate(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, *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 = _extract_google_usage(response)
                model = getattr(self, "model_name", None) or getattr(self, "_model_name", None)
                # Normalize model name (remove "models/" prefix)
                if model and model.startswith("models/"):
                    model = model[len("models/"):]
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="google",
                    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, inp, out),
                )
                sd.set_model_call(model_name=model)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture google_genai span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model_name", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture google_genai error span", exc_info=True)
            raise

    return wrapper


def _wrap_agenerate(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(self, *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 = _extract_google_usage(response)
                model = getattr(self, "model_name", None) or getattr(self, "_model_name", None)
                if model and model.startswith("models/"):
                    model = model[len("models/"):]
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="google",
                    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, inp, out),
                )
                sd.set_model_call(model_name=model)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async google_genai span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model_name", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async google_genai error span", exc_info=True)
            raise

    return wrapper


class GoogleGenAIInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_generate: Any = None
        self._original_agenerate: Any = None

    def instrument(self) -> None:
        import google.generativeai as genai

        model_cls = genai.GenerativeModel
        self._original_generate = model_cls.generate_content
        model_cls.generate_content = _wrap_generate(self._original_generate)

        self._original_agenerate = model_cls.generate_content_async
        model_cls.generate_content_async = _wrap_agenerate(self._original_agenerate)

    def uninstrument(self) -> None:
        if self._original_generate is None:
            return
        import google.generativeai as genai

        genai.GenerativeModel.generate_content = self._original_generate
        genai.GenerativeModel.generate_content_async = self._original_agenerate
        self._original_generate = None
        self._original_agenerate = None
