"""Mingx client: TracerProvider, OTLP export, and observation APIs."""

import base64
import os
from contextlib import contextmanager
from time import time_ns
from typing import Any, Optional

from opentelemetry import trace as otel_trace_api
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SpanProcessor
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator

from mingx._client.attributes import (
    MingxOtelSpanAttributes,
    create_trace_attributes,
)
from mingx._client.constants import MINGX_TRACER_NAME, get_observation_types_list
from mingx._client.span import MingxEvent, MingxGeneration, MingxSpan
from mingx.version import __version__


def _default_span_exporter(
    *,
    endpoint: Optional[str] = None,
    token: Optional[str] = None,
) -> Optional[SpanExporter]:
    """Create default OTLP HTTP span exporter. Lazy import to avoid loading when tracing disabled."""
    try:
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
    except Exception:
        return None
    headers: Optional[dict[str, str]] = None
    if token:
        # HTTP Basic: if token contains ":", use as "username:password"; else use ":token" (password only)
        credential = token if ":" in token else f":{token}"
        b64 = base64.b64encode(credential.encode()).decode()
        headers = {"Authorization": f"Basic {b64}"}
    return OTLPSpanExporter(endpoint=endpoint, headers=headers)


class MingxClient:
    """Client for Mingx tracing. Uses OpenTelemetry with OTLP export (env or constructor)."""

    def __init__(
        self,
        *,
        tracing_enabled: Optional[bool] = None,
        service_name: Optional[str] = None,
        environment: Optional[str] = None,
        tracer_provider: Optional[TracerProvider] = None,
        span_exporter: Optional[SpanExporter] = None,
        endpoint: Optional[str] = None,
        token: Optional[str] = None,
    ) -> None:
        """Initialize the client.

        endpoint: OTLP traces endpoint URL (e.g. https://ingest.example.com/v1/traces).
          If not set, OTEL_EXPORTER_OTLP_TRACES_ENDPOINT or OTEL_EXPORTER_OTLP_ENDPOINT is used.
        token: For HTTP Basic auth. If it contains ":", treated as "username:password"
          (e.g. "user@example.com:secret"); otherwise sent as password with empty username.
          If not set, OTEL_EXPORTER_OTLP_HEADERS can be used for auth.
        """
        env_enabled = os.environ.get("MINGX_TRACING_ENABLED", "true").lower() not in ("false", "0")
        self._tracing_enabled = tracing_enabled if tracing_enabled is not None else env_enabled
        self._environment = environment or os.environ.get("OTEL_SERVICE_NAME", "default")
        self._service_name = service_name or os.environ.get("OTEL_SERVICE_NAME", "mingx")

        if not self._tracing_enabled or os.environ.get("OTEL_SDK_DISABLED", "false").lower() == "true":
            self._tracer = otel_trace_api.NoOpTracer()
            self._tracer_provider = None
            self._processor = None
            return

        if tracer_provider is not None:
            self._tracer_provider = tracer_provider
            self._tracer = self._tracer_provider.get_tracer(MINGX_TRACER_NAME, __version__)
            self._processor = None
            return

        resource = Resource.create({
            "service.name": self._service_name,
        })
        provider = TracerProvider(resource=resource, id_generator=RandomIdGenerator())
        exporter = span_exporter or _default_span_exporter(endpoint=endpoint, token=token)
        if exporter is not None:
            processor = BatchSpanProcessor(exporter)
            provider.add_span_processor(processor)
            self._processor = processor
        else:
            self._processor = None

        otel_trace_api.set_tracer_provider(provider)
        self._tracer_provider = provider
        self._tracer = provider.get_tracer(MINGX_TRACER_NAME, __version__)

    def _create_observation(
        self,
        otel_span: otel_trace_api.Span,
        as_type: str,
        *,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[Any] = None,
        model: Optional[str] = None,
        model_parameters: Optional[dict[str, Any]] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
    ) -> MingxSpan | MingxGeneration | MingxEvent:
        if as_type == "event":
            return MingxEvent(
                otel_span=otel_span,
                client=self,
                input=input,
                output=output,
                metadata=metadata,
                environment=self._environment,
                version=version,
                level=level,
                status_message=status_message,
            )
        if as_type == "generation":
            return MingxGeneration(
                otel_span=otel_span,
                client=self,
                input=input,
                output=output,
                metadata=metadata,
                environment=self._environment,
                version=version,
                level=level,
                status_message=status_message,
                completion_start_time=completion_start_time,
                model=model,
                model_parameters=model_parameters,
                usage_details=usage_details,
                cost_details=cost_details,
            )
        return MingxSpan(
            otel_span=otel_span,
            client=self,
            input=input,
            output=output,
            metadata=metadata,
            environment=self._environment,
            version=version,
            level=level,
            status_message=status_message,
        )

    @contextmanager
    def start_as_current_observation(
        self,
        name: str,
        *,
        as_type: str = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[Any] = None,
        model: Optional[str] = None,
        model_parameters: Optional[dict[str, Any]] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
    ):
        """Context manager: create and set current observation (span/generation/event)."""
        if as_type not in get_observation_types_list():
            as_type = "span"
        with self._tracer.start_as_current_span(name) as otel_span:
            obs = self._create_observation(
                otel_span,
                as_type,
                input=input,
                output=output,
                metadata=metadata,
                version=version,
                level=level,
                status_message=status_message,
                completion_start_time=completion_start_time,
                model=model,
                model_parameters=model_parameters,
                usage_details=usage_details,
                cost_details=cost_details,
            )
            yield obs

    def start_span(
        self,
        name: str,
        *,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ) -> MingxSpan:
        """Start a span (not set as current). Caller must call .end()."""
        otel_span = self._tracer.start_span(name)
        return self._create_observation(
            otel_span, "span", input=input, output=output, metadata=metadata, **kwargs
        )

    def start_observation(
        self,
        name: str,
        *,
        as_type: str = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[Any] = None,
        model: Optional[str] = None,
        model_parameters: Optional[dict[str, Any]] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
    ) -> MingxSpan | MingxGeneration | MingxEvent:
        """Start an observation under current context. Caller must call .end()."""
        if as_type not in get_observation_types_list():
            as_type = "span"
        otel_span = self._tracer.start_span(name)
        return self._create_observation(
            otel_span,
            as_type,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            completion_start_time=completion_start_time,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
        )

    @contextmanager
    def start_as_current_span(
        self,
        name: str,
        *,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ):
        """Context manager: create and set current span."""
        with self.start_as_current_observation(
            name, as_type="span", input=input, output=output, metadata=metadata, **kwargs
        ) as span:
            yield span

    def update_current_span(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
    ) -> None:
        """Update the current active span."""
        current = otel_trace_api.get_current_span()
        if current is None or current is otel_trace_api.INVALID_SPAN or not current.is_recording():
            return
        if name:
            current.update_name(name)
        attrs = {}
        if input is not None:
            attrs[MingxOtelSpanAttributes.OBSERVATION_INPUT] = (
                input if isinstance(input, str) else __import__("json").dumps(input, default=str)
            )
        if output is not None:
            attrs[MingxOtelSpanAttributes.OBSERVATION_OUTPUT] = (
                output if isinstance(output, str) else __import__("json").dumps(output, default=str)
            )
        if level is not None:
            attrs[MingxOtelSpanAttributes.OBSERVATION_LEVEL] = level
        if status_message is not None:
            attrs[MingxOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE] = status_message
        if attrs:
            current.set_attributes(attrs)

    def update_current_trace(
        self,
        *,
        name: Optional[str] = None,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        version: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        tags: Optional[list[str]] = None,
    ) -> None:
        """Update trace-level attributes on the current span."""
        current = otel_trace_api.get_current_span()
        if current is None or current is otel_trace_api.INVALID_SPAN or not current.is_recording():
            return
        attrs = create_trace_attributes(
            name=name,
            user_id=user_id,
            session_id=session_id,
            version=version,
            input=input,
            output=output,
            metadata=metadata,
            tags=tags,
        )
        if attrs:
            current.set_attributes(attrs)

    def create_event(
        self,
        name: str,
        *,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ) -> MingxEvent:
        """Create a point-in-time event (start_time = end_time)."""
        ts = time_ns()
        otel_span = self._tracer.start_span(name, start_time=ts)
        event = MingxEvent(
            otel_span=otel_span,
            client=self,
            input=input,
            output=output,
            metadata=metadata,
            environment=self._environment,
            **kwargs,
        )
        event.end(end_time=ts)
        return event

    def flush(self) -> None:
        """Flush pending spans to the exporter."""
        if self._processor is not None:
            self._processor.force_flush()

    def shutdown(self) -> None:
        """Shutdown and flush."""
        if self._processor is not None:
            self._processor.shutdown()
