bcms 0.1.12

use crate::api::BackendAPI;
use crate::bcms_capnp::b_c_m_s as bcms;
use crate::ble::BleManager;
use crate::data::{DataStore, IoTData};
use crate::devices::{DeviceMemory, WorkingMode};
use crate::{Error, Result};
use capnp::capability::Promise;
use capnp_rpc::{pry, rpc_twoparty_capnp, twoparty, RpcSystem};
use futures_util::io::AsyncReadExt;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Mutex;
use tracing::{error, info, warn};

/// RPC Server implementation
pub struct RpcServer {
    bind_address: SocketAddr,
    device_memory: Arc<DeviceMemory>,
    ble_manager: Arc<BleManager>,
    data_store: Arc<DataStore>,
    api_client: Option<Arc<Mutex<BackendAPI>>>,
}

impl RpcServer {
    pub fn new(
        bind_address: SocketAddr,
        device_memory: Arc<DeviceMemory>,
        ble_manager: Arc<BleManager>,
        data_store: Arc<DataStore>,
    ) -> Self {
        Self {
            bind_address,
            device_memory,
            ble_manager,
            data_store,
            api_client: None,
        }
    }

    pub fn with_api_client(
        bind_address: SocketAddr,
        device_memory: Arc<DeviceMemory>,
        ble_manager: Arc<BleManager>,
        data_store: Arc<DataStore>,
        api_client: Arc<Mutex<BackendAPI>>,
    ) -> Self {
        Self {
            bind_address,
            device_memory,
            ble_manager,
            data_store,
            api_client: Some(api_client),
        }
    }

    /// Start the RPC server
    pub async fn start(&self) -> Result<()> {
        let listener = TcpListener::bind(&self.bind_address).await?;
        info!("RPC server listening on {}", self.bind_address);

        loop {
            match listener.accept().await {
                Ok((stream, addr)) => {
                    info!("New RPC connection from {}", addr);
                    let device_memory = Arc::clone(&self.device_memory);
                    let ble_manager = Arc::clone(&self.ble_manager);
                    let data_store = Arc::clone(&self.data_store);

                    // Handle connection directly without spawning since RpcSystem is not Send
                    if let Err(e) = Self::handle_connection(stream, device_memory, ble_manager, data_store, self.api_client.clone()).await {
                        error!("RPC connection error: {}", e);
                    }
                }
                Err(e) => {
                    error!("Failed to accept RPC connection: {}", e);
                }
            }
        }
    }

    /// Handle a single RPC connection
    async fn handle_connection(
        stream: TcpStream,
        device_memory: Arc<DeviceMemory>,
        ble_manager: Arc<BleManager>,
        data_store: Arc<DataStore>,
        api_client: Option<Arc<Mutex<BackendAPI>>>,
    ) -> Result<()> {
        tracing::debug!("RPC server: handling new connection");
        let (reader, writer) = tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();

        let network = twoparty::VatNetwork::new(
            reader,
            writer,
            rpc_twoparty_capnp::Side::Server,
            Default::default(),
        );

        let bcms_impl = BcmsImpl::new(device_memory, ble_manager, data_store, api_client);
        let bcms_client: bcms::Client = capnp_rpc::new_client(bcms_impl);

        tracing::debug!("RPC server: starting RPC system");
        let mut rpc_system = RpcSystem::new(Box::new(network), Some(bcms_client.client));
        rpc_system.await?;
        tracing::debug!("RPC server: connection closed");

        Ok(())
    }
}

/// Implementation of the BCMS Cap'n Proto interface
struct BcmsImpl {
    device_memory: Arc<DeviceMemory>,
    ble_manager: Arc<BleManager>,
    data_store: Arc<DataStore>,
    api_client: Option<Arc<Mutex<BackendAPI>>>,
}

impl BcmsImpl {
    fn new(device_memory: Arc<DeviceMemory>, ble_manager: Arc<BleManager>, data_store: Arc<DataStore>, api_client: Option<Arc<Mutex<BackendAPI>>>) -> Self {
        Self {
            device_memory,
            ble_manager,
            data_store,
            api_client,
        }
    }

    /// Convert WorkingMode to Cap'n Proto enum
    fn working_mode_to_capnp(mode: WorkingMode) -> crate::bcms_capnp::BCMSWorkingMode {
        match mode {
            WorkingMode::Any => crate::bcms_capnp::BCMSWorkingMode::Any,
            WorkingMode::Approved => crate::bcms_capnp::BCMSWorkingMode::Approved,
        }
    }

