#!/usr/bin/env python
from __future__ import annotations

import logging
import os
import sys
from importlib.metadata import PackageNotFoundError, version as pkg_version

from snowflake.connector.telemetry import (
    TelemetryClient as PCTelemetryClient,
    TelemetryData as PCTelemetryData,
)
from snowflake.connector.time_util import get_time_millis
from snowflake.snowpark import Session

#
# Lightweight telemetry helper for snowpark-submit.
#
# Scope: emit a few events (invocation, completion, cli_operation) with minimal
# fields and synchronous delivery through the Snowflake connector telemetry
# client when available. If no telemetry client is present or telemetry is
# disabled, this becomes a no-op.
#


logger = logging.getLogger("snowpark-submit")

DISABLE_ENV = "SNOWPARK_SUBMIT_DISABLE_TELEMETRY"
EVENT_INVOCATION = "scos_submit_invocation"
EVENT_COMPLETION = "scos_submit_completion"
EVENT_CLI_OPERATION = "scos_submit_cli_operation"
ERROR_SNIPPET_LIMIT = 120

try:
    CLI_VERSION = pkg_version("snowpark-submit")
except PackageNotFoundError:
    CLI_VERSION = "unknown"


def _extract_telemetry_client(session: Session) -> PCTelemetryClient | None:
    """Drill down session internals to find the connector telemetry client."""
    server_connection = getattr(session, "_conn", None)
    if not server_connection:
        return None
    snowflake_connection = getattr(server_connection, "_conn", None)
    if not snowflake_connection:
        return None
    return getattr(snowflake_connection, "_telemetry", None)


def _trim(text: str, limit: int = ERROR_SNIPPET_LIMIT) -> str:
    return text if len(text) <= limit else text[: limit - 3] + "..."


class _NoOpSink:
    def add(self, payload: dict, timestamp: int) -> None:  # pragma: no cover
        return

    def flush(self) -> None:  # pragma: no cover
        return


class _ClientSink:
    def __init__(self, telemetry_client: PCTelemetryClient) -> None:
        self._client = telemetry_client

    def add(self, payload: dict, timestamp: int) -> None:
        self._client.try_add_log_to_batch(
            PCTelemetryData(message=payload, timestamp=timestamp)
        )

    def flush(self) -> None:
        self._client.send_batch()


class SubmitTelemetry:
    def __init__(self) -> None:
        self._enabled = DISABLE_ENV not in os.environ
        self._initialized = False
        self._sink = _NoOpSink()
        self._static: dict = {}
        # Idempotency guard: invocation event should only be sent once per run.
        self._invocation_sent = False

    def initialize(
        self,
        session: Session,
        source: str = "snowpark-submit",
        compute_pool: str | None = None,
    ) -> None:
        if not self._enabled:
            return
        # Idempotent: safe to call multiple times.
        if self._initialized:
            return
        self._initialized = True

        telemetry_client = _extract_telemetry_client(session)
        if not telemetry_client:
            self._enabled = False
            return

        self._sink = _ClientSink(telemetry_client)

        py_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
        self._static = {
            "source": source,
            "cli_version": CLI_VERSION,
            "python_version": py_version,
            "os": os.name,
            "account": self._safe(session.get_current_account),
            "user": self._safe(session.get_current_user),
            "role": self._safe(session.get_current_role),
            "warehouse": self._safe(session.get_current_warehouse),
            "compute_pool": compute_pool,
        }

    def _safe(self, fn, key: str | None = None) -> str | None:
        try:
            return fn() if key is None else fn().get(key)
        except Exception:
            return None

    def send_invocation(
        self,
        workload_name: str | None,
        operation: str,
        language: str | None = None,
        correlation_id: str | None = None,
    ) -> None:
        """Emit invocation telemetry once per run. Idempotent.

        ``language`` is passed here (not at initialize-time) because the workload
        language is only known after the submitted filename is inspected, which
        happens after telemetry is initialized.
        """
        if not self._enabled or not self._initialized:
            return
        if self._invocation_sent:
            return
        self._invocation_sent = True
        payload = {
            **self._static,
            "type": EVENT_INVOCATION,
            "workload_name": workload_name,
            "operation": operation,
        }
        if language:
            payload["language"] = language
        if correlation_id:
            payload["correlation_id"] = correlation_id
        self._sink.add(payload, get_time_millis())
        self._sink.flush()

    def send_completion(
        self,
        workload_name: str | None,
        workload_status: str | None,
        job_exit_code: int | None,
        error: str | None = None,
        correlation_id: str | None = None,
    ) -> None:
        if not self._enabled or not self._initialized:
            return

        payload = {
            **self._static,
            "type": EVENT_COMPLETION,
            "workload_name": workload_name,
            "workload_status": workload_status,
            "job_exit_code": job_exit_code,
        }

        if error:
            payload["error"] = _trim(error)
        if correlation_id:
            payload["correlation_id"] = correlation_id

        self._sink.add(payload, get_time_millis())
        self._sink.flush()

    def send_cli_operation(
        self,
        operation: str,
        workload_name: str | None,
        exit_code: int,
        error: str | None = None,
        correlation_id: str | None = None,
    ) -> None:
        """
        Send telemetry for CLI operations like status, kill, list.
        These are NOT job completions - they're CLI command results.
        """
        if not self._enabled or not self._initialized:
            return

        payload = {
            **self._static,
            "type": EVENT_CLI_OPERATION,
            "operation": operation,
            "workload_name": workload_name,
            "exit_code": exit_code,
        }

        if error:
            payload["error"] = _trim(error)
        if correlation_id:
            payload["correlation_id"] = correlation_id

        self._sink.add(payload, get_time_millis())
        self._sink.flush()


# Global instance
telemetry = SubmitTelemetry()
