// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "telemetry.h"
#include "device_info.h"
#include "telemetry_context.h"
#include "telemetry_environment.h"
#include "telemetry_redaction.h"
#include "telemetry_sampling.h"
#include "../logging.h"

#if defined(ORTGENAI_ENABLE_TELEMETRY)
#include "Enums.hpp"
#include "ILogConfiguration.hpp"
#include "LogManager.hpp"
#include "LogManagerProvider.hpp"
// Generated by CMake into the build tree. Defines ORTGENAI_TELEMETRY_TENANT_TOKEN only when an override
// was supplied via CMake or environment; keeps the token off the compiler command line.
#include "ortgenai_telemetry_tenant_token.h"
#endif

#include <cstdio>
#include <filesystem>
#include <random>
#include <string>

// Version string defined by the build system
#ifndef ORTGENAI_VERSION
#define ORTGENAI_VERSION "unknown"
#endif

#if defined(ORTGENAI_ENABLE_TELEMETRY) && defined(__ANDROID__)
extern "C" bool OrtGenAIIsAndroidTelemetryReady() noexcept;
#endif

namespace Generators {

namespace {

#if defined(ORTGENAI_ENABLE_TELEMETRY)

// Decode a base64-encoded string at runtime.
std::string DecodeBase64(const std::string& encoded) {
  auto DecodeChar = [](char c) -> int {
    if (c >= 'A' && c <= 'Z') return c - 'A';
    if (c >= 'a' && c <= 'z') return c - 'a' + 26;
    if (c >= '0' && c <= '9') return c - '0' + 52;
    if (c == '+') return 62;
    if (c == '/') return 63;
    return -1;
  };

  std::string result;
  result.reserve(encoded.size() * 3 / 4);
  uint32_t accum = 0;
  int bits = 0;
  for (char c : encoded) {
    int val = DecodeChar(c);
    if (val < 0) continue;
    accum = (accum << 6) | static_cast<uint32_t>(val);
    bits += 6;
    if (bits >= 8) {
      bits -= 8;
      result.push_back(static_cast<char>((accum >> bits) & 0xFF));
    }
  }
  return result;
}

std::string GetToken() {
  static constexpr char kMaskKey[] = "OnnxRuntimeGenAI";
  constexpr size_t klen = sizeof(kMaskKey) - 1;
  std::string decoded = DecodeBase64(
      "LA1WHDYXWhILD1FwUQhyLXdfX0g3RFoRX19Vf1IPcnFiCFcaZUAIEgxAUH5cX2x9eQpYVTNHC0REVVclVFx4fC0LDU9hWFlHXlo=");
  for (size_t i = 0; i < decoded.size(); ++i)
    decoded[i] = static_cast<char>(decoded[i] ^ kMaskKey[i % klen]);
  return decoded;
}

bool PrepareSampledEvent(MAT::EventProperties& event, std::string_view app_session_guid,
                         uint32_t session_id) {
  if (!TelemetryInternal::ShouldSampleSession(app_session_guid, session_id)) return false;
  event.SetPopsample(TelemetryInternal::kModelSessionSampleRatePercent);
  return true;
}

// Generate a random v4 UUID as a hex string (e.g. "f81d4fae-7dec-41d0-8f12-00a0c91e6bf6").
// Used for the process-wide AppSessionGuid (Tier 1 identity).
std::string GenerateGuidV4() {
  // Draw directly from std::random_device rather than seeding a PRNG from one sample.
  std::random_device rd;
  uint64_t hi = (static_cast<uint64_t>(rd()) << 32) | rd();
  uint64_t lo = (static_cast<uint64_t>(rd()) << 32) | rd();
  // Set version (4) and variant (10xx) bits.
  hi = (hi & 0xFFFFFFFFFFFF0FFFULL) | 0x0000000000004000ULL;
  lo = (lo & 0x3FFFFFFFFFFFFFFFULL) | 0x8000000000000000ULL;

  char buf[37];
  std::snprintf(buf, sizeof(buf),
                "%08x-%04x-%04x-%04x-%012llx",
                static_cast<uint32_t>(hi >> 32),
                static_cast<uint32_t>((hi >> 16) & 0xFFFF),
                static_cast<uint32_t>(hi & 0xFFFF),
                static_cast<uint32_t>(lo >> 48),
                static_cast<unsigned long long>(lo & 0xFFFFFFFFFFFFULL));
  return std::string(buf);
}

#endif  // ORTGENAI_ENABLE_TELEMETRY

}  // namespace

#if defined(ORTGENAI_ENABLE_TELEMETRY)

struct GenAiTelemetry::Impl {
  MAT::ILogConfiguration config;  // must outlive log_manager (SDK holds a reference)
  MAT::ILogManager* log_manager{nullptr};
  MAT::ILogger* logger{nullptr};
};

#else

struct GenAiTelemetry::Impl {};

#endif  // ORTGENAI_ENABLE_TELEMETRY

// Guard against access after static destruction
static std::atomic<bool> s_instance_destroyed{false};

GenAiTelemetry& GenAiTelemetry::Instance() {
  static GenAiTelemetry instance;
  return instance;
}

bool GenAiTelemetry::IsDestroyed() {
  return s_instance_destroyed.load();
}

void GenAiTelemetry::Initialize() {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  std::unique_lock<std::shared_mutex> lock(mutex_);
  if (initialized_.load()) return;

  // Full suppression is process-wide and irreversible: never create the uploader, emit an event, or
  // persist a device id. Latching also covers later GenAI reinitialization after an environment change.
  if (telemetry_disabled_.load() ||
      TelemetryInternal::IsRunningInCI() ||
      TelemetryInternal::IsRunningUnitTests() ||
      TelemetryInternal::IsTelemetryDisabledByEnvironment()) {
    telemetry_disabled_.store(true);
    enabled_.store(false);
    return;
  }

#if defined(__ANDROID__)
  // The Android SDK requires its Java HTTP bridge before native initialization. AAR builds install
  // it automatically; native-only consumers without a Java context safely run without telemetry.
  if (!OrtGenAIIsAndroidTelemetryReady()) {
    const bool should_warn =
        g_log.enabled && g_log.warning && !android_readiness_warning_logged_.exchange(true);
    lock.unlock();
    if (should_warn) {
      try {
        Log("warning", "Android telemetry is waiting for the 1DS Java HttpClient");
      } catch (...) {
        // Diagnostics must not affect host control flow.
      }
    }
    return;
  }
#endif

  // Telemetry must never affect host control flow: any SDK failure here is
  // swallowed and simply leaves telemetry uninitialized (a later call retries).
  // enabled_ is not cleared on failure paths, so a successful retry still emits.
  std::unique_ptr<Impl> pending_impl;
  const auto release_pending = [&pending_impl]() noexcept {
    if (pending_impl && pending_impl->log_manager) {
      try {
        MAT::LogManagerProvider::Release(pending_impl->config);
      } catch (...) {
      }
      pending_impl->log_manager = nullptr;
      pending_impl->logger = nullptr;
    }
    pending_impl.reset();
  };

  try {
    // Ingestion token: an official build's override from the generated header
    // (ORTGENAI_TELEMETRY_TENANT_TOKEN) if present, otherwise the encoded in-repo default.
#if defined(ORTGENAI_TELEMETRY_TENANT_TOKEN)
    const std::string ikey = ORTGENAI_TELEMETRY_TENANT_TOKEN;
#else
    const std::string ikey = GetToken();
#endif
    if (ikey.empty()) return;  // no token -> cannot send; stay uninitialized

    pending_impl = std::make_unique<Impl>();

    // Configuration is owned by the pending Impl so it outlives the log manager (the SDK holds a reference to it).
    auto& config = pending_impl->config;
    config[MAT::CFG_STR_COLLECTOR_URL] = "https://mobile.events.data.microsoft.com/OneCollector/1.0";
    config[MAT::CFG_STR_PRIMARY_TOKEN] = ikey;
    config[MAT::CFG_INT_TRACE_LEVEL_MASK] = 0;
    config[MAT::CFG_INT_SDK_MODE] = MAT::SdkModeTypes::SdkModeTypes_CS;
    config[MAT::CFG_INT_RAM_QUEUE_SIZE] = 512 * 1024;

    // Do not block process teardown to upload: persisted events are sent on the next
    // run. 0 keeps Shutdown non-blocking and avoids adding exit latency to host apps.
    config[MAT::CFG_INT_MAX_TEARDOWN_TIME] = 0;

#if !defined(__ANDROID__)
    const std::filesystem::path cache_dir = GetTelemetryStorageDir();
    if (!cache_dir.empty()) {
      // Product-specific cache file name: the .onnxruntime directory is shared with ONNX Runtime (and
      // Olive) for the device id, but each product needs its OWN 1DS offline cache so two log managers
      // with different tenant tokens never share one SQLite db (which could mix or misroute events).
      config[MAT::CFG_STR_CACHE_FILE_PATH] = (cache_dir / "onnxruntime-genai.db").string();
    }
#endif

    // Create a dedicated log manager via LogManagerProvider (recommended over the
    // LOGMANAGER_INSTANCE singleton for library use). wantController=true makes this
    // instance the host (name == host), so Flush/FlushAndTeardown act on it.
    MAT::status_t status = MAT::STATUS_SUCCESS;
    pending_impl->log_manager =
        MAT::LogManagerProvider::CreateLogManager("OnnxRuntimeGenAI", true, config, status);
    if (status != MAT::STATUS_SUCCESS || pending_impl->log_manager == nullptr) {
      release_pending();
      return;
    }
    pending_impl->logger = pending_impl->log_manager->GetLogger(ikey);
    if (pending_impl->logger == nullptr) {
      release_pending();
      return;
    }
    (void)TelemetryInternal::TrySuppressContext([&]() {
      TelemetryInternal::SuppressUnneededCommonContext(
          *pending_impl->logger->GetSemanticContext());
    });
    bool network_context_suppressed = false;
    if (process_info_logged_.load()) {
      // ProcessInfo is once per process; a reinitialized logger no longer needs network context.
      network_context_suppressed = TelemetryInternal::TrySuppressContext([&]() {
        TelemetryInternal::SuppressNetworkContext(
            *pending_impl->logger->GetSemanticContext());
      });
    }
    network_context_suppressed_.store(network_context_suppressed);
    pending_impl->log_manager->SetTransmitProfile(MAT::TransmitProfile_BestEffort);

    // Process-wide AppSessionGuid stamped on every event as logger context (not a
    // per-event property). An explicit, SDK-independent correlator that answers
    // "which events came from the same process run?" and makes the lightweight
    // per-event sessionId / generatorId counters globally unique.
    if (app_session_guid_.empty()) app_session_guid_ = GenerateGuidV4();
    pending_impl->logger->SetContext("AppSessionGuid", app_session_guid_);
    pending_impl->logger->SetContext("LibraryVersion", ORTGENAI_VERSION);

    const auto& device = GetDeviceInfo();
    // Desktop builds override the 1DS SDK's device id with a privacy-preserving generated-id hash.
    // Android/iOS leave this empty so the SDK uses its platform device id.
    if (!device.device_id.empty()) {
      pending_impl->logger->GetSemanticContext()->SetDeviceId(device.device_id);
    }

    impl_ = std::move(pending_impl);
    initialized_.store(true);
  } catch (...) {
    // Never propagate a telemetry-init failure to the host.
    release_pending();
  }
#endif
}

void GenAiTelemetry::Shutdown() {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  std::unique_lock<std::shared_mutex> lock(mutex_);
  if (!initialized_.load()) return;

  // The exclusive lock waits for active LogEvent calls to release their shared locks and prevents
  // new ones from observing the logger while teardown is in progress.
  initialized_.store(false);

  // 1DS-recommended shutdown sequence. Each operation is independent and best effort so Release()
  // still runs if flushing or teardown fails.
  if (impl_ && impl_->log_manager) {
    try {
      impl_->log_manager->Flush();
    } catch (...) {
    }
    try {
      impl_->log_manager->FlushAndTeardown();
    } catch (...) {
    }
    try {
      MAT::LogManagerProvider::Release(impl_->config);
    } catch (...) {
    }
    impl_->log_manager = nullptr;
    impl_->logger = nullptr;
  }
  impl_.reset();
#endif
}
GenAiTelemetry::~GenAiTelemetry() {
  s_instance_destroyed.store(true);
  Shutdown();
}

bool GenAiTelemetry::IsEnabled() const {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  return enabled_.load() && initialized_.load();
#else
  return false;
#endif
}

void GenAiTelemetry::SetEnabled(bool enabled) {
  if (!telemetry_disabled_.load()) {
    enabled_.store(enabled);
  }
}

uint32_t GenAiTelemetry::AllocateSessionId() {
  return next_session_id_.fetch_add(1);
}

#if defined(ORTGENAI_ENABLE_TELEMETRY)
void GenAiTelemetry::TrySuppressNetworkContext() {
  if (network_context_suppressed_.load()) return;

  if (TelemetryInternal::TrySuppressContext([&]() {
        TelemetryInternal::SuppressNetworkContext(*impl_->logger->GetSemanticContext());
      })) {
    network_context_suppressed_.store(true);
  }
}

std::shared_lock<std::shared_mutex> GenAiTelemetry::LockForLogging(bool require_enabled) {
  std::shared_lock<std::shared_mutex> lock(mutex_);
  if ((require_enabled && !enabled_.load()) || !initialized_.load() || !impl_ || !impl_->logger)
    return {};  // empty (unlocked) -> caller bails
  return lock;  // locked; ownership moves to caller
}

#endif

void GenAiTelemetry::LogProcessInfo() {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  bool emitted = false;
  bool warn_device_id_fallback = false;
  // Runtime API suppression leaves the uploader live, so ProcessInfo still fires. Full process
  // suppression never initializes a logger and returns here.
  RunLocked([&] {
    // Only log once per process, and only mark it logged while telemetry is live.
    bool expected = false;
    if (!process_info_logged_.compare_exchange_strong(expected, true)) {
      emitted = true;
      return;
    }

    const auto& device = GetDeviceInfo();
    warn_device_id_fallback = device.device_id_status == "Failed";

    MAT::EventProperties event("OnnxRuntimeGenAI.ProcessInfo");
    event.SetPopsample(100.0);
    // sessionId 0 = process scope (model sessions are numbered from 1); ProcessInfo
    // correlates with model/generate events via the AppSessionGuid logger context.
    event.SetProperty("sessionId", static_cast<int64_t>(0));
    event.SetProperty("libraryVersion", ORTGENAI_VERSION);
    // Device id is sent as ext.device.localId by the SDK (desktop override in Initialize, platform id
    // on Android/iOS), so it is not duplicated as a custom property here.
    event.SetProperty("osArchitecture", device.os_architecture);
    event.SetProperty("processorCount", static_cast<int64_t>(device.processor_count));
    event.SetProperty("totalMemoryMB", static_cast<int64_t>(device.total_memory_mb));
    event.SetProperty("cpuModel", device.cpu_model);
    event.SetProperty("deviceIdStatus", device.device_id_status);

    impl_->logger->LogEvent(event);
    // ProcessInfo captures PAL network context. Clearing it afterward is best effort.
    TrySuppressNetworkContext();
    emitted = true;
  },
            /*require_enabled=*/false);
  if (warn_device_id_fallback && g_log.enabled && g_log.warning) {
    try {
      Log("warning", "Failed to persist telemetry device ID; using an in-memory identifier");
    } catch (...) {
      // Diagnostics must not affect host control flow.
    }
  }
  // Allow a retry if telemetry was not live (or threw before emitting).
  if (!emitted) process_info_logged_.store(false);
#endif
}

bool GenAiTelemetry::LogModelLoadStart(uint32_t session_id) {
  bool emitted = false;
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.ModelLoadStart");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));

    impl_->logger->LogEvent(event);
    emitted = true;
  });
