"""Telemetry initialization — TracerProvider + APM exporter setup."""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Any, Optional

from opentelemetry import trace as otel_trace
from opentelemetry.context import Context, get_current
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import Span, Tracer, set_span_in_context

from .apm.config import resolve_apm_endpoint, resolve_service_name
from .apm.metrics_bridge import ApmMetricsBridge
from .apm.span_exporter import ApmSpanExporter, TranslatorContext
from ._compat import apply_compat_patches
from .context_propagator import AgentContextPropagator


@dataclass
class Telemetry:
    tracer_provider: TracerProvider
    tracer: Tracer
    metrics_bridge: ApmMetricsBridge
    _root_span_ids: set = field(default_factory=set)
    _root_spans_with_children: set = field(default_factory=set)

    def start_request_span(
        self, name: str, attributes: dict[str, Any] | None = None
    ) -> tuple[Span, Context]:
        span = self.tracer.start_span(name, attributes=attributes or {})
        self._root_span_ids.add(format(span.get_span_context().span_id, "016x"))
        ctx = set_span_in_context(span, get_current())
        return span, ctx

    def end_request_span(self, span: Span) -> None:
        span_id = format(span.get_span_context().span_id, "016x")
        span.end()
        # Flush pending spans BEFORE removing the root span ID.
        # BatchSpanProcessor exports asynchronously — child spans queued
        # earlier need the root ID still present when _is_trace_entry runs
        # inside the ApmSpanExporter so they get gen_ai.is_entry=true.
        try:
            self.tracer_provider.force_flush(timeout_millis=5000)
        except Exception:
            pass
        self._root_span_ids.discard(span_id)
        self._root_spans_with_children.discard(span_id)

    def shutdown(self) -> None:
        self._root_span_ids.clear()
        self._root_spans_with_children.clear()
        try:
            self.metrics_bridge.shutdown()
        except Exception:
            pass
        try:
            self.tracer_provider.shutdown()
        except Exception:
            pass


_telemetry: Telemetry | None = None


def setup_telemetry(
    *,
    tracer_provider: TracerProvider | None = None,
    endpoint: str | None = None,
    service_name: str | None = None,
    apm_token: str | None = None,
) -> Telemetry:
    """Initialize telemetry. Called once at process startup.

    Args:
        tracer_provider: Override for testing (InMemoryExporter).
                         If None, creates production provider with APM exporter.
    """
    global _telemetry
    if _telemetry is not None:
        return _telemetry

    apply_compat_patches()

    endpoint = endpoint or os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") or resolve_apm_endpoint()
    service_name = service_name or resolve_service_name()
    environment = os.environ.get("NODE_ENV") or "development"
    apm_token = apm_token or os.environ.get("SVC_TRACE_KEY_LB6V03", "")

    # Dev dual-export: optionally send to remote APM alongside local Studio.
    dev_apm_token = os.environ.get("EDGEONE_DEV_APM_TOKEN", "")
    dev_apm_endpoint = os.environ.get("EDGEONE_DEV_APM_ENDPOINT") or resolve_apm_endpoint()

    # 使用 Resource.create() 以保留 SDK 默认属性（telemetry.sdk.language 等），
    # 否则 inferRuntime() 无法识别 Python runtime，会误判为 agent-node。
    # Resource.create() 会自动合并 OTEL_RESOURCE_ATTRIBUTES 和 SDK 默认属性。
    resource = Resource.create({
        "service.name": service_name,
        "deployment.environment": environment,
        "token": dev_apm_token or apm_token,
    })

    root_span_ids: set[str] = set()
    root_spans_with_children: set[str] = set()

    if tracer_provider is None:
        translator_ctx = TranslatorContext(
            is_root_span_id=lambda sid: sid in root_span_ids,
            root_span_has_children=lambda sid: sid in root_spans_with_children,
        )
        raw_exporter = OTLPSpanExporter(endpoint=f"{endpoint.rstrip('/')}/v1/traces")
        wrapped_exporter = ApmSpanExporter(raw_exporter, translator_ctx)

        tracer_provider = TracerProvider(resource=resource)
        tracer_provider.add_span_processor(BatchSpanProcessor(wrapped_exporter))

        # Dev APM dual-export processor (if EDGEONE_DEV_APM_TOKEN is set)
        if dev_apm_token:
            dev_apm_raw = OTLPSpanExporter(endpoint=f"{dev_apm_endpoint.rstrip('/')}/v1/traces")
            dev_apm_wrapped = ApmSpanExporter(dev_apm_raw, translator_ctx)
            tracer_provider.add_span_processor(BatchSpanProcessor(dev_apm_wrapped))
            print(f"[observability] APM dual-export enabled → {dev_apm_endpoint}")

        metrics_bridge = ApmMetricsBridge(
            endpoint=dev_apm_endpoint if dev_apm_token else endpoint,
            apm_token=dev_apm_token or apm_token,
            resource=resource,
        )
        tracer_provider.add_span_processor(metrics_bridge.span_processor)

        # Inject agent.run_id (= traceId) and agent.conversation_id into every span.
        # Mirrors Node's AgentContextPropagator so langfuse.agent.run_id aligns with traceId.
        tracer_provider.add_span_processor(AgentContextPropagator())

        # Track which root spans have children (for gen_ai.is_entry on root span)
        tracer_provider.add_span_processor(
            _RootChildrenTracker(root_span_ids, root_spans_with_children)
        )

        otel_trace.set_tracer_provider(tracer_provider)
    else:
        # Test mode — no APM exporter, no metrics bridge, don't touch global state
        metrics_bridge = _NoOpMetricsBridge()
        # 测试模式下也注册 propagator，让 InMemoryExporter 能验证注入行为。
        try:
            tracer_provider.add_span_processor(AgentContextPropagator())
        except Exception:
            # 某些测试 provider 可能不支持重复添加，忽略即可。
            pass

    tracer = tracer_provider.get_tracer("edgeone-agent-observability", "0.1.0")

    _telemetry = Telemetry(
        tracer_provider=tracer_provider,
        tracer=tracer,
        metrics_bridge=metrics_bridge,
        _root_span_ids=root_span_ids,
        _root_spans_with_children=root_spans_with_children,
    )
    return _telemetry


def get_telemetry() -> Telemetry | None:
    return _telemetry


def reset_telemetry() -> None:
    """Reset global state. For testing only."""
    global _telemetry
    _telemetry = None


class _NoOpMetricsBridge:
    """Stub for test mode."""
    @property
    def span_processor(self):
        return None

    def shutdown(self) -> None:
        pass


class _RootChildrenTracker:
    """SpanProcessor that marks root spans as having children when a child starts."""

    def __init__(self, root_span_ids: set[str], root_spans_with_children: set[str]) -> None:
        self._root_span_ids = root_span_ids
        self._root_spans_with_children = root_spans_with_children

    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
        parent = span.parent if hasattr(span, "parent") else None
        if parent is None:
            return
        try:
            parent_id = format(parent.span_id, "016x")
        except (AttributeError, TypeError):
            return
        if parent_id in self._root_span_ids:
            self._root_spans_with_children.add(parent_id)

    def _on_ending(self, span) -> None:
        pass

    def on_end(self, span) -> None:
        pass

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True