"""OpenTelemetry setup and the GenAI instruments used everywhere else (lessons 16 and 17).

    telemetry = Telemetry.setup("docs-copilot")    # OTLP if OTEL_EXPORTER_OTLP_ENDPOINT is set,
                                                   # otherwise in-memory (tests, the offline demo)
    with telemetry.tracer.start_as_current_span("POST /ask"): ...

Span and metric names follow the OpenTelemetry GenAI semantic conventions (status: Development).
Keeping them in this one module means a convention rename is a one-line change.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

from opentelemetry import metrics, trace
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import InMemoryMetricReader, MetricReader
from opentelemetry.sdk.metrics.view import ExplicitBucketHistogramAggregation, View
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

TOKEN_BUCKETS = [1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864]
DURATION_BUCKETS = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92]

_globals_set = False


@dataclass
class Telemetry:
    tracer_provider: TracerProvider
    meter_provider: MeterProvider
    spans: InMemorySpanExporter | None           # set only in in-memory mode
    metric_reader: MetricReader

    @classmethod
    def setup(cls, service_name: str = "docs-copilot", otlp: bool | None = None) -> "Telemetry":
        global _globals_set
        otlp = bool(os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")) if otlp is None else otlp
        resource = Resource.create({"service.name": os.environ.get("OTEL_SERVICE_NAME", service_name),
                                    "service.version": "1.0.0"})
        tracer_provider = TracerProvider(resource=resource)
        if otlp:                                  # production: batch export to a Collector
            from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
            from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
            tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))
            reader: MetricReader = PeriodicExportingMetricReader(OTLPMetricExporter(), export_interval_millis=15_000)
            spans = None
        else:                                     # tests and the demo: keep everything in memory
            spans = InMemorySpanExporter()
            tracer_provider.add_span_processor(SimpleSpanProcessor(spans))
            reader = InMemoryMetricReader()
        meter_provider = MeterProvider(resource=resource, metric_readers=[reader], views=[
            View(instrument_name="gen_ai.client.token.usage",
                 aggregation=ExplicitBucketHistogramAggregation(TOKEN_BUCKETS)),
            View(instrument_name="gen_ai.client.operation.duration",
                 aggregation=ExplicitBucketHistogramAggregation(DURATION_BUCKETS)),
            View(instrument_name="app.request.duration",
                 aggregation=ExplicitBucketHistogramAggregation(DURATION_BUCKETS))])
        if not _globals_set:                      # global providers can only be set once per process
            trace.set_tracer_provider(tracer_provider)
            metrics.set_meter_provider(meter_provider)
            _globals_set = True
        return cls(tracer_provider, meter_provider, spans, reader)

    def __post_init__(self) -> None:
        self.tracer = self.tracer_provider.get_tracer("copilot")
        meter = self.meter_provider.get_meter("copilot")
        self.token_usage = meter.create_histogram("gen_ai.client.token.usage", unit="{token}",
                                                  description="Number of input and output tokens used.")
        self.operation_duration = meter.create_histogram("gen_ai.client.operation.duration", unit="s",
                                                         description="GenAI operation duration.")
        self.request_duration = meter.create_histogram("app.request.duration", unit="s")
        self.answers = meter.create_counter("app.answers", unit="{answer}",
                                            description="Answers, by grounded=true/false")
        self.cost = meter.create_counter("app.cost_usd", unit="{USD}",      # Prometheus: app_cost_usd_total
                                         description="Model spend in US dollars")
        self.guardrail_hits = meter.create_counter("app.guardrail.hits", unit="{hit}")

    def shutdown(self) -> None:
        """Flush and stop. Call before a short-lived process exits (lesson 16)."""
        self.tracer_provider.shutdown()
        self.meter_provider.shutdown()

    # ------------------------------------------------------------------ inspection helpers (in-memory mode)
    def finished_spans(self) -> list:
        return list(self.spans.get_finished_spans()) if self.spans else []

    def metric_points(self) -> dict[str, list]:
        out: dict[str, list] = {}
        data = self.metric_reader.get_metrics_data() if isinstance(self.metric_reader, InMemoryMetricReader) else None
        for rm in (data.resource_metrics if data else []):
            for sm in rm.scope_metrics:
                for m in sm.metrics:
                    out[m.name] = list(m.data.data_points)
        return out