#endif
  return emitted;
}

void GenAiTelemetry::LogModelLoad(uint32_t session_id, const ModelLoadInfo& info) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.ModelLoad");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    event.SetProperty("modelType", info.model_type);
    event.SetProperty("modelFamily", info.model_family);
    event.SetProperty("modality", info.modality);
    event.SetProperty("executionProviders", info.execution_providers);
    event.SetProperty("selectedDevice", info.selected_device);
    event.SetProperty("vocabSize", static_cast<int64_t>(info.vocab_size));
    event.SetProperty("contextLength", static_cast<int64_t>(info.context_length));
    event.SetProperty("numHiddenLayers", static_cast<int64_t>(info.num_hidden_layers));
    event.SetProperty("hiddenSize", static_cast<int64_t>(info.hidden_size));
    event.SetProperty("numAttentionHeads", static_cast<int64_t>(info.num_attention_heads));
    event.SetProperty("numKeyValueHeads", static_cast<int64_t>(info.num_key_value_heads));
    event.SetProperty("intraOpNumThreads", static_cast<int64_t>(info.intra_op_num_threads));
    event.SetProperty("graphOptimizationLevel", static_cast<int64_t>(info.graph_optimization_level));
    event.SetProperty("gpuMemoryMB", static_cast<int64_t>(info.gpu_memory_mb));
    event.SetProperty("isInMemory", info.is_in_memory);

    impl_->logger->LogEvent(event);
  });
