streaming_telemetry/src/grpc/telemetry_server.rs

use crate::app_state::AppState;
use futures::StreamExt;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_stream::Stream;
use tonic::{Request, Response, Status};

use crate::grpc::third_party_voyager::{
    machine_telemetry_server::MachineTelemetry, request_fqp, DataPoint, FqpType,
    Request as TelemetryRequest, RequestFqp, SetRequest, TypedStruct, TypedValue, Update,
};
use crate::handlers::xpath::{get_xpath_urls, process_event_sources, SubscribeParams};
use crate::sensor_database::sensor_database::sensor_database_main::handle_subscribe_inner;

use super::third_party_voyager;

use redfish_codegen::models::metric_report_definition::v1_4_3::MetricReportDefinitionType;
use redfish_codegen::models::odata_v4;

use crate::events::{EventResponse, EventSource, EventSources};
use crate::grpc::server_config::voyager_server_config::ServerConfig as ProtobufServerConfig;
use prost::Message;
use protobuf::text_format::parse_from_str;
use protobuf::Message as ProtobufMessage;
use third_party_voyager::ServerConfig as ProstServerConfig;
use third_party_voyager::{typed_value, Threshold, Thresholds};

#[derive(Debug, Default)]
pub struct BmcTelemetryService {
    pub state: Arc<AppState>,
    pub server_config: Arc<tokio::sync::RwLock<ProstServerConfig>>,
}

pub fn load_server_config(
    file_path: &str,
) -> Result<ProstServerConfig, Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string(file_path)?;

    // Use protobuf creat to parse a textproto string to protobuf crate format ServerConfig object..
    let config: ProtobufServerConfig = parse_from_str(&content)?;

    // Encode the protobuf crate format ServerConfig object into binary format
    let mut binary_data = Vec::new();
    config
        .write_to_vec(&mut binary_data)
        .expect("Failed to encode ProtobufServerConfig object to binary");

    // Use this trick to convert protobuf crate format message to prost crate format message
    let config = ProstServerConfig::decode(&*binary_data)
        .expect("Failed to decode binary data into ProstServerConfig");
    println!("Loaded ServerConfig: {:#?}", config);

    Ok(config)
}

fn get_double_value(typed_value: &TypedValue) -> Option<f64> {
    typed_value.value.as_ref().and_then(|value| {
        if let typed_value::Value::DoubleVal(v) = value {
            Some(*v)
        } else {
            None
        }
    })
}

pub fn get_threshold_config(reading: f64, config: &Thresholds) -> Option<&Threshold> {
    let thresholds = &config.threshold;

    if thresholds.is_empty() {
        return None;
    }

    let mut current_index = 0;

    while current_index < thresholds.len() {
        let threshold = &thresholds[current_index];
        let cross_above = threshold
            .cross_above_value
            .as_ref()
            .and_then(get_double_value);
        let cross_below = threshold
            .cross_below_value
            .as_ref()
            .and_then(get_double_value);

        match (cross_above, cross_below) {
            (Some(above), Some(below)) => {
                if reading > above && current_index < thresholds.len() - 1 {
                    current_index += 1;
                } else if reading < below && current_index > 0 {
                    current_index -= 1;
                } else {
                    return Some(threshold);
                }
            }
            (Some(above), None) => {
                if reading > above && current_index < thresholds.len() - 1 {
                    current_index += 1;
                } else {
                    return Some(threshold);
                }
            }
            (None, Some(below)) => {
                if reading < below && current_index > 0 {
                    current_index -= 1;
                } else {
                    return Some(threshold);
                }
            }
            (None, None) => return Some(threshold),
        }
    }

    // If we've gone through all thresholds, return the last one
    thresholds.last()
}

fn create_subscribe_params(req_fqp: &RequestFqp) -> SubscribeParams {
    SubscribeParams {
        sample_mode: Some(
            match req_fqp.mode {
                1 => "SAMPLING_MODE_ON_CHANGE",
                2 => "SAMPLING_MODE_PERIODIC",
                3 => "SAMPLING_MODE_THRESHOLD",
                _ => "Unspecified",
            }
            .to_string(),
        ),
        sampling_frequency_ns: Some(req_fqp.sample_frequency_expect_ns),
        export_frequency_ns: Some(req_fqp.export_frequency_ns),
        suppress_redundant: Some(req_fqp.suppress_redundant),
        ..Default::default()
    }
}

