bcms 0.1.12

use bcms::api::BackendAPI;
use bcms::ble::BleManager;
use bcms::config::Config;
use bcms::data::DataStore;
use bcms::devices::DeviceMemory;
use bcms::rpc::RpcServer;
use bcms::submission::SubmissionTracker;
use bcms::Result;
use ones_oidc::{OpenIdconnectClient, read_private_key, load_device_config};
use openidconnect::{ClientId, IssuerUrl};
use openidconnect::core::CoreProviderMetadata;
use openidconnect::reqwest::async_http_client;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::{interval, sleep, Duration};
use tracing::{error, info, warn};
use clap::Parser;

#[derive(Parser, Debug)]
#[command(name = "bcms-daemon")]
#[command(about = "Bluetooth Client Manager Service Daemon")]
struct Args {
    #[arg(long, default_value = "false")]
    debug: bool,
    
    #[arg(long, help = "Client identifier for API access (overrides device config)")]
    identifier: Option<String>,
    
    #[arg(long)]
    config_file: Option<PathBuf>,
    
    #[arg(long, help = "Enable device identity with OIDC authentication")]
    use_device_identity: bool,
    
    #[arg(long, default_value = "/root/.local/share/px-device-identity/private.pem", help = "Path to device private key file")]
    device_private_key: Option<PathBuf>,
    
    #[arg(long, help = "OIDC issuer URL (overrides device config)")]
    issuer_url: Option<String>,
    
    #[arg(long, default_value = "/etc/px-device-identity/device.yml", help = "Path to device identity config file")]
    device_config_file: Option<PathBuf>,
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();
    
    // Initialize logging
    init_logging(args.debug);
    
    info!("Starting BCMS Daemon...");
    
    // Load configuration
    let config = if let Some(ref config_path) = args.config_file {
        let config_str = std::fs::read_to_string(config_path)
            .map_err(|e| bcms::Error::Config(format!("Failed to read config file: {}", e)))?;
        serde_json::from_str(&config_str)
            .map_err(|e| bcms::Error::Config(format!("Failed to parse config file: {}", e)))?
    } else {
        Config::default()
    };
    
    // Initialize device memory with persistent storage
    let device_file = get_device_file_path();
    let device_memory = Arc::new(DeviceMemory::new(Some(device_file)));
    
    // Load existing devices
    if let Err(e) = device_memory.load_devices().await {
        warn!("Failed to load existing devices: {}", e);
    }
    
    // Initialize data store
    let data_store = Arc::new(DataStore::new());
    
    // Initialize BLE manager
    let ble_manager = Arc::new(
        BleManager::new(config.clone(), Arc::clone(&device_memory), Arc::clone(&data_store))
            .await
            .expect("Failed to initialize BLE manager")
    );
    
    // Get identifier from CLI args (used for API calls)
    let identifier = args.identifier.clone();

    // Initialize OIDC client if device identity is enabled
    let oidc_client = if args.use_device_identity {
        match initialize_oidc_client(&args, &config).await {
            Ok(client) => {
                info!("OIDC client initialized successfully");
                Some(Arc::new(client))
            }
            Err(e) => {
                error!("Failed to initialize OIDC client: {}", e);
                warn!("Continuing in local-only mode without device identity");
                None
            }
        }
    } else {
        None
    };

    // Initialize API client if device identity is enabled, or if mock mode is configured
    let mock_mode = config.mock_backend_url.is_some() && config.mock_bearer_token.is_some();
    let api_client = if mock_mode || (args.use_device_identity && oidc_client.is_some()) {
        if mock_mode {
            info!("Running in mock mode with backend: {}", config.mock_backend_url.as_ref().unwrap());
        }
        Some(Arc::new(Mutex::new(BackendAPI::new(
            identifier.clone(),
            config.http_timeout(),
            oidc_client.clone(),
            config.mock_backend_url.clone(),
            config.mock_bearer_token.clone(),
        ))))
    } else {
        None
    };
    
    // Initialize RPC server with optional API client
    let rpc_server = if let Some(ref api_client) = api_client {
        RpcServer::with_api_client(
            config.rpc_bind_address().parse()
                .expect("Invalid RPC bind address"),
            Arc::clone(&device_memory),
            Arc::clone(&ble_manager),
            Arc::clone(&data_store),
            Arc::clone(api_client),
        )
    } else {
        RpcServer::new(
            config.rpc_bind_address().parse()
                .expect("Invalid RPC bind address"),
            Arc::clone(&device_memory),
            Arc::clone(&ble_manager),
            Arc::clone(&data_store),
        )
    };
    
