//! Application state management for Tauri
//!
//! This module provides thread-safe wrappers around the HemodynamicsService,
//! MreService, BioheatService, SlidescopeService, NeuralOperatorDemo, PiddmDemo,
//! and DigitalTwinDemo that can be accessed by Tauri commands via the state management API.

use std::sync::Arc;
use tokio::sync::RwLock;

use rtx_hemodynamics_server::{
    BioheatService, HemodynamicsService, MreService, ServiceConfig, SlidescopeService,
    SlidescopeServiceConfig,
};
use rtx_neural_operator_demo::{NeuralOperatorDemo, TrainingSession};
use rtx_neural_operator_shared::config::PDEConfig;
use rtx_neural_operator_shared::ipc::{TrainingConfig, TrainingProgress};

// PIDDM demo imports
use rtx_piddm_demo::{PiddmTrainer, PiddmSampler};

// Digital Twin demo imports
use rtx_digital_twin_demo::TwinSimulator;

// Image Classifier demo imports
use rtx_image_classifier_demo::ImageClassifier;

// Time Series Forecast demo imports
use rtx_timeseries_demo::TimeSeriesForecaster;

// Portfolio Optimizer demo imports
use rtx_portfolio_demo::PortfolioOptimizer;

// Risk Analyzer demo imports
use rtx_risk_analyzer::RiskAnalyzer;

// PINN Benchmark demo imports
use rtx_pinn_benchmark::BenchmarkRunner;

/// Tauri application state containing the simulation services
pub struct AppState {
    /// The hemodynamics service instance
    service: Arc<RwLock<HemodynamicsService>>,
    /// The MRE elastography service instance
    mre_service: Arc<MreService>,
    /// The thermal ablation bioheat service instance
    bioheat_service: Arc<BioheatService>,
    /// The SlideScope pathology service instance
    slidescope_service: Arc<SlidescopeService>,
    /// The Neural Operator demo instance
    neural_operator_demo: Arc<RwLock<NeuralOperatorDemo>>,
    /// The Neural Operator training session (optional)
    neural_operator_training: Arc<RwLock<Option<TrainingSession>>>,
    /// The PIDDM trainer instance
    piddm_trainer: Arc<RwLock<Option<PiddmTrainer>>>,
    /// The PIDDM sampler instance
    piddm_sampler: Arc<RwLock<Option<PiddmSampler>>>,
    /// The Digital Twin simulator instance
    digital_twin_simulator: Arc<RwLock<Option<TwinSimulator>>>,
    /// The Image Classifier instance
    image_classifier: Arc<RwLock<Option<ImageClassifier>>>,
    /// The Time Series Forecaster instance
    timeseries_forecaster: Arc<RwLock<Option<TimeSeriesForecaster>>>,
    /// The Portfolio Optimizer instance
    portfolio_optimizer: Arc<RwLock<Option<PortfolioOptimizer>>>,
    /// The Risk Analyzer instance
    risk_analyzer: Arc<RwLock<Option<RiskAnalyzer>>>,
    /// The PINN Benchmark runner instance
    pinn_benchmark: Arc<RwLock<Option<BenchmarkRunner>>>,
}

impl AppState {
    /// Creates a new application state with default configuration
    #[must_use]
    pub fn new() -> Self {
        let config = ServiceConfig::default();
        let service = HemodynamicsService::new(config);
        let mre_service = MreService::with_defaults();
        let bioheat_service = BioheatService::with_defaults();
        let slidescope_service = SlidescopeService::new(SlidescopeServiceConfig::default());
        let neural_operator_demo = NeuralOperatorDemo::new();
        Self {
            service: Arc::new(RwLock::new(service)),
            mre_service: Arc::new(mre_service),
            bioheat_service: Arc::new(bioheat_service),
            slidescope_service: Arc::new(slidescope_service),
            neural_operator_demo: Arc::new(RwLock::new(neural_operator_demo)),
            neural_operator_training: Arc::new(RwLock::new(None)),
            piddm_trainer: Arc::new(RwLock::new(None)),
            piddm_sampler: Arc::new(RwLock::new(None)),
            digital_twin_simulator: Arc::new(RwLock::new(None)),
            image_classifier: Arc::new(RwLock::new(None)),
            timeseries_forecaster: Arc::new(RwLock::new(None)),
            portfolio_optimizer: Arc::new(RwLock::new(None)),
            risk_analyzer: Arc::new(RwLock::new(None)),
            pinn_benchmark: Arc::new(RwLock::new(None)),
        }
    }

    /// Creates a new application state with custom configuration
    #[must_use]
    pub fn with_config(config: ServiceConfig) -> Self {
        let service = HemodynamicsService::new(config);
        let mre_service = MreService::with_defaults();
        let bioheat_service = BioheatService::with_defaults();
        let slidescope_service = SlidescopeService::new(SlidescopeServiceConfig::default());
        let neural_operator_demo = NeuralOperatorDemo::new();
        Self {
            service: Arc::new(RwLock::new(service)),
            mre_service: Arc::new(mre_service),
            bioheat_service: Arc::new(bioheat_service),
            slidescope_service: Arc::new(slidescope_service),
            neural_operator_demo: Arc::new(RwLock::new(neural_operator_demo)),
            neural_operator_training: Arc::new(RwLock::new(None)),
            piddm_trainer: Arc::new(RwLock::new(None)),
            piddm_sampler: Arc::new(RwLock::new(None)),
            digital_twin_simulator: Arc::new(RwLock::new(None)),
            image_classifier: Arc::new(RwLock::new(None)),
            timeseries_forecaster: Arc::new(RwLock::new(None)),
            portfolio_optimizer: Arc::new(RwLock::new(None)),
            risk_analyzer: Arc::new(RwLock::new(None)),
            pinn_benchmark: Arc::new(RwLock::new(None)),
        }
    }