async fn create_response_stream(
    state: Arc<AppState>,
    req_id: String,
    event_sources: EventSources,
    server_config: Arc<tokio::sync::RwLock<ProstServerConfig>>,
    tx: mpsc::Sender<Result<Update, Status>>,
) -> Result<(), Status> {
    let mut rx_stream =
        handle_subscribe_inner(state.clone(), event_sources, Some(server_config)).await;

    while let Some(batch_events) = rx_stream.next().await {
        let mut current_update: Option<Update> = None;
        for event in batch_events {
            if let Some(_sensor_value) = &event.sensor_value {
                let data_point = create_sensor_data_point(&event);

                if let Some(update) = &mut current_update {
                    update.data_points.push(data_point);
                } else {
                    current_update = Some(Update {
                        req_id: req_id.clone(),
                        data_points: vec![data_point],
                        ..Default::default()
                    });
                }
            }
        }

        if let Some(update) = current_update.take() {
            if let Err(e) = tx.send(Ok(update)).await {
                eprintln!("Error sending update: {:?}", e);
                return Err(Status::internal("Failed to send update"));
            }
        }
    }
    Ok(())
}

async fn handle_fqp(
    state: Arc<AppState>,
    server_config: Arc<tokio::sync::RwLock<ProstServerConfig>>,
    req_fqp: &RequestFqp,
    req_id: String,
    tx: mpsc::Sender<Result<Update, Status>>,
) -> Result<(), Status> {
    let fqp = match &req_fqp.fqp {
        Some(fqp) => fqp,
        None => return Err(Status::invalid_argument("Address not specified")),
    };
    let identifiers: HashMap<String, Vec<String>> = fqp
        .identifiers
        .iter()
        .map(|(k, v)| (k.clone(), vec![v.clone()]))
        .collect();
    let invariable_filters = HashMap::new();
    // TODO: need adapt to proto encoded filters
    // let mut variable_filters = HashMap::new();
    // for (k, v) in fqp.filters.iter() {
    //     if is_invariable_filter(v) {
    //         invariable_filters.insert(k.clone(), v.clone());
    //     } else {
    //         variable_filters.insert(k.clone(), v.clone());
    //     }
    // }

    let segments: Vec<&str> = fqp.specifier.split('/').collect();
    let urls =
        match get_xpath_urls(&state, &identifiers, &Some(invariable_filters), &segments).await {
            Ok(urls) => urls,
            Err(_) => return Err(Status::internal("Failed to get XPath URLs")),
        };

    let params = create_subscribe_params(req_fqp);
    let event_sources = process_event_sources(&params, urls);
    create_response_stream(state, req_id, event_sources, server_config, tx).await
}

async fn get_selected_fqps(
    server_config: &Arc<tokio::sync::RwLock<ProstServerConfig>>,
    req_config_group: &str,
) -> Vec<(String, Option<Threshold>)> {
    let server_config = server_config.read().await;
    let mut result = Vec::new();

    // Step 1: Get the ConfigGroup from the top-level map
    if let Some(config_group) = server_config.cfg_groups.get(req_config_group) {
        // Step 2: Iterate through req_fqp_names in the ConfigGroup
        for req_fqp_name in &config_group.req_fqp_names {
            // Step 3: Get the ReqFqpConfig from the second-level map
            if let Some(req_fqp_config) = server_config.req_fqp_configs.get(req_fqp_name) {
                // Step 4: Iterate through RequestFqp in ReqFqpConfig
                for req_fqp in &req_fqp_config.req_fqp {
                    if let Some(fqp) = &req_fqp.fqp {
                        let specifier = fqp.specifier.clone();

                        // Step 5: Get the matching Threshold from the threshold_config map
                        let threshold =
                            if let Some(request_fqp::Config::ThresholdConfig(threshold_config)) =
                                &req_fqp.config
                            {
                                server_config
                                    .threshold_config
                                    .get(&req_fqp.req_fqp_name)
                                    .and_then(|thresholds| {
                                        thresholds
                                            .threshold
                                            .iter()
                                            .find(|t| t.name == *threshold_config)
                                            .cloned()
                                    })
                            } else {
                                None
                            };

                        result.push((specifier, threshold));
                    }
                }
            }
        }
    }

    result
}

