// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! `dial9-tokio-telemetry` integration, compiled when the `dial9` feature is on.
//!
//! Captures Tokio runtime-level events (poll start/end, worker park/unpark,
//! task spawn/terminate) into rotating binary trace segments. This is an
//! on-demand profiler rather than always-on telemetry: it is disabled by
//! default and requires a `--cfg tokio_unstable` build.

use super::config::Dial9Config;
use super::state::{dial9_runtime_state, measure_disk_usage_bytes};
use super::{EVENT_DIAL9_STATE, LOG_COMPONENT_OBS, LOG_SUBSYSTEM_DIAL9};
use crate::TelemetryError;
use dial9_tokio_telemetry::telemetry::{
    Dial9Handle, Dial9HandleTokioExt, DiskBuffer, ProcessResourceUsageConfig, RecorderPerfExt, TokioAttachOptions, recorder,
};
use std::time::Duration;
use tracing::{info, warn};

pub type TelemetryGuard = Dial9Handle;

type ShutdownRecorder = Box<dyn FnOnce() + Send + 'static>;

/// How often the background refresher restates trace-file disk usage.
const DISK_USAGE_REFRESH_INTERVAL: Duration = Duration::from_secs(60);

/// Maximum time spent flushing the recorder during graceful shutdown.
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);

/// Name recorded in segment metadata so the trace viewer can label workers.
const RUNTIME_NAME: &str = "rustfs-worker";

/// Owns a live dial9 telemetry session.
///
/// Dropping this guard flushes buffered events and seals the active segment.
/// It must be dropped before the process exits: a guard parked in a `static` is
/// never dropped, and the final buffered events — usually the interesting ones —
/// are lost.
pub struct Dial9SessionGuard {
    guard: TelemetryGuard,
    shutdown: Option<ShutdownRecorder>,
    config: Dial9Config,
}

impl Dial9SessionGuard {
    /// Whether the underlying telemetry session is recording.
    pub fn is_active(&self) -> bool {
        self.guard.is_enabled() && self.guard.is_connected() && !self.guard.is_stopped()
    }
}

impl std::fmt::Debug for Dial9SessionGuard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Dial9SessionGuard")
            .field("active", &self.is_active())
            .field("output_dir", &self.config.output_dir)
            .finish()
    }
}

impl Drop for Dial9SessionGuard {
    fn drop(&mut self) {
        dial9_runtime_state().record_runtime_stopped();
        info!(
            event = EVENT_DIAL9_STATE,
            component = LOG_COMPONENT_OBS,
            subsystem = LOG_SUBSYSTEM_DIAL9,
            state = "flushed",
            "dial9 state changed"
        );

        if let Some(shutdown) = self.shutdown.take() {
            shutdown();
        }
    }
}

/// Build a Tokio runtime with dial9 telemetry attached and recording started.
///
/// # Errors
///
/// Returns an error when dial9 is not enabled, when the output directory cannot
/// be created, or when the traced runtime fails to build. Callers are expected
/// to fall back to a standard runtime.
pub fn build_traced_runtime(
    builder: tokio::runtime::Builder,
) -> Result<(tokio::runtime::Runtime, Dial9SessionGuard), TelemetryError> {
    let config = Dial9Config::from_env();
    if !config.enabled {
        return Err(TelemetryError::Io("dial9 telemetry is not enabled".to_string()));
    }

    std::fs::create_dir_all(&config.output_dir).map_err(|e| {
        dial9_runtime_state().record_runtime_error(&config);
        TelemetryError::Io(format!("Failed to create dial9 output directory '{}': {e}", config.output_dir))
    })?;

    let writer = DiskBuffer::builder()
        .base_path(config.base_path())
        .max_file_size(config.max_file_size)
        .max_total_size(config.total_disk_budget())
        .build()
        .map_err(|e| {
            dial9_runtime_state().record_runtime_error(&config);
            TelemetryError::Io(format!("Failed to create dial9 DiskBuffer: {e}"))
        })?;

    let recorder = recorder(writer)
        .with_process_resource_usage(ProcessResourceUsageConfig::default())
        .build();
    let guard = recorder.handle().clone();
    let shutdown: ShutdownRecorder = Box::new(move || recorder.graceful_shutdown(SHUTDOWN_TIMEOUT));

    let attached = guard
        .attach_tokio_runtime(
            builder,
            TokioAttachOptions::builder()
                .runtime_name(RUNTIME_NAME)
                .task_tracking_enabled(true)
                .build(),
        )
        .map(|runtime| (runtime, guard, shutdown));

    // No task dumps here. dial9 captures a task dump only for futures it
    // wrapped itself, i.e. those spawned via `dial9_tokio_telemetry::spawn`;
    // `tokio::spawn` gets no wrapper. RustFS spawns with `tokio::spawn`
    // throughout, so enabling task dumps records nothing. Measured on an
    // identical workload: 0 dumps via `tokio::spawn`, 14709 via `dial9::spawn`.
    // See rustfs/backlog#1157 (D9-16) and dial9-rs/dial9#477.
    //
    // No `with_s3_uploader` here: dial9's `worker-s3` feature carries a
    // vulnerable TLS stack. See the note in `crates/obs/Cargo.toml`.
    finish_traced_runtime(attached, config)
}

