/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * documentdb_gateway_core/src/requests/request_tracker.rs
 *
 *-------------------------------------------------------------------------
 */

use std::sync::atomic::{AtomicU64, Ordering};

use tokio::time::Instant;

use crate::telemetry::utils;

#[derive(Debug)]
pub enum RequestIntervalKind {
    /// Time spent reading stream from request body.
    ReadRequest,

    /// Interval kind for the overall request processing duration, which includes `FormatRequest`, and `HandleRequest` via backend.
    /// `ReadRequest` and `WriteResponse` are not part of `HandleMessage`.
    HandleMessage,

    /// Time spent formatting and parsing the incoming request.
    FormatRequest,

    /// Time spent handling the request, which includes `ProcessRequest` and, if applicable,
    /// `PostgresBeginTransaction`, `PostgresSetStatementTimeout`, and `PostgresCommitTransaction`.
    HandleRequest,

    /// Time spent in network transport and Postgres processing.
    ProcessRequest,

    /// Time spent beginning a Postgres transaction.
    PostgresBeginTransaction,

    /// Time spent setting statement timeout parameters in Postgres.
    PostgresSetStatementTimeout,

    /// Time spent committing a Postgres transaction.
    PostgresCommitTransaction,

    /// Time spent acquiring a connection from the Postgres connection pool.
    OpenBackendConnection,

    /// Time spent writing the response to the stream.
    WriteResponse,

    /// Special value used to define the size of the metrics array.
    MaxUnused,
}

#[derive(Debug)]
pub struct RequestTracker {
    pub request_interval_metrics_array: [AtomicU64; RequestIntervalKind::MaxUnused as usize],
}

impl Default for RequestTracker {
    fn default() -> Self {
        Self::new()
    }
}

impl RequestTracker {
    #[must_use]
    pub fn new() -> Self {
        Self {
            request_interval_metrics_array: std::array::from_fn(|_| AtomicU64::new(0)),
        }
    }

    #[expect(clippy::cast_possible_truncation, reason = "nanoseconds fit in u64")]
    pub fn record_duration(&self, interval: RequestIntervalKind, start_time: Instant) {
        let elapsed = start_time.elapsed();
        self.request_interval_metrics_array[interval as usize]
            .fetch_add(elapsed.as_nanos() as u64, Ordering::Relaxed);
    }

    pub fn get_interval_elapsed_time(&self, interval: RequestIntervalKind) -> u64 {
        self.request_interval_metrics_array[interval as usize].load(Ordering::Relaxed)
    }

    pub fn get_interval_elapsed_time_ms(&self, interval: RequestIntervalKind) -> u64 {
        utils::ns_to_ms(self.get_interval_elapsed_time(interval))
    }
}
