"""Telemetry gateway — single routing layer for all editions.

Maps ``TelemetryCase`` → ``CaseRoute`` and dispatches events to either
the CE/EE trace pipeline (``TraceQueueManager``) or the enterprise-only
metric/log Celery queue.

This module lives in ``core/`` so both CE and EE share one routing table
and one ``emit()`` entry point.  No separate enterprise gateway module is
needed — enterprise-specific dispatch (Celery task, payload offloading)
is handled here behind lazy imports that no-op in CE.
"""

from __future__ import annotations

import json
import logging
import uuid
from typing import TYPE_CHECKING, Any

from core.ops.entities.trace_entity import TraceTaskName
from enterprise.telemetry.contracts import CaseRoute, SignalType
from extensions.ext_storage import storage

if TYPE_CHECKING:
    from core.ops.ops_trace_manager import TraceQueueManager
    from enterprise.telemetry.contracts import TelemetryCase

logger = logging.getLogger(__name__)

PAYLOAD_SIZE_THRESHOLD_BYTES = 1 * 1024 * 1024

# ---------------------------------------------------------------------------
# Routing table — authoritative mapping for all editions
# ---------------------------------------------------------------------------

_case_to_trace_task: dict[TelemetryCase, TraceTaskName] | None = None
_case_routing: dict[TelemetryCase, CaseRoute] | None = None


def _get_case_to_trace_task() -> dict[TelemetryCase, TraceTaskName]:
    global _case_to_trace_task
    if _case_to_trace_task is None:
        from enterprise.telemetry.contracts import TelemetryCase

        _case_to_trace_task = {
            TelemetryCase.WORKFLOW_RUN: TraceTaskName.WORKFLOW_TRACE,
            TelemetryCase.MESSAGE_RUN: TraceTaskName.MESSAGE_TRACE,
            TelemetryCase.NODE_EXECUTION: TraceTaskName.NODE_EXECUTION_TRACE,
            TelemetryCase.DRAFT_NODE_EXECUTION: TraceTaskName.DRAFT_NODE_EXECUTION_TRACE,
            TelemetryCase.PROMPT_GENERATION: TraceTaskName.PROMPT_GENERATION_TRACE,
            TelemetryCase.TOOL_EXECUTION: TraceTaskName.TOOL_TRACE,
            TelemetryCase.MODERATION_CHECK: TraceTaskName.MODERATION_TRACE,
            TelemetryCase.SUGGESTED_QUESTION: TraceTaskName.SUGGESTED_QUESTION_TRACE,
            TelemetryCase.DATASET_RETRIEVAL: TraceTaskName.DATASET_RETRIEVAL_TRACE,
            TelemetryCase.GENERATE_NAME: TraceTaskName.GENERATE_NAME_TRACE,
        }
    return _case_to_trace_task


def get_trace_task_to_case() -> dict[TraceTaskName, TelemetryCase]:
    """Return TraceTaskName → TelemetryCase (inverse of _get_case_to_trace_task)."""
    return {v: k for k, v in _get_case_to_trace_task().items()}


def _get_case_routing() -> dict[TelemetryCase, CaseRoute]:
    global _case_routing
    if _case_routing is None:
        from enterprise.telemetry.contracts import CaseRoute, SignalType, TelemetryCase

        _case_routing = {
            # TRACE — CE-eligible (flow in both CE and EE)
            TelemetryCase.WORKFLOW_RUN: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.MESSAGE_RUN: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.TOOL_EXECUTION: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.MODERATION_CHECK: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.SUGGESTED_QUESTION: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.DATASET_RETRIEVAL: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            TelemetryCase.GENERATE_NAME: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=True),
            # TRACE — enterprise-only
            TelemetryCase.NODE_EXECUTION: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=False),
            TelemetryCase.DRAFT_NODE_EXECUTION: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=False),
            TelemetryCase.PROMPT_GENERATION: CaseRoute(signal_type=SignalType.TRACE, ce_eligible=False),
            # METRIC_LOG — enterprise-only (signal-driven, not trace)
            TelemetryCase.APP_CREATED: CaseRoute(signal_type=SignalType.METRIC_LOG, ce_eligible=False),
            TelemetryCase.APP_UPDATED: CaseRoute(signal_type=SignalType.METRIC_LOG, ce_eligible=False),
            TelemetryCase.APP_DELETED: CaseRoute(signal_type=SignalType.METRIC_LOG, ce_eligible=False),
            TelemetryCase.FEEDBACK_CREATED: CaseRoute(signal_type=SignalType.METRIC_LOG, ce_eligible=False),
        }
    return _case_routing