    info!("All components initialized. Starting async loops...");

    // Start async loops conditionally based on API client availability (device identity or mock mode)
    if let Some(api_client) = api_client {
        info!("API client available - starting all loops including API submission");
            tokio::select! {
                result = device_discovery_loop(Arc::clone(&ble_manager), &config) => {
                    error!("Device discovery loop ended: {:?}", result);
                }
                result = data_submission_loop(Arc::clone(&data_store), Arc::clone(&device_memory), Arc::clone(&ble_manager), Arc::clone(&api_client), &config) => {
                    error!("Data submission loop ended: {:?}", result);
                }
                result = cache_cleanup_loop(Arc::clone(&data_store), &config) => {
                    error!("Cache cleanup loop ended: {:?}", result);
                }
                result = rpc_server.start() => {
                    error!("RPC server ended: {:?}", result);
                }
                result = device_registration_loop(Arc::clone(&device_memory), Arc::clone(&data_store), Arc::clone(&api_client), &config) => {
                    error!("Device registration loop ended: {:?}", result);
                }
                result = ble_manager.start_advertisement_monitoring_loop() => {
                    error!("Advertisement monitoring loop ended: {:?}", result);
                }
            }
    } else {
        info!("Device identity disabled - running local-only mode");
        tokio::select! {
            result = device_discovery_loop(Arc::clone(&ble_manager), &config) => {
                error!("Device discovery loop ended: {:?}", result);
            }
            result = cache_cleanup_loop(Arc::clone(&data_store), &config) => {
                error!("Cache cleanup loop ended: {:?}", result);
            }
            result = rpc_server.start() => {
                error!("RPC server ended: {:?}", result);
            }
            result = ble_manager.start_advertisement_monitoring_loop() => {
                error!("Advertisement monitoring loop ended: {:?}", result);
            }
        }
    }
    
    Ok(())
}

/// Initialize logging based on debug flag
fn init_logging(debug: bool) {
    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
    
    let filter = if debug {
        "debug"
    } else {
        "info"
    };
    
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(filter))
        )
        .with(tracing_subscriber::fmt::layer().with_target(false))
        .init();
}

/// Get the device file path for persistent storage
fn get_device_file_path() -> PathBuf {
    let mut path = dirs::data_local_dir()
        .unwrap_or_else(|| PathBuf::from("."));
    path.push("bluetooth-client-manager-service");
    path.push("device.json");
    path
}

/// Device discovery loop - scans for BLE devices every 5s
async fn device_discovery_loop(ble_manager: Arc<BleManager>, config: &Config) -> Result<()> {
    info!("Starting device discovery loop");
    let mut interval = interval(config.bluetooth_scan_duration());
    
    loop {
        interval.tick().await;
        
        if let Err(e) = ble_manager.scan_for_devices().await {
            error!("Device discovery error: {}", e);
        }
    }
}