    /// Convert Cap'n Proto enum to WorkingMode
    fn working_mode_from_capnp(mode: crate::bcms_capnp::BCMSWorkingMode) -> WorkingMode {
        match mode {
            crate::bcms_capnp::BCMSWorkingMode::Any => WorkingMode::Any,
            crate::bcms_capnp::BCMSWorkingMode::Approved => WorkingMode::Approved,
        }
    }
}

impl bcms::Server for BcmsImpl {
    fn list(
        &mut self,
        params: bcms::ListParams,
        mut results: bcms::ListResults,
    ) -> Promise<(), capnp::Error> {
        tracing::debug!("RPC server: list method called");
        let device_memory = Arc::clone(&self.device_memory);
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let only_approved = params.get_only_approved();

            let mut devices = device_memory.get_devices_filtered(only_approved).await;
            tracing::debug!("RPC list: only_approved={}, found {} devices", only_approved, devices.len());

            // Update pairing and connection status from actual Bluetooth stack for each device
            for device in &mut devices {
                // Check actual pairing status
                match ble_manager.is_device_paired(&device.address).await {
                    Ok(actual_paired) => {
                        tracing::debug!("  Device: {} ({}), approved={}, paired (memory)={}, paired (bluetooth)={}",
                            device.name, device.address, device.approved, device.paired, actual_paired);
                        // Use actual Bluetooth pairing status
                        device.paired = actual_paired;
                    }
                    Err(e) => {
                        tracing::warn!("Failed to check actual pairing status for {}: {}", device.address, e);
                        // Keep the memory status if we can't check actual status
                    }
                }

                // Check actual connection status
                match ble_manager.is_device_connected(&device.address).await {
                    Ok(actual_connected) => {
                        tracing::debug!("  Device: {} connected (bluetooth)={}", device.address, actual_connected);
                        device.connected = actual_connected;
                    }
                    Err(e) => {
                        tracing::debug!("Failed to check connection status for {}: {}", device.address, e);
                        // Keep memory status on error
                    }
                }
            }

            let mut results_builder = results.get();
            {
                let mut devices_builder = results_builder.reborrow().init_devices(devices.len() as u32);

                for (i, device) in devices.iter().enumerate() {
                    let mut device_builder = devices_builder.reborrow().get(i as u32);
                    device_builder.set_address(&device.address);
                    device_builder.set_name(&device.name);
                    device_builder.set_approved(device.approved);
                    device_builder.set_paired(device.paired);
                    device_builder.set_device_class(device.device_class);
                    device_builder.set_appearance(device.appearance);
                    device_builder.set_connected(device.connected);
                }
            }
            results_builder.init_errors(0);

            Ok(())
        })
    }

    fn approve(
        &mut self,
        params: bcms::ApproveParams,
        mut results: bcms::ApproveResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);
        let api_client = self.api_client.clone();

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            match device_memory.approve_device(&address.to_string()?).await {
                Ok(success) if success => {
                    // Device approved successfully - trigger immediate API registration if available
                    if let Some(api_client) = api_client {
                        match device_memory.get_device(&address.to_string()?).await {
                            Some(device) => {
                                // Use tokio::spawn to prevent panics from blocking the RPC response
                                let api_client_clone = Arc::clone(&api_client);
                                let device_address = device.address.clone();
                                let device_memory_clone = Arc::clone(&device_memory);
                                tokio::spawn(async move {
                                    let mut api = api_client_clone.lock().await;
                                    match api.create_iot_device_if_not_exists(
                                        &device_address,
                                        Some(&device.name),
                                        device.connection_type.as_deref(),
                                        None,
                                        None,
                                    ).await {
                                        Ok(response) => {
                                            info!("Device {} immediately registered with backend API after approval", device_address);

                                            // Extract and store device ID
                                            if let Some(device_id) = response.get("id").and_then(|id| id.as_str()) {
                                                if let Err(e) = device_memory_clone.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 approved device {}", device_id, device_address);
                                                }
                                            }
                                        }
                                        Err(e) => {
                                            warn!("Failed to immediately register approved device {}: {}", device_address, e);
                                        }
                                    }
                                });
                            }
                            None => {
                                warn!("Approved device {} not found in memory", address.to_string()?);
                            }
                        }
                    }

                    let mut results_builder = results.get();
                    results_builder.set_status(true);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Ok(_) => {
                    // Device not found or approval failed
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, "Device not found or already approved");
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to approve device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn remove(
        &mut self,
        params: bcms::RemoveParams,
        mut results: bcms::RemoveResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            match device_memory.remove_device(&address.to_string()?).await {
                Ok(success) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(success);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to remove device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn mode(
        &mut self,
        _params: bcms::ModeParams,
        mut results: bcms::ModeResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);

        Promise::from_future(async move {
            let mode = device_memory.get_working_mode().await;
            let mut results_builder = results.get();

            let capnp_mode = Self::working_mode_to_capnp(mode);
            results_builder.set_mode(capnp_mode);

            Ok(())
        })
    }

    fn set_mode(
        &mut self,
        params: bcms::SetModeParams,
        mut results: bcms::SetModeResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);

        Promise::from_future(async move {
            let params = params.get()?;
            let mode_capnp = params.get_mode()?;
            let mode = Self::working_mode_from_capnp(mode_capnp);

            match device_memory.set_working_mode(mode).await {
                Ok(()) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(true);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to set mode: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn pair(
        &mut self,
        params: bcms::PairParams,
        mut results: bcms::PairResults,
    ) -> Promise<(), capnp::Error> {
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            match ble_manager.pair_device(&address.to_string()?).await {
                Ok(success) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(success);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to pair device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn unpair(
        &mut self,
        params: bcms::UnpairParams,
        mut results: bcms::UnpairResults,
    ) -> Promise<(), capnp::Error> {
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            match ble_manager.unpair_device(&address.to_string()?).await {
                Ok(success) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(success);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to unpair device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn is_paired(
        &mut self,
        params: bcms::IsPairedParams,
        mut results: bcms::IsPairedResults,
    ) -> Promise<(), capnp::Error> {
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            match ble_manager.is_device_paired(&address.to_string()?).await {
                Ok(is_paired) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(is_paired);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to check pairing status: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn connect(
        &mut self,
        params: bcms::ConnectParams,
        mut results: bcms::ConnectResults,
    ) -> Promise<(), capnp::Error> {
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            info!("RPC connect: attempting to connect to device {}", address.to_str()?);

            match ble_manager.connect_device(&address.to_string()?).await {
                Ok(success) => {
                    info!("RPC connect: device {} connection result: {}", address.to_str()?, success);
                    let mut results_builder = results.get();
                    results_builder.set_status(success);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    error!("RPC connect: failed to connect to device {}: {}", address.to_str()?, e);
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to connect to device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn disconnect(
        &mut self,
        params: bcms::DisconnectParams,
        mut results: bcms::DisconnectResults,
    ) -> Promise<(), capnp::Error> {
        let ble_manager = Arc::clone(&self.ble_manager);

        Promise::from_future(async move {
            let params = params.get()?;
            let address = params.get_address()?;

            info!("RPC disconnect: attempting to disconnect from device {}", address.to_str()?);

            match ble_manager.disconnect_device(&address.to_string()?).await {
                Ok(success) => {
                    info!("RPC disconnect: device {} disconnection result: {}", address.to_str()?, success);
                    let mut results_builder = results.get();
                    results_builder.set_status(success);
                    results_builder.init_errors(0);
                    Ok(())
                }
                Err(e) => {
                    error!("RPC disconnect: failed to disconnect from device {}: {}", address.to_str()?, e);
                    let mut results_builder = results.get();
                    results_builder.set_status(false);
                    let mut errors_builder = results_builder.init_errors(1);
                    errors_builder.set(0, &format!("Failed to disconnect from device: {}", e));
                    Ok(())
                }
            }
        })
    }

    fn submit_data(
        &mut self,
        params: bcms::SubmitDataParams,
        mut results: bcms::SubmitDataResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);
        let data_store = Arc::clone(&self.data_store);

        Promise::from_future(async move {
            let params = params.get()?;
            let payload = params.get_payload()?;

            // Extract device info
            let device_info = payload.get_device()?;
            let address = device_info.get_address()?.to_string()?;
            let name = device_info.get_name()?.to_string()?;
            let connection_type = device_info.get_connection_type()?.to_string()?;

            // Auto-approve the device since it's from RPC (local submission)
            if let Err(e) = device_memory.add_device_with_auto_approval(&address, &name, Some(connection_type)).await {
                warn!("Failed to auto-approve RPC device {}: {}", address, e);
            }

            // Parse IoT data from JSON strings
            let data_list = payload.get_data()?;
            let mut iot_data_batch = Vec::new();
            let mut errors = Vec::new();

            info!("RPC submit_data: {} data points from {} ({})", data_list.len(), name, address);
            for i in 0..data_list.len() {
                let data_json = data_list.get(i)?.to_string()?;
                info!("  -> {}", data_json);
                match serde_json::from_str::<IoTData>(&data_json) {
                    Ok(iot_data) => {
                        iot_data_batch.push(iot_data);
                    }
                    Err(e) => {
                        errors.push(format!("Failed to parse IoT data at index {}: {}", i, e));
                    }
                }
            }

            // Store the data in DataStore
            if !iot_data_batch.is_empty() {
                if let Err(e) = data_store.add_data_batch(iot_data_batch.clone()).await {
                    errors.push(format!("Failed to store IoT data: {}", e));
                } else {
                    // Update device last_seen timestamp based on latest data timestamp
                    let latest_timestamp = iot_data_batch.iter()
                        .map(|data| data.timestamp())
                        .max()
                        .unwrap_or_else(|| chrono::Utc::now().timestamp());

                    if let Err(e) = device_memory.update_last_seen(&address, latest_timestamp).await {
                        warn!("Failed to update last_seen for device {}: {}", address, e);
                    }
                }
            }

            // Build response
            let mut results_builder = results.get();
            results_builder.set_status(errors.is_empty());

            if !errors.is_empty() {
                let mut errors_builder = results_builder.init_errors(errors.len() as u32);
                for (i, error) in errors.iter().enumerate() {
                    errors_builder.set(i as u32, error);
                }
            } else {
                results_builder.init_errors(0);
            }

            Ok(())
        })
    }

    fn stats(
        &mut self,
        _params: bcms::StatsParams,
        mut results: bcms::StatsResults,
    ) -> Promise<(), capnp::Error> {
        let device_memory = Arc::clone(&self.device_memory);
        let data_store = Arc::clone(&self.data_store);

        Promise::from_future(async move {
            let mut errors: Vec<String> = Vec::new();
            let mut total_data_points = 0u32;

            // Get only approved devices (the ones we track)
            let devices = device_memory.get_approved_devices().await;
            let mut device_stats = Vec::new();

            for device in devices {
                // Get data statistics for this device
                let data_summary = data_store.get_device_data_summary(&device.address).await;

                let mut data_type_stats = Vec::new();
                for (data_type, (count, last_timestamp)) in data_summary {
                    data_type_stats.push((data_type, count, last_timestamp));
                    total_data_points += count;
                }

                device_stats.push((device, data_type_stats));
            }

            // Build response
            let mut results_builder = results.get();
            let mut stats_result_builder = results_builder.init_result();

            // Set devices
            let mut devices_builder = stats_result_builder.reborrow().init_devices(device_stats.len() as u32);
            for (i, (device, data_types)) in device_stats.iter().enumerate() {
                let mut device_builder = devices_builder.reborrow().get(i as u32);
                device_builder.set_address(&device.address);
                device_builder.set_name(&device.name);
                device_builder.set_approved(device.approved);
                device_builder.set_paired(device.paired);
                device_builder.set_connection_type(&device.connection_type.clone().unwrap_or_else(|| "bluetooth_le".to_string()));
                device_builder.set_last_seen(device.last_seen.unwrap_or(0));
                device_builder.set_backend_registered(device.backend_device_id.is_some());

                // Set data type stats
                let mut data_types_builder = device_builder.init_data_types(data_types.len() as u32);
                for (j, (data_type, count, last_timestamp)) in data_types.iter().enumerate() {
                    let mut data_type_builder = data_types_builder.reborrow().get(j as u32);
                    data_type_builder.set_data_type(data_type);
                    data_type_builder.set_count(*count);
                    data_type_builder.set_last_timestamp(*last_timestamp);
                }
            }

            // Set total data points
            stats_result_builder.set_total_data_points(total_data_points);

            // Set errors
            if !errors.is_empty() {
                let mut errors_builder = stats_result_builder.init_errors(errors.len() as u32);
                for (i, error) in errors.iter().enumerate() {
                    errors_builder.set(i as u32, error);
                }
            } else {
                stats_result_builder.init_errors(0);
            }

            Ok(())
        })
    }
}