from __future__ import annotations

import json
import warnings
from collections.abc import Callable, Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from decimal import Decimal
from typing import Any

from opentelemetry.util.types import AttributeValue

from pydantic_ai._instrumentation import (
    ANY_ADAPTER,
    GEN_AI_PROVIDER_NAME_ATTRIBUTE,
    GEN_AI_REQUEST_MODEL_ATTRIBUTE,
    record_uncaught_errors,
    server_attributes,
)
from pydantic_ai._warnings import CostCalculationFailedWarning
from pydantic_ai.models.instrumented import InstrumentationSettings

from .base import EmbeddingModel
from .result import EmbeddingResult, EmbedInputType
from .settings import EmbeddingSettings
from .wrapper import WrapperEmbeddingModel

__all__ = 'instrument_embedding_model', 'InstrumentedEmbeddingModel'


def instrument_embedding_model(model: EmbeddingModel, instrument: InstrumentationSettings | bool) -> EmbeddingModel:
    """Instrument an embedding model with OpenTelemetry/logfire."""
    if instrument and not isinstance(model, InstrumentedEmbeddingModel):
        if instrument is True:
            instrument = InstrumentationSettings()

        model = InstrumentedEmbeddingModel(model, instrument)

    return model


@dataclass(init=False)
class InstrumentedEmbeddingModel(WrapperEmbeddingModel):
    """Embedding model which wraps another model so that requests are instrumented with OpenTelemetry.

    See the [Debugging and Monitoring guide](https://pydantic.dev/docs/ai/integrations/logfire/) for more info.
    """

    instrumentation_settings: InstrumentationSettings
    """Instrumentation settings for this model."""

    def __init__(
        self,
        wrapped: EmbeddingModel | str,
        options: InstrumentationSettings | None = None,
    ) -> None:
        super().__init__(wrapped)
        self.instrumentation_settings = options or InstrumentationSettings()

    async def embed(
        self, inputs: str | Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
    ) -> EmbeddingResult:
        inputs, settings = self.prepare_embed(inputs, settings)
        with self._instrument(inputs, input_type, settings) as finish:
            result = await super().embed(inputs, input_type=input_type, settings=settings)
            finish(result)
            return result

    @staticmethod
    def _recorded_settings(settings: EmbeddingSettings | None) -> dict[str, object]:
        """The settings recorded on the span: all but the request header and body overrides."""
        return {key: value for key, value in (settings or {}).items() if key not in ('extra_headers', 'extra_body')}

    @contextmanager
    def _instrument(
        self,
        inputs: list[str],
        input_type: EmbedInputType,
        settings: EmbeddingSettings | None,
    ) -> Generator[Callable[[EmbeddingResult], None]]:
        operation = 'embeddings'
        span_name = f'{operation} {self.model_name}'

        inputs_count = len(inputs)

        attributes: dict[str, AttributeValue] = {
            'gen_ai.operation.name': operation,
            **self.model_attributes(self.wrapped),
            'input_type': input_type,
            'inputs_count': inputs_count,
        }

        if recorded_settings := self._recorded_settings(settings):
            attributes['embedding_settings'] = json.dumps(self.serialize_any(recorded_settings))

        if self.instrumentation_settings.include_content:
            attributes['inputs'] = json.dumps(inputs)

        attributes['logfire.json_schema'] = json.dumps(
            {
                'type': 'object',
                'properties': {
                    'input_type': {'type': 'string'},
                    'inputs_count': {'type': 'integer'},
                    'embedding_settings': {'type': 'object'},
                    **(
                        {'inputs': {'type': ['array']}, 'embeddings': {'type': 'array'}}
                        if self.instrumentation_settings.include_content
                        else {}
                    ),
                },
            }
        )

        record_metrics: Callable[[], None] | None = None
        try:
            with (
                self.instrumentation_settings.tracer.start_as_current_span(
                    span_name, attributes=attributes, record_exception=False, set_status_on_exception=False
                ) as span,
                record_uncaught_errors(span, include_content=self.instrumentation_settings.include_content),
            ):

                def finish(result: EmbeddingResult):
                    # Prepare metric recording closure first so metrics are recorded
                    # even if the span is not recording.
                    provider_name = attributes[GEN_AI_PROVIDER_NAME_ATTRIBUTE]
                    request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE]
                    response_model = result.model_name or request_model
                    cost = result.usage.cost

                    def _record_metrics():
                        # Reads `cost` when it runs, after it has been estimated below.
                        self._record_result_metrics(
                            result, provider_name, operation, request_model, response_model, cost
                        )

                    nonlocal record_metrics
                    record_metrics = _record_metrics

                    cost = _embedding_cost(result)

                    if not span.is_recording():
                        return  # pragma: lax no cover

                    attributes_to_set: dict[str, AttributeValue] = {
                        **result.usage.opentelemetry_attributes(),
                        'gen_ai.response.model': response_model,
                    }

                    if cost is not None:
                        attributes_to_set['operation.cost'] = float(cost)

                    embeddings = result.embeddings
                    if embeddings:  # pragma: no branch
                        attributes_to_set['gen_ai.embeddings.dimension.count'] = len(embeddings[0])
                        if self.instrumentation_settings.include_content:
                            attributes_to_set['embeddings'] = json.dumps(embeddings)

                    if result.provider_response_id is not None:
                        attributes_to_set['gen_ai.response.id'] = result.provider_response_id

                    span.set_attributes(attributes_to_set)

                yield finish
        finally:
            if record_metrics:  # pragma: no branch
                # Record metrics after the span finishes to avoid duplication.
                record_metrics()

    def _record_result_metrics(
        self,
        result: EmbeddingResult,
        provider_name: AttributeValue,
        operation: str,
        request_model: AttributeValue,
        response_model: AttributeValue,
        cost: Decimal | None,
    ) -> None:
        metric_attributes = {
            GEN_AI_PROVIDER_NAME_ATTRIBUTE: provider_name,
            'gen_ai.operation.name': operation,
            GEN_AI_REQUEST_MODEL_ATTRIBUTE: request_model,
            'gen_ai.response.model': response_model,
        }
        token_attributes = {
            **metric_attributes,
            'gen_ai.token.type': 'input',
        }
        tokens = result.usage.input_tokens or 0
        if tokens:  # pragma: no branch
            self.instrumentation_settings.tokens_histogram.record(tokens, token_attributes)
            if cost is not None:
                self.instrumentation_settings.cost_histogram.record(float(cost), metric_attributes)

    @staticmethod
    def model_attributes(model: EmbeddingModel) -> dict[str, AttributeValue]:
        return {
            GEN_AI_PROVIDER_NAME_ATTRIBUTE: model.system,
            GEN_AI_REQUEST_MODEL_ATTRIBUTE: model.model_name,
            **server_attributes(model.base_url),
        }

    @staticmethod
    def serialize_any(value: Any) -> str:
        try:
            return ANY_ADAPTER.dump_python(value, mode='json')
        except Exception:  # pragma: no cover
            try:
                return str(value)
            except Exception as e:
                return f'Unable to serialize: {e}'


def _embedding_cost(result: EmbeddingResult) -> Decimal | None:
    """The cost the usage reports, else the estimate: `None` when it's unknown, or when unmeasured usage estimates to zero."""
    cost = result.usage.cost
    try:
        # A cost already on the usage, e.g. reported by the Pydantic AI Gateway, wins over the estimate.
        price = cost if cost is not None else result.cost().total_price
        # Match best_effort_price's policy while retaining embedding-specific lookup and warnings.
        return None if cost is None and result.usage.unmeasured_requests and price == 0 else price
    except LookupError:
        # The cost of this provider/model is unknown, which is common.
        return cost
    except Exception as e:
        warnings.warn(
            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning, stacklevel=2
        )
        return cost