"""Optional OpenTelemetry hooks for DataForge backend surfaces."""

from __future__ import annotations

import os
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from importlib import import_module
from typing import Any, Literal

_SENSITIVE_ATTR_FRAGMENTS = ("authorization", "cookie", "token", "key", "secret", "password")
RepairTelemetryStage = Literal[
    "detect",
    "propose",
    "safety_gate",
    "smt_verify",
    "transaction_create",
    "transaction_apply",
    "receipt",
    "revert",
]
REPAIR_TELEMETRY_SPANS: dict[str, str] = {
    "detect": "dataforge.repair.detect",
    "propose": "dataforge.repair.propose",
    "safety_gate": "dataforge.repair.safety_gate",
    "smt_verify": "dataforge.repair.smt_verify",
    "transaction_create": "dataforge.repair.transaction_create",
    "transaction_apply": "dataforge.repair.transaction_apply",
    "receipt": "dataforge.repair.receipt",
    "revert": "dataforge.repair.revert",
}
REPAIR_TELEMETRY_STAGES = frozenset(REPAIR_TELEMETRY_SPANS)


def _otel_enabled() -> bool:
    """Return whether optional OpenTelemetry instrumentation is enabled."""
    return os.environ.get("DATAFORGE_OTEL_ENABLED", "").strip().lower() in {
        "1",
        "true",
        "yes",
        "on",
    }


def _safe_attrs(attributes: dict[str, Any]) -> dict[str, str | int | float | bool]:
    """Keep only scalar, non-sensitive telemetry attributes."""
    safe: dict[str, str | int | float | bool] = {}
    for key, value in attributes.items():
        lowered = key.lower()
        if any(fragment in lowered for fragment in _SENSITIVE_ATTR_FRAGMENTS):
            continue
        if lowered in {"row_values", "rows", "payload", "source_bytes", "csv"}:
            continue
        if isinstance(value, str | int | float | bool):
            safe[key] = value
    return safe


def configure_fastapi_observability(app: Any, *, service_name: str) -> bool:
    """Instrument a FastAPI app when OpenTelemetry is explicitly enabled."""
    if not _otel_enabled():
        return False
    try:
        fastapi_instrumentation = import_module("opentelemetry.instrumentation.fastapi")
        trace_module = import_module("opentelemetry.trace")
    except ImportError:
        return False

    app.state.dataforge_service_name = service_name
    fastapi_instrumentation.FastAPIInstrumentor.instrument_app(
        app,
        tracer_provider=trace_module.get_tracer_provider(),
        excluded_urls="/api/docs,/docs,/redoc,/openapi.json",
    )
    return True


def _span_name(stage: str) -> str:
    """Return the span name for a public repair telemetry stage."""
    if stage in REPAIR_TELEMETRY_SPANS:
        return REPAIR_TELEMETRY_SPANS[stage]
    return stage


@contextmanager
def repair_stage_span(stage: RepairTelemetryStage | str, **attributes: Any) -> Iterator[None]:
    """Create a repair-stage span when OpenTelemetry is available."""
    if not _otel_enabled():
        with nullcontext():
            yield
        return

    try:
        trace_module = import_module("opentelemetry.trace")
    except ImportError:
        with nullcontext():
            yield
        return

    tracer = trace_module.get_tracer("dataforge.repair")
    with tracer.start_as_current_span(_span_name(stage)) as span:
        if stage in REPAIR_TELEMETRY_SPANS:
            span.set_attribute("dataforge.repair.stage", stage)
        for key, value in _safe_attrs(attributes).items():
            span.set_attribute(key, value)
        yield