/// Data submission loop - submits collected data to backend every 30s
async fn data_submission_loop(
    data_store: Arc<DataStore>,
    device_memory: Arc<DeviceMemory>,
    ble_manager: Arc<BleManager>,
    api_client: Arc<Mutex<BackendAPI>>,
    config: &Config,
) -> Result<()> {
    info!("Starting data submission loop");
    let mut interval = interval(config.data_submission_duration());
    let mut tracker = SubmissionTracker::new();
    let base_interval_secs = config.data_submission_interval;
    let max_backoff_secs = 300.0; // Cap backoff at 5 minutes

    loop {
        interval.tick().await;

        // Collect data from approved devices
        let approved_devices = device_memory.get_approved_devices().await;

        for device in approved_devices {
            if device.paired {
                // Collect data from the device
                match ble_manager.connect_and_collect_data(&device.address).await {
                    Ok(data) => {
                        if !data.is_empty() {
                            // Add to data store
                            if let Err(e) = data_store.add_data_batch(data.clone()).await {
                                error!("Failed to add data to store: {}", e);
                            } else {
                                info!("Collected {} data points from device {}", data.len(), device.address);
                            }
                        }
                    }
                    Err(e) => {
                        warn!("Failed to collect data from device {}: {}", device.address, e);
                    }
                }
            }
        }

        // Enforce memory limits before submission
        if let Err(e) = data_store.enforce_limits(
            config.max_data_points_per_device,
            config.max_total_data_points
        ).await {
            warn!("Failed to enforce data limits: {}", e);
        }

        // Get only NEW data per device (each device tracked independently)
        let new_data = data_store.get_data_since_per_device(tracker.get_all_timestamps()).await;

        if !new_data.is_empty() {
            info!("Found {} new data points across devices", new_data.len());

            // Apply rate limiting (1 sample per device+type every N seconds)
            let rate_limited_data = DataStore::limit_sample_rate(new_data.clone(), config.samples_every_seconds);
            info!("After rate limiting: {} data points (reduced from {})", rate_limited_data.len(), new_data.len());

            if !rate_limited_data.is_empty() {
                // Get registered devices for API submission
                let registered_devices = device_memory.get_registered_devices().await;

                // Build set of device addresses that will actually be submitted
                let submittable_addresses: std::collections::HashSet<String> = registered_devices
                    .iter()
                    .filter(|d| d.approved && d.backend_device_id.is_some())
                    .map(|d| d.address.clone())
                    .collect();

                // Filter data to only include devices that will actually be submitted
                let submitted_data: Vec<_> = rate_limited_data
                    .iter()
                    .filter(|d| submittable_addresses.contains(d.address()))
                    .collect();

                // Transform data to API format
                let api_payload = BackendAPI::dump_iot_data_for_api_submission(
                    rate_limited_data.clone(),
                    &registered_devices,
                );

                // Only submit if we have data from registered devices
                if !api_payload.data.is_empty() {
                    let mut api = api_client.lock().await;
                    match api.submit_iot_data_formatted(api_payload).await {
                        Ok(()) => {
                            info!("Successfully submitted {} data points", submitted_data.len());

                            // Update tracker for submitted devices only
                            tracker.mark_submitted(&submitted_data);

                            // Log pending data count (devices not yet registered)
                            let pending_count = rate_limited_data.len() - submitted_data.len();
                            if pending_count > 0 {
                                info!("{} data points pending from unregistered devices", pending_count);
                            }

                            // Reset failure tracking on success
                            let submitted_addrs: Vec<&str> = submitted_data.iter()
                                .map(|d| d.address())
                                .collect::<std::collections::HashSet<_>>()
                                .into_iter()
                                .collect();
                            tracker.mark_batch_recovered(&submitted_addrs);

                            let stats = tracker.stats();
                            if stats.devices_with_failures > 0 {
                                info!("Submission recovered for {} devices", submitted_addrs.len());
                            }
                        }
                        Err(e) => {
                            // Mark all devices in this batch as failed
                            let failed_addrs: Vec<&str> = submitted_data.iter()
                                .map(|d| d.address())
                                .collect::<std::collections::HashSet<_>>()
                                .into_iter()
                                .collect();
                            tracker.mark_batch_failed(&failed_addrs);

                            let backoff = tracker.get_global_backoff(base_interval_secs, max_backoff_secs);
                            let stats = tracker.stats();

                            error!("Failed to submit data (max failures: {}): {}", stats.max_failures, e);
                            warn!("Applying exponential backoff: waiting {:.1}s before next retry", backoff.as_secs_f64());

                            // Apply backoff delay
                            sleep(backoff).await;
                        }
                    }
                } else {
                    // No submittable devices have data - warn about pending
                    let pending_device_count = rate_limited_data
                        .iter()
                        .map(|d| d.address())
                        .collect::<std::collections::HashSet<_>>()
                        .len();
                    warn!("{} data points from {} devices pending - devices not yet registered with backend",
                          rate_limited_data.len(), pending_device_count);
                }
            } else {
                info!("No data remaining after rate limiting");
            }
        }

        // Periodic cleanup of old data to prevent memory growth
        if let Err(e) = data_store.cleanup_old_data(config.max_data_age_seconds).await {
            warn!("Failed to cleanup old data: {}", e);
        }
    }
}

/// Cache cleanup loop - cleans old data every 180s
async fn cache_cleanup_loop(data_store: Arc<DataStore>, config: &Config) -> Result<()> {
    info!("Starting cache cleanup loop");
    let mut interval = interval(config.cache_clear_duration());
    
    loop {
        interval.tick().await;
        
        // Clean up data older than configured max age
        let max_age_seconds = config.max_data_age_seconds;
        match data_store.cleanup_old_data(max_age_seconds).await {
            Ok(removed_count) => {
                if removed_count > 0 {
                    info!("Cleaned up {} old data points", removed_count);
                }
            }
            Err(e) => {
                error!("Failed to cleanup old data: {}", e);
            }
        }
    }
}