#endif
}

void GenAiTelemetry::LogModelLoadEnd(uint32_t session_id, bool is_success,
                                     double load_time_ms,
                                     const std::string& error_message) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.ModelLoadEnd");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    event.SetProperty("isSuccess", is_success);
    event.SetProperty("loadTimeMs", load_time_ms);
    if (!error_message.empty()) {
      event.SetProperty("errorMessage", ScrubStringForTelemetry(error_message));
    }

    impl_->logger->LogEvent(event);
  });
#endif
}

void GenAiTelemetry::LogGeneratorCreate(uint32_t session_id, uint32_t generator_id,
                                        int batch_size, int num_beams, int max_length,
                                        int top_k, float top_p, float temperature,
                                        bool do_sample, bool use_graph_capture, bool has_guidance) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.GeneratorCreate");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    event.SetProperty("generatorId", static_cast<int64_t>(generator_id));
    event.SetProperty("batchSize", static_cast<int64_t>(batch_size));
    event.SetProperty("numBeams", static_cast<int64_t>(num_beams));
    event.SetProperty("maxLength", static_cast<int64_t>(max_length));
    event.SetProperty("topK", static_cast<int64_t>(top_k));
    event.SetProperty("topP", static_cast<double>(top_p));
    event.SetProperty("temperature", static_cast<double>(temperature));
    event.SetProperty("doSample", do_sample);
    event.SetProperty("useGraphCapture", use_graph_capture);
    event.SetProperty("hasGuidance", has_guidance);

    impl_->logger->LogEvent(event);
  });