fn get_event_sources(selected: &Vec<(String, Option<Threshold>)>) -> EventSources {
    let mut event_sources = EventSources::default();
    for (url, threshold) in selected {
        // TODO: set default reporting interval
        let mut new_event_source = EventSource {
            odata_id: Some(odata_v4::Id(url.to_owned())),
            sampling_type: Some(MetricReportDefinitionType::Periodic),
            sampling_rate: Some(1),
            ..Default::default()
        };

        if let Some(threshold) = threshold {
            new_event_source.sampling_rate =
                Some((1_000_000_000 / threshold.sample_frequency_expect_ns) as i64);
        }
        event_sources.0.push(new_event_source);
    }
    event_sources
}

async fn handle_config_group(
    state: Arc<AppState>,
    server_config: Arc<tokio::sync::RwLock<ProstServerConfig>>,
    req: TelemetryRequest,
    tx: mpsc::Sender<Result<Update, Status>>,
) -> Result<(), Status> {
    let selected = get_selected_fqps(&server_config, &req.req_config_group).await;
    println!(
        "handle_config_group {} selected Fqp and Theshold: {:#?}",
        &req.req_config_group, selected
    );
    let event_sources = get_event_sources(&selected);
    create_response_stream(state, req.req_id.clone(), event_sources, server_config, tx).await
}

fn create_sensor_data_point(event: &EventResponse) -> DataPoint {
    DataPoint {
        timestamp_ns: event.timestamp.unwrap_or_default() as u64,
        data: Some(third_party_voyager::data_point::Data::KeyValue(
            TypedStruct {
                fields: HashMap::from([
                    (
                        "@odata.id".to_string(),
                        TypedValue {
                            value: Some(third_party_voyager::typed_value::Value::StringVal(
                                event.sensor_value.as_ref().unwrap().odata_id.0.clone(),
                            )),
                        },
                    ),
                    (
                        "SensorValue".to_string(),
                        TypedValue {
                            value: Some(third_party_voyager::typed_value::Value::DoubleVal(
                                event
                                    .sensor_value
                                    .as_ref()
                                    .unwrap()
                                    .reading
                                    .expect("reading value is required"),
                            )),
                        },
                    ),
                    (
                        "Status.Health".to_string(),
                        TypedValue {
                            value: Some(third_party_voyager::typed_value::Value::StringVal(
                                "OK".to_owned(),
                            )),
                        },
                    ),
                ]),
            },
        )),
        ..Default::default()
    }
}

async fn handle_sensor_request(
    state: Arc<AppState>,
    req_fqp: &RequestFqp,
    req_id: String,
    tx: mpsc::Sender<Result<Update, Status>>,
) -> Result<(), Status> {
    // Use fixed identifiers
    let identifiers: HashMap<String, Vec<String>> = HashMap::from([
        ("ChassisId".to_string(), vec!["*".to_string()]),
        ("SensorId".to_string(), vec!["*".to_string()]),
    ]);
    let segments: Vec<&str> = "/redfish/v1/Chassis/{ChassisId}/Sensors/{SensorId}"
        .split('/')
        .collect();

    let urls = match get_xpath_urls(&state, &identifiers, &None, &segments).await {
        Ok(urls) => urls,
        Err(_) => return Err(Status::internal("Failed to get XPath URLs")),
    };

    let params = create_subscribe_params(req_fqp);
    let event_sources = process_event_sources(&params, urls);
    let mut rx_stream = handle_subscribe_inner(state.clone(), event_sources, None).await;

    let mut current_update: Option<Update> = None;

    while let Some(batch_events) = rx_stream.next().await {
        for event in batch_events {
            if let Some(_sensor_value) = &event.sensor_value {
                let data_point = create_sensor_data_point(&event);

                if let Some(update) = &mut current_update {
                    update.data_points.push(data_point);
                } else {
                    current_update = Some(Update {
                        req_id: req_id.clone(),
                        data_points: vec![data_point],
                        ..Default::default()
                    });
                }
            }
        }

        if let Some(update) = current_update.take() {
            if let Err(e) = tx.send(Ok(update)).await {
                eprintln!("Error sending update: {:?}", e);
                return Err(Status::internal("Failed to send update"));
            }
        }
    }
    Ok(())
}

