import base64
import json
import os
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Optional

from litellm._logging import verbose_logger
from litellm.integrations.arize import _utils
from litellm.integrations.langfuse.langfuse_otel_attributes import (
    LangfuseLLMObsOTELAttributes,
)
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.types.integrations.langfuse_otel import (
    LangfuseSpanAttributes,
)
from litellm.types.utils import StandardCallbackDynamicParams

if TYPE_CHECKING:
    from opentelemetry.trace import Span as _Span

    Span = _Span | Any
else:
    Span = Any


LANGFUSE_CLOUD_EU_ENDPOINT: Final = "https://cloud.langfuse.com/api/public/otel"
LANGFUSE_CLOUD_US_ENDPOINT: Final = "https://us.cloud.langfuse.com/api/public/otel"
LANGFUSE_INGESTION_VERSION_HEADER: Final = "x-langfuse-ingestion-version"
LANGFUSE_INGESTION_VERSION: Final = "4"


class LangfuseOtelLogger(OpenTelemetry):
    def __init__(self, config=None, *args, **kwargs):
        # Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually
        # and passing it to the parent OpenTelemetry class
        if config is None:
            config = self._create_open_telemetry_config_from_langfuse_env()
        super().__init__(config=config, *args, **kwargs)

    @staticmethod
    def set_langfuse_otel_attributes(span: Span, kwargs, response_obj):
        """
        Sets OpenTelemetry span attributes for Langfuse observability.
        Uses the same attribute setting logic as Arize Phoenix for consistency.
        """

        _utils.set_attributes(span, kwargs, response_obj, LangfuseLLMObsOTELAttributes)
        span.set_attribute("langfuse.observation.type", "generation")

        #########################################################
        # Set Langfuse specific attributes
        #########################################################
        LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj)

    @staticmethod
    def _extract_langfuse_metadata(kwargs: dict) -> dict:
        """
        Extracts Langfuse metadata from the standard LiteLLM kwargs structure.

        1. Reads kwargs["litellm_params"]["metadata"] if present and is a dict.
        2. Enriches it with any `langfuse_*` request-header params via the
           existing LangFuseLogger.add_metadata_from_header helper so that proxy
           users get identical behaviour across vanilla and OTEL integrations.
        """
        litellm_params: Final = kwargs.get("litellm_params", {}) or {}
        metadata = litellm_params.get("metadata") or {}
        # Ensure we only work with dicts
        if metadata is None or not isinstance(metadata, dict):
            metadata = {}

        # Re-use header extraction logic from the vanilla logger if available
        try:
            from litellm.integrations.langfuse.langfuse import (
                LangFuseLogger as _LFLogger,
            )

            metadata = _LFLogger.add_metadata_from_header(litellm_params, metadata)
        except Exception:
            # Fallback silently if import fails; header enrichment just won't happen
            pass

        return metadata

    @staticmethod
    def _set_metadata_attributes(span: Span, metadata: dict):
        """Helper to set metadata attributes from mapping."""
        from litellm.integrations.arize._utils import safe_set_attribute

        mapping: Final = {
            "generation_name": LangfuseSpanAttributes.GENERATION_NAME,
            "generation_id": LangfuseSpanAttributes.GENERATION_ID,
            "parent_observation_id": LangfuseSpanAttributes.PARENT_OBSERVATION_ID,
            "mask_input": LangfuseSpanAttributes.MASK_INPUT,
            "mask_output": LangfuseSpanAttributes.MASK_OUTPUT,
            "trace_user_id": LangfuseSpanAttributes.TRACE_USER_ID,
            "session_id": LangfuseSpanAttributes.SESSION_ID,
            "tags": LangfuseSpanAttributes.TAGS,
            "trace_name": LangfuseSpanAttributes.TRACE_NAME,
            "trace_id": LangfuseSpanAttributes.TRACE_ID,
            "trace_metadata": LangfuseSpanAttributes.TRACE_METADATA,
            "trace_release": LangfuseSpanAttributes.RELEASE,
            "existing_trace_id": LangfuseSpanAttributes.EXISTING_TRACE_ID,
            "update_trace_keys": LangfuseSpanAttributes.UPDATE_TRACE_KEYS,
            "debug_langfuse": LangfuseSpanAttributes.DEBUG_LANGFUSE,
        }

        version: Final = (
            metadata.get("trace_version") if metadata.get("trace_version") is not None else metadata.get("version")
        )
        if version is not None:
            safe_set_attribute(span, LangfuseSpanAttributes.VERSION.value, version)

        for key, enum_attr in mapping.items():
            if key in metadata and metadata[key] is not None:
                value = metadata[key]
                if key == "trace_id" and isinstance(value, str):
                    # trace_id must be 32 hex char no dashes for langfuse : Litellm sends uuid with dashes (might be breaking at some point)
                    value = value.replace("-", "")

                if isinstance(value, (list, dict)):
                    try:
                        value = json.dumps(value)
                    except Exception:
                        value = str(value)
                safe_set_attribute(span, enum_attr.value, value)

    @staticmethod
    def _set_observation_output(span: Span, response_obj):
        """Helper to set observation output attributes."""
        from litellm.integrations.arize._utils import safe_set_attribute
        from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

        if not response_obj or not hasattr(response_obj, "get"):
            return

        choices: Final = response_obj.get("choices", [])
        if choices:
            first_choice: Final = choices[0]
            message: Final = first_choice.get("message", {})
            tool_calls: Final = message.get("tool_calls")
            if tool_calls:
                transformed_tool_calls: Final = []
                for tool_call in tool_calls:
                    function = tool_call.get("function", {})
                    arguments_str = function.get("arguments", "{}")
                    try:
                        arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
                    except json.JSONDecodeError:
                        arguments_obj = {}
                    langfuse_tool_call = {
                        "id": response_obj.get("id", ""),
                        "name": function.get("name", ""),
                        "call_id": tool_call.get("id", ""),
                        "type": "function_call",
                        "arguments": arguments_obj,
                    }
                    transformed_tool_calls.append(langfuse_tool_call)
                safe_set_attribute(
                    span,
                    LangfuseSpanAttributes.OBSERVATION_OUTPUT.value,
                    safe_dumps(transformed_tool_calls),
                )
            else:
                output_data: Final = {}
                if message.get("role"):
                    output_data["role"] = message.get("role")
                if message.get("content") is not None:
                    output_data["content"] = message.get("content")
                if output_data:
                    safe_set_attribute(
                        span,
                        LangfuseSpanAttributes.OBSERVATION_OUTPUT.value,
                        safe_dumps(output_data),
                    )

        output: Final = response_obj.get("output", [])
        if output:
            output_items_data: Final[list[dict]] = []
            for item in output:
                if hasattr(item, "type"):
                    item_type = item.type
                    if item_type == "reasoning" and hasattr(item, "summary"):
                        for summary in item.summary:
                            if hasattr(summary, "text"):
                                output_items_data.append(
                                    {
                                        "role": "reasoning_summary",
                                        "content": summary.text,
                                    }
                                )
                    elif item_type == "message":
                        output_items_data.append(
                            {
                                "role": getattr(item, "role", "assistant"),
                                "content": getattr(getattr(item, "content", [{}])[0], "text", ""),
                            }
                        )
                    elif item_type == "function_call":
                        arguments_str = getattr(item, "arguments", "{}")
                        arguments_obj = (
                            safe_json_loads(arguments_str, default={})
                            if isinstance(arguments_str, str)
                            else arguments_str
                        )
                        langfuse_tool_call = {
                            "id": getattr(item, "id", ""),
                            "name": getattr(item, "name", ""),
                            "call_id": getattr(item, "call_id", ""),
                            "type": "function_call",
                            "arguments": arguments_obj,
                        }
                        output_items_data.append(langfuse_tool_call)
            if output_items_data:
                safe_set_attribute(
                    span,
                    LangfuseSpanAttributes.OBSERVATION_OUTPUT.value,
                    safe_dumps(output_items_data),
                )

    @staticmethod
    def _set_langfuse_specific_attributes(span: Span, kwargs, response_obj):
        """
        Sets Langfuse specific metadata attributes onto the OTEL span.

        All keys supported by the vanilla Langfuse integration are mapped to
        OTEL-safe attribute names defined in LangfuseSpanAttributes.  Complex
        values (lists/dicts) are serialised to JSON strings for OTEL
        compatibility.
        """
        from litellm.integrations.arize._utils import safe_set_attribute
        from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

        dynamic_params: Final = kwargs.get("standard_callback_dynamic_params")
        langfuse_environment: Final = (
            dynamic_params.get("langfuse_environment") if dynamic_params else None
        ) or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
        if langfuse_environment:
            safe_set_attribute(
                span,
                LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value,
                langfuse_environment,
            )

        metadata: Final = LangfuseOtelLogger._extract_langfuse_metadata(kwargs)
        LangfuseOtelLogger._set_metadata_attributes(span=span, metadata=metadata)

        messages: Final = kwargs.get("messages")
        if messages:
            safe_set_attribute(
                span,
                LangfuseSpanAttributes.OBSERVATION_INPUT.value,
                safe_dumps(messages),
            )

        LangfuseOtelLogger._set_observation_output(span=span, response_obj=response_obj)

    @staticmethod
    def _get_langfuse_otel_host() -> str | None:
        """
        Returns the Langfuse OTEL host based on environment variables.

        Returned in the following order of precedence:
        1. LANGFUSE_OTEL_HOST
        2. LANGFUSE_HOST
        """
        return os.environ.get("LANGFUSE_OTEL_HOST") or os.environ.get("LANGFUSE_HOST")

    def _create_open_telemetry_config_from_langfuse_env(self) -> OpenTelemetryConfig:
        """
        Creates OpenTelemetryConfig from Langfuse environment variables.
        Does NOT modify global environment variables.
        """
        from litellm.integrations.opentelemetry import OpenTelemetryConfig

        public_key: Final = os.environ.get("LANGFUSE_PUBLIC_KEY", None)
        secret_key: Final = os.environ.get("LANGFUSE_SECRET_KEY", None)

        if not public_key or not secret_key:
            # If no keys, return default from env (likely logging to console or something else)
            return OpenTelemetryConfig.from_env()

        return LangfuseOtelLogger._build_langfuse_otel_config(
            public_key=public_key,
            secret_key=secret_key,
            langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
        )

    @staticmethod
    def get_langfuse_otel_config() -> "OpenTelemetryConfig":
        """
        Retrieves the Langfuse OpenTelemetry configuration based on environment variables.

        Environment Variables:
            LANGFUSE_PUBLIC_KEY: Required. Langfuse public key for authentication.
            LANGFUSE_SECRET_KEY: Required. Langfuse secret key for authentication.
            LANGFUSE_HOST: Optional. Custom Langfuse host URL. Defaults to US cloud.

        Returns:
            OpenTelemetryConfig: A Pydantic model containing Langfuse OTEL configuration.

        Raises:
            ValueError: If required keys are missing.
        """
        public_key: Final = os.environ.get("LANGFUSE_PUBLIC_KEY", None)
        secret_key: Final = os.environ.get("LANGFUSE_SECRET_KEY", None)

        if not public_key or not secret_key:
            raise ValueError(
                "LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set for Langfuse OpenTelemetry integration."
            )

        return LangfuseOtelLogger._build_langfuse_otel_config(
            public_key=public_key,
            secret_key=secret_key,
            langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
        )

    @staticmethod
    def _build_langfuse_otel_config(
        public_key: str, secret_key: str, langfuse_host: str | None
    ) -> "OpenTelemetryConfig":
        """
        Builds an OTLP HTTP config pointing at the Langfuse OTEL endpoint for the
        given host (US cloud when no host is provided), authorized with the given keys.
        """
        if langfuse_host:
            normalized_host: Final = langfuse_host if langfuse_host.startswith("http") else f"https://{langfuse_host}"
            endpoint = f"{normalized_host.rstrip('/')}/api/public/otel"
            verbose_logger.debug("Using Langfuse OTEL endpoint from host: %s", endpoint)
        else:
            endpoint = LANGFUSE_CLOUD_US_ENDPOINT
            verbose_logger.debug("Using Langfuse US cloud endpoint: %s", endpoint)

        auth_header: Final = LangfuseOtelLogger._get_langfuse_authorization_header(
            public_key=public_key, secret_key=secret_key
        )

        return OpenTelemetryConfig(
            exporter="otlp_http",
            endpoint=endpoint,
            headers=LangfuseOtelLogger._format_otel_headers(
                LangfuseOtelLogger._build_langfuse_otel_headers(auth_header)
            ),
        )

    @staticmethod
    def _get_langfuse_authorization_header(public_key: str, secret_key: str) -> str:
        """
        Get the authorization header for Langfuse OpenTelemetry.
        """
        auth_string: Final = f"{public_key}:{secret_key}"
        auth_header: Final = base64.b64encode(auth_string.encode()).decode()
        return f"Basic {auth_header}"

    @staticmethod
    def _build_langfuse_otel_headers(auth_header: str) -> dict[str, str]:
        """
        Build the OTLP header set Langfuse expects.

        `x-langfuse-ingestion-version: 4` selects Langfuse's v4 ingestion path;
        without it spans fall back to the older transformation path.
        """
        return {
            "Authorization": auth_header,
            LANGFUSE_INGESTION_VERSION_HEADER: LANGFUSE_INGESTION_VERSION,
        }

    @staticmethod
    def _format_otel_headers(headers: dict[str, str]) -> str:
        """
        Serialize a header mapping into the comma-separated OTLP header string
        """
        return ",".join(f"{key}={value}" for key, value in headers.items())

    def construct_dynamic_otel_headers(
        self, standard_callback_dynamic_params: StandardCallbackDynamicParams
    ) -> dict | None:
        """
        Construct dynamic Langfuse headers from standard callback dynamic params

        This is used for team/key based logging.

        Returns:
            dict: A dictionary of dynamic Langfuse headers
        """
        dynamic_headers: Final = {}

        dynamic_langfuse_public_key: Final = standard_callback_dynamic_params.get("langfuse_public_key")
        dynamic_langfuse_secret_key: Final = standard_callback_dynamic_params.get("langfuse_secret_key")
        if dynamic_langfuse_public_key and dynamic_langfuse_secret_key:
            auth_header: Final = LangfuseOtelLogger._get_langfuse_authorization_header(
                public_key=dynamic_langfuse_public_key,
                secret_key=dynamic_langfuse_secret_key,
            )
            dynamic_headers.update(LangfuseOtelLogger._build_langfuse_otel_headers(auth_header))

        return dynamic_headers

    def construct_dynamic_otel_config(
        self, standard_callback_dynamic_params: StandardCallbackDynamicParams
    ) -> Optional["OpenTelemetryConfig"]:
        """
        Build a full per-request OTLP config from team/key dynamic Langfuse credentials.

        Key-scoped credentials must define the export target, not just the auth
        headers: without this, a proxy with no global LANGFUSE_* env vars keeps its
        init-time fallback exporter (console), so key-level langfuse_otel silently
        never reaches Langfuse.
        """
        public_key: Final = standard_callback_dynamic_params.get("langfuse_public_key")
        secret_key: Final = standard_callback_dynamic_params.get("langfuse_secret_key")
        if not public_key or not secret_key:
            return None

        langfuse_host: Final = standard_callback_dynamic_params.get("langfuse_host") or self._get_langfuse_otel_host()
        return LangfuseOtelLogger._build_langfuse_otel_config(
            public_key=public_key,
            secret_key=secret_key,
            langfuse_host=langfuse_host,
        )

    def create_litellm_proxy_request_started_span(
        self,
        start_time: datetime,
        headers: dict,
    ) -> Span | None:
        """
        Override to prevent creating empty proxy request spans.

        Langfuse should only receive spans for actual LLM calls, not for
        internal proxy operations (auth, postgres, proxy_pre_call, etc.).

        By returning None, we prevent the parent span from being created,
        which in turn prevents empty traces from being sent to Langfuse.
        """
        return None

    async def async_service_success_hook(self, *args, **kwargs):
        """
        Langfuse should not receive service success logs.
        """

    async def async_service_failure_hook(self, *args, **kwargs):
        """
        Langfuse should not receive service failure logs.
        """