def __getattr__(name: str) -> Any:
    """Lazy module-level access to routing tables."""
    if name == "CASE_ROUTING":
        return _get_case_routing()
    if name == "CASE_TO_TRACE_TASK":
        return _get_case_to_trace_task()
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def is_enterprise_telemetry_enabled() -> bool:
    try:
        from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled

        return is_enterprise_telemetry_enabled()
    except Exception:
        return False


def _handle_payload_sizing(
    payload: dict[str, Any],
    tenant_id: str,
    event_id: str,
) -> tuple[dict[str, Any], str | None]:
    """Inline or offload payload based on size.

    Returns ``(payload_for_envelope, storage_key | None)``.  Payloads
    exceeding ``PAYLOAD_SIZE_THRESHOLD_BYTES`` are written to object
    storage and replaced with an empty dict in the envelope.
    """
    try:
        payload_json = json.dumps(payload)
        payload_size = len(payload_json.encode("utf-8"))
    except (TypeError, ValueError):
        logger.warning("Failed to serialize payload for sizing: event_id=%s", event_id)
        return payload, None

    if payload_size <= PAYLOAD_SIZE_THRESHOLD_BYTES:
        return payload, None

    storage_key = f"telemetry/{tenant_id}/{event_id}.json"
    try:
        storage.save(storage_key, payload_json.encode("utf-8"))
        logger.debug("Stored large payload to storage: key=%s, size=%d", storage_key, payload_size)
        return {}, storage_key
    except Exception:
        logger.warning("Failed to store large payload, inlining instead: event_id=%s", event_id, exc_info=True)
        return payload, None


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def emit(
    case: TelemetryCase,
    context: dict[str, Any],
    payload: dict[str, Any],
    trace_manager: TraceQueueManager | None = None,
) -> None:
    """Route a telemetry event to the correct pipeline.

    TRACE events are enqueued into ``TraceQueueManager`` (works in both CE
    and EE).  Enterprise-only traces are silently dropped when EE is
    disabled.

    METRIC_LOG events are dispatched to the enterprise Celery queue;
    silently dropped when enterprise telemetry is unavailable.
    """
    route = _get_case_routing().get(case)
    if route is None:
        logger.warning("Unknown telemetry case: %s, dropping event", case)
        return

    if not route.ce_eligible and not is_enterprise_telemetry_enabled():
        logger.debug("Dropping EE-only event: case=%s (EE disabled)", case)
        return

    if route.signal_type == SignalType.TRACE:
        _emit_trace(case, context, payload, trace_manager)
    else:
        _emit_metric_log(case, context, payload)


def _emit_trace(
    case: TelemetryCase,
    context: dict[str, Any],
    payload: dict[str, Any],
    trace_manager: TraceQueueManager | None,
) -> None:
    from core.ops.ops_trace_manager import TraceQueueManager as LocalTraceQueueManager
    from core.ops.ops_trace_manager import TraceTask

    trace_task_name = _get_case_to_trace_task().get(case)
    if trace_task_name is None:
        logger.warning("No TraceTaskName mapping for case: %s", case)
        return

    queue_manager = trace_manager or LocalTraceQueueManager(
        app_id=context.get("app_id"),
        user_id=context.get("user_id"),
    )
    queue_manager.add_trace_task(TraceTask(trace_task_name, user_id=context.get("user_id"), **payload))
    logger.debug("Enqueued trace task: case=%s, app_id=%s", case, context.get("app_id"))


def _emit_metric_log(
    case: TelemetryCase,
    context: dict[str, Any],
    payload: dict[str, Any],
) -> None:
    """Build envelope and dispatch to enterprise Celery queue.

    No-ops when the enterprise telemetry task is not importable (CE mode).
    """
    try:
        from tasks.enterprise_telemetry_task import process_enterprise_telemetry
    except ImportError:
        logger.debug("Enterprise metric/log dispatch unavailable, dropping: case=%s", case)
        return

    tenant_id = context.get("tenant_id") or ""
    event_id = str(uuid.uuid4())

    payload_for_envelope, payload_ref = _handle_payload_sizing(payload, tenant_id, event_id)

    from enterprise.telemetry.contracts import TelemetryEnvelope

    envelope = TelemetryEnvelope(
        case=case,
        tenant_id=tenant_id,
        event_id=event_id,
        payload=payload_for_envelope,
        metadata={"payload_ref": payload_ref} if payload_ref else None,
    )

    process_enterprise_telemetry.delay(envelope.model_dump_json())
    logger.debug(
        "Enqueued metric/log event: case=%s, tenant_id=%s, event_id=%s",
        case,
        tenant_id,
        event_id,
    )