/// Publish the outcome of a traced-runtime build and start the background
/// disk-usage refresher.
fn finish_traced_runtime(
    started: std::io::Result<(tokio::runtime::Runtime, TelemetryGuard, ShutdownRecorder)>,
    config: Dial9Config,
) -> Result<(tokio::runtime::Runtime, Dial9SessionGuard), TelemetryError> {
    let (runtime, guard, shutdown) = started.map_err(|e| {
        dial9_runtime_state().record_runtime_error(&config);
        TelemetryError::Io(format!("Failed to attach dial9 runtime telemetry: {e}"))
    })?;

    if !guard.is_enabled() {
        dial9_runtime_state().record_runtime_error(&config);
        return Err(TelemetryError::Io("dial9 runtime telemetry attached with recording disabled".to_string()));
    }

    dial9_runtime_state().record_runtime_started(&config);
    runtime.spawn(refresh_disk_usage_loop());

    info!(
        event = EVENT_DIAL9_STATE,
        component = LOG_COMPONENT_OBS,
        subsystem = LOG_SUBSYSTEM_DIAL9,
        state = "recording",
        output_dir = %config.output_dir,
        file_prefix = %config.file_prefix,
        disk_budget_bytes = config.total_disk_budget(),
        "dial9 state changed"
    );

    Ok((
        runtime,
        Dial9SessionGuard {
            guard,
            shutdown: Some(shutdown),
            config,
        },
    ))
}

/// Periodically restate trace-file disk usage so the metrics collector can read
/// it from an atomic instead of walking the directory on its own thread.
async fn refresh_disk_usage_loop() {
    let mut ticker = tokio::time::interval(DISK_USAGE_REFRESH_INTERVAL);
    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    loop {
        ticker.tick().await;

        let Some(location) = dial9_runtime_state().trace_location() else {
            continue;
        };

        // `measure_disk_usage_bytes` walks a directory and stats every entry,
        // so it must stay off the async workers.
        let measured =
            tokio::task::spawn_blocking(move || measure_disk_usage_bytes(&location.output_dir, &location.file_prefix)).await;

        match measured {
            Ok(bytes) => dial9_runtime_state().set_disk_usage_bytes(bytes),
            Err(e) => warn!(
                event = EVENT_DIAL9_STATE,
                component = LOG_COMPONENT_OBS,
                subsystem = LOG_SUBSYSTEM_DIAL9,
                result = "disk_usage_refresh_failed",
                error = %e,
                "dial9 state changed"
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::state::{reset_for_test, runtime_stats_snapshot};
    use super::*;

    #[test]
    fn build_traced_runtime_refuses_when_disabled() {
        reset_for_test();
        // dial9 is not configured in the unit-test environment, so
        // `Dial9Config::from_env()` yields a disabled config.
        let result = build_traced_runtime(tokio::runtime::Builder::new_current_thread());
        assert!(result.is_err(), "disabled dial9 must not build a traced runtime");
    }

    #[test]
    fn recording_session_reports_active_until_stopped() {
        reset_for_test();
        let config = Dial9Config {
            enabled: true,
            ..Dial9Config::default()
        };

        dial9_runtime_state().record_runtime_started(&config);
        assert_eq!(runtime_stats_snapshot().active_sessions, 1);

        // Mirrors what `Dial9SessionGuard::drop` publishes.
        dial9_runtime_state().record_runtime_stopped();
        assert_eq!(runtime_stats_snapshot().active_sessions, 0);
    }
}