    /// Returns a reference to the hemodynamics service
    #[must_use]
    pub fn service(&self) -> &Arc<RwLock<HemodynamicsService>> {
        &self.service
    }

    /// Returns a reference to the MRE service
    #[must_use]
    pub fn mre_service(&self) -> &Arc<MreService> {
        &self.mre_service
    }

    /// Returns a reference to the bioheat service
    #[must_use]
    pub fn bioheat_service(&self) -> &Arc<BioheatService> {
        &self.bioheat_service
    }

    /// Returns a reference to the SlideScope service
    #[must_use]
    pub fn slidescope_service(&self) -> &Arc<SlidescopeService> {
        &self.slidescope_service
    }

    /// Returns a reference to the Neural Operator demo
    #[must_use]
    pub fn neural_operator_demo(&self) -> &Arc<RwLock<NeuralOperatorDemo>> {
        &self.neural_operator_demo
    }

    /// Returns a reference to the Neural Operator training session
    #[must_use]
    pub fn neural_operator_training(&self) -> &Arc<RwLock<Option<TrainingSession>>> {
        &self.neural_operator_training
    }

    /// Returns a reference to the PIDDM trainer
    #[must_use]
    pub fn piddm_trainer(&self) -> &Arc<RwLock<Option<PiddmTrainer>>> {
        &self.piddm_trainer
    }

    /// Returns a reference to the PIDDM sampler
    #[must_use]
    pub fn piddm_sampler(&self) -> &Arc<RwLock<Option<PiddmSampler>>> {
        &self.piddm_sampler
    }

    /// Returns a reference to the Digital Twin simulator
    #[must_use]
    pub fn digital_twin_simulator(&self) -> &Arc<RwLock<Option<TwinSimulator>>> {
        &self.digital_twin_simulator
    }

    /// Returns a reference to the Image Classifier
    #[must_use]
    pub fn image_classifier(&self) -> &Arc<RwLock<Option<ImageClassifier>>> {
        &self.image_classifier
    }

    /// Returns a reference to the Time Series Forecaster
    #[must_use]
    pub fn timeseries_forecaster(&self) -> &Arc<RwLock<Option<TimeSeriesForecaster>>> {
        &self.timeseries_forecaster
    }

    /// Returns a reference to the Portfolio Optimizer
    #[must_use]
    pub fn portfolio_optimizer(&self) -> &Arc<RwLock<Option<PortfolioOptimizer>>> {
        &self.portfolio_optimizer
    }

    /// Returns a reference to the Risk Analyzer
    #[must_use]
    pub fn risk_analyzer(&self) -> &Arc<RwLock<Option<RiskAnalyzer>>> {
        &self.risk_analyzer
    }

    /// Returns a reference to the PINN Benchmark runner
    #[must_use]
    pub fn pinn_benchmark(&self) -> &Arc<RwLock<Option<BenchmarkRunner>>> {
        &self.pinn_benchmark
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_app_state_creation() {
        let _state = AppState::new();
    }

    #[test]
    fn test_app_state_default() {
        let state = AppState::default();
        assert!(Arc::strong_count(state.service()) >= 1);
        assert!(Arc::strong_count(state.mre_service()) >= 1);
        assert!(Arc::strong_count(state.bioheat_service()) >= 1);
        assert!(Arc::strong_count(state.slidescope_service()) >= 1);
        assert!(Arc::strong_count(state.neural_operator_demo()) >= 1);
        assert!(Arc::strong_count(state.neural_operator_training()) >= 1);
        assert!(Arc::strong_count(state.piddm_trainer()) >= 1);
        assert!(Arc::strong_count(state.piddm_sampler()) >= 1);
        assert!(Arc::strong_count(state.digital_twin_simulator()) >= 1);
        assert!(Arc::strong_count(state.image_classifier()) >= 1);
        assert!(Arc::strong_count(state.timeseries_forecaster()) >= 1);
    }

    #[test]
    fn test_app_state_with_config() {
        let config = ServiceConfig::default();
        let state = AppState::with_config(config);
        assert!(Arc::strong_count(state.service()) >= 1);
        assert!(Arc::strong_count(state.mre_service()) >= 1);
        assert!(Arc::strong_count(state.bioheat_service()) >= 1);
        assert!(Arc::strong_count(state.slidescope_service()) >= 1);
        assert!(Arc::strong_count(state.neural_operator_demo()) >= 1);
        assert!(Arc::strong_count(state.neural_operator_training()) >= 1);
        assert!(Arc::strong_count(state.piddm_trainer()) >= 1);
        assert!(Arc::strong_count(state.piddm_sampler()) >= 1);
        assert!(Arc::strong_count(state.digital_twin_simulator()) >= 1);
        assert!(Arc::strong_count(state.image_classifier()) >= 1);
        assert!(Arc::strong_count(state.timeseries_forecaster()) >= 1);
    }
}