#endif
}

void GenAiTelemetry::LogGeneration(uint32_t session_id, uint32_t generator_id,
                                   int64_t prompt_tokens, const std::string& input_modality,
                                   const GenerateEndInfo& info,
                                   int64_t start_timestamp_ms, int64_t end_timestamp_ms) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties start_event("OnnxRuntimeGenAI.GenerateStart");
    if (!PrepareSampledEvent(start_event, app_session_guid_, session_id)) return;
    start_event.SetTimestamp(start_timestamp_ms);
    start_event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    start_event.SetProperty("generatorId", static_cast<int64_t>(generator_id));
    start_event.SetProperty("promptTokens", prompt_tokens);
    start_event.SetProperty("inputModality", input_modality);
    impl_->logger->LogEvent(start_event);

    MAT::EventProperties end_event("OnnxRuntimeGenAI.GenerateEnd");
    if (!PrepareSampledEvent(end_event, app_session_guid_, session_id)) return;
    end_event.SetTimestamp(end_timestamp_ms);
    end_event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    end_event.SetProperty("generatorId", static_cast<int64_t>(generator_id));
    end_event.SetProperty("totalTokens", info.total_tokens);
    end_event.SetProperty("generatedTokens", info.generated_tokens);
    end_event.SetProperty("rewindCount", info.rewind_count);
    end_event.SetProperty("rewoundTokens", info.rewound_tokens);
    end_event.SetProperty("timeToFirstTokenMs", info.time_to_first_token_ms);
    end_event.SetProperty("totalTimeMs", info.total_time_ms);
    end_event.SetProperty("tokensPerSecond", info.tokens_per_second);
    impl_->logger->LogEvent(end_event);
  });
#endif
}

void GenAiTelemetry::LogAdapterActivated(uint32_t session_id, uint32_t generator_id) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.AdapterActivated");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    event.SetProperty("generatorId", static_cast<int64_t>(generator_id));

    impl_->logger->LogEvent(event);
  });
#endif
}

void GenAiTelemetry::LogRuntimeError(uint32_t session_id,
                                     const std::string& error_type,
                                     const std::string& error_message,
                                     const std::string& context) {
#if defined(ORTGENAI_ENABLE_TELEMETRY)
  RunLocked([&] {
    MAT::EventProperties event("OnnxRuntimeGenAI.RuntimeError");
    if (!PrepareSampledEvent(event, app_session_guid_, session_id)) return;
    event.SetProperty("sessionId", static_cast<int64_t>(session_id));
    event.SetProperty("errorType", error_type);
    event.SetProperty("errorMessage", ScrubStringForTelemetry(error_message));
    event.SetProperty("context", context);

    impl_->logger->LogEvent(event);
  });
#endif
}

}  // namespace Generators
