"""Shared instrumentation wrapper factory for LLM providers."""
from __future__ import annotations

import functools
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from fibinaci.tracing.span_builder import SpanData, estimate_cost

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


@dataclass
class ProviderConfig:
    """Provider-specific field mapping for span extraction."""
    name: str  # "openai" or "anthropic"
    input_tokens_field: str  # "prompt_tokens" or "input_tokens"
    output_tokens_field: str  # "completion_tokens" or "output_tokens"
    finish_reason_getter: Any = None  # callable(response) -> str | None


def _extract_span(
    config: ProviderConfig,
    kwargs: dict,
    response: Any,
    started_at: datetime,
    ended_at: datetime,
) -> dict:
    from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id

    usage = getattr(response, "usage", None)
    tokens_input = getattr(usage, config.input_tokens_field, None) if usage else None
    tokens_output = getattr(usage, config.output_tokens_field, None) if usage else None
    model = getattr(response, "model", None) or kwargs.get("model")

    finish_reason = None
    if config.finish_reason_getter:
        finish_reason = config.finish_reason_getter(response)

    sd = SpanData(
        trace_id=get_current_trace_id() or _new_trace_id(),
        span_id=_new_span_id(),
        span_kind="model_call",
        provider=config.name,
        model_type="llm",
        started_at=started_at,
        ended_at=ended_at,
        duration_ms=(ended_at - started_at).total_seconds() * 1000,
        status="ok",
        tokens_input=tokens_input,
        tokens_output=tokens_output,
        tokens_total=(tokens_input or 0) + (tokens_output or 0) if tokens_input else None,
        cost_usd=estimate_cost(model, tokens_input, tokens_output),
    )
    sd.set_model_call(
        model_name=model,
        temperature=kwargs.get("temperature"),
        max_tokens=kwargs.get("max_tokens"),
        finish_reason=finish_reason,
    )
    return sd.to_dict()


def _make_error_span(config: ProviderConfig, kwargs: dict, exc: Exception, started_at: datetime, ended_at: datetime) -> dict:
    from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id

    sd = SpanData(
        trace_id=get_current_trace_id() or _new_trace_id(),
        span_id=_new_span_id(),
        span_kind="model_call",
        provider=config.name,
        model_type="llm",
        started_at=started_at,
        ended_at=ended_at,
        duration_ms=(ended_at - started_at).total_seconds() * 1000,
        status="error",
        error_type=type(exc).__name__,
        error_message=str(exc),
    )
    sd.set_model_call(model_name=kwargs.get("model"))
    return sd.to_dict()


def wrap_sync(original: Any, config: ProviderConfig) -> 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
                get_buffer().push(_extract_span(config, kwargs, response, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture %s span", config.name, 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 %s error span", config.name, exc_info=True)
            raise

    return wrapper


def wrap_async(original: Any, config: ProviderConfig) -> 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
                get_buffer().push(_extract_span(config, kwargs, response, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async %s span", config.name, 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 %s error span", config.name, exc_info=True)
            raise

    return wrapper