#[tonic::async_trait]
impl MachineTelemetry for BmcTelemetryService {
    type SubscribeV2Stream = ReceiverStream<Result<Update, Status>>;
    type SubscribeStream = Pin<Box<dyn Stream<Item = Result<Update, Status>> + Send + 'static>>;

    async fn subscribe_v2(
        &self,
        request: Request<tonic::Streaming<TelemetryRequest>>,
    ) -> Result<Response<Self::SubscribeV2Stream>, Status> {
        let mut stream = request.into_inner();
        let state = self.state.clone();
        let server_config = self.server_config.clone();

        let (tx, rx) = mpsc::channel(16); // Adjust channel size as needed

        tokio::spawn(async move {
            while let Some(req) = match stream.message().await {
                Ok(req) => req,
                Err(e) => {
                    eprintln!("Error receiving stream message: {:?}", e);
                    return;
                }
            } {
                if !req.req_config_group.is_empty() {
                    if let Err(e) =
                        handle_config_group(state.clone(), server_config.clone(), req, tx.clone())
                            .await
                    {
                        eprintln!("Handler error: {:?}", e);
                    }
                    continue;
                }
                for req_fqp in &req.req_fqp {
                    if let Some(fqp) = &req_fqp.fqp {
                        println!("fqp: {:?}", fqp);
                        match FqpType::try_from(fqp.r#type) {
                            Ok(fqp_type) => {
                                match fqp_type {
                                    FqpType::NotSet => {
                                        if let Err(e) = handle_fqp(
                                            state.clone(),
                                            server_config.clone(),
                                            req_fqp,
                                            req.req_id.clone(),
                                            tx.clone(),
                                        )
                                        .await
                                        {
                                            eprintln!("Handler error: {:?}", e);
                                        }
                                    }
                                    FqpType::RedfishResource => {
                                        // Handle RedfishResource case, use Sensor as example
                                        if fqp.specifier.contains("Sensor") {
                                            if let Err(e) = handle_sensor_request(
                                                state.clone(),
                                                req_fqp,
                                                req.req_id.clone(),
                                                tx.clone(),
                                            )
                                            .await
                                            {
                                                eprintln!("Handler error: {:?}", e);
                                            }
                                        } else {
                                            eprintln!("Subscribe by RedfishResource for {} not implemented", fqp.specifier);
                                        }
                                    }
                                }
                            }
                            Err(e) => {
                                eprintln!("Invalid FqpType value: {}. Error: {:?}", fqp.r#type, e);
                                continue;
                            }
                        }
                    }
                }
            }
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn get(&self, _: Request<TelemetryRequest>) -> Result<Response<Update>, Status> {
        unimplemented!()
    }

    async fn put(&self, _: Request<SetRequest>) -> Result<Response<Update>, Status> {
        unimplemented!()
    }

    async fn post(&self, _: Request<SetRequest>) -> Result<Response<Update>, Status> {
        unimplemented!()
    }

    async fn patch(&self, _: Request<SetRequest>) -> Result<Response<Update>, Status> {
        unimplemented!()
    }

    async fn delete(&self, _: Request<SetRequest>) -> Result<Response<Update>, Status> {
        unimplemented!()
    }

    async fn subscribe(
        &self,
        _: Request<TelemetryRequest>,
    ) -> Result<Response<Self::SubscribeStream>, Status> {
        unimplemented!()
    }
}