/// Device registration loop - registers approved devices with backend API
async fn device_registration_loop(
    device_memory: Arc<DeviceMemory>,
    data_store: Arc<DataStore>,
    api_client: Arc<Mutex<BackendAPI>>,
    config: &Config,
) -> Result<()> {
    info!("Starting device registration loop");
    let mut interval = interval(Duration::from_secs(30)); // Check every 30 seconds

    loop {
        interval.tick().await;

        let approved_devices = device_memory.get_approved_devices().await;

        for device in approved_devices {
            // Infer supportedDataTypes from data received for this device
            let data_types = data_store.get_device_data_types(&device.address).await;
            let supported_data_types = if data_types.is_empty() {
                None
            } else {
                Some(data_types)
            };

            // Check if device exists in backend
            let mut api = api_client.lock().await;
            match api.create_iot_device_if_not_exists(
                &device.address,
                Some(&device.name),
                device.connection_type.as_deref(),
                None,
                supported_data_types,
            ).await {
                Ok(response) => {
                    // Check if backend wants to remember this device
                    let remember_device = response.get("remember_device")
                        .and_then(|v| v.as_bool())
                        .unwrap_or(true); // Default to true if field is missing
                    
                    if !remember_device {
                        // Backend doesn't want to remember this device - remove approval
                        warn!("Backend reported remember_device=false for approved device {} ({}). Removing approval to stop data collection.", 
                              device.name, device.address);
                        
                        if let Err(e) = device_memory.remove_device(&device.address).await {
                            warn!("Failed to remove device {} after remember_device=false: {}", device.address, e);
                        } else {
                            info!("Successfully removed device {} due to remember_device=false", device.address);
                        }
                        continue; // Skip to next device
                    }
                    
                    info!("Device {} is registered with backend", device.address);
                    
                    // Extract device ID from response and store it
                    if let Some(device_id) = response.get("id").and_then(|id| id.as_str()) {
                        if let Err(e) = device_memory.set_backend_device_id(&device.address, device_id.to_string()).await {
                            warn!("Failed to store backend device ID for {}: {}", device.address, e);
                        } else {
                            info!("Stored backend device ID {} for device {}", device_id, device.address);
                        }
                    }
                }
                Err(e) => {
                    warn!("Failed to register device {}: {}", device.address, e);
                }
            }
        }
    }
}

/// Initialize OIDC client from CLI arguments and configuration
async fn initialize_oidc_client(args: &Args, config: &Config) -> Result<OpenIdconnectClient> {
    // Load device config from file
    let default_device_config_path = PathBuf::from("/etc/px-device-identity/device.yml");
    let device_config_path = args.device_config_file.as_ref()
        .unwrap_or(&default_device_config_path);
    
    let device_config = load_device_config(&device_config_path.to_string_lossy())
        .map_err(|e| bcms::Error::Config(format!("Failed to load device config: {}", e)))?;

    // Get issuer URL from CLI args, config, or device config
    let issuer_url = args.issuer_url.as_ref()
        .or(config.issuer_url.as_ref())
        .cloned()
        .unwrap_or_else(|| device_config.get_issuer_url());

    // Get device private key path from CLI args, config, or default
    let private_key_path = args.device_private_key.as_ref()
        .cloned()
        .or_else(|| config.device_private_key_path.as_ref().map(|s| PathBuf::from(s)))
        .unwrap_or_else(|| PathBuf::from("/root/.local/share/px-device-identity/private.pem"));

    // Get client_id from device config (always a UUID)
    let client_id_str = device_config.client_id.to_string();

    info!("Initializing OIDC client with issuer: {}", issuer_url);
    info!("Using client_id: {}", client_id_str);
    info!("Using device config: {:?}", device_config);

    // Load private key from file
    let private_key = read_private_key(&private_key_path.to_string_lossy())
        .map_err(|e| bcms::Error::Config(format!("Failed to load private key: {}", e)))?;

    // Create issuer URL
    let issuer_url = IssuerUrl::new(issuer_url.clone())
        .map_err(|e| bcms::Error::Config(format!("Invalid issuer URL: {}", e)))?;

    // Create client ID
    let client_id = ClientId::new(client_id_str.clone());

    // Discover provider metadata using OpenID Connect async discovery
    let provider_metadata = CoreProviderMetadata::discover_async(issuer_url.clone(), async_http_client).await
        .map_err(|e| bcms::Error::Config(format!("Failed to discover provider metadata: {}", e)))?;

    // Create OIDC client
    let oidc_client = OpenIdconnectClient::new(
        client_id,
        issuer_url,
        provider_metadata,
        private_key,
    );

    info!("OIDC client initialized successfully");
    Ok(oidc_client)
}