//! MRE Elastography service implementation
//!
//! This module provides the main service interface for MRE (Magnetic Resonance
//! Elastography) inverse solver. It handles initialization, training steps,
//! and snapshot generation for visualization.

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

use mre_shared::ipc::PhantomType;
use mre_shared::{LossRecord, MreSnapshot, MreStatus, PhantomConfig, WaveField};
use rtx_mre::config::MreConfig;
use rtx_mre::phantom::PhantomGenerator;
use rtx_mre::solver::MreSolver;

use crate::error::{ServerError, ServerResult};

/// MRE service configuration
#[derive(Debug, Clone)]
pub struct MreServiceConfig {
    /// Wave Net layers
    pub wave_net_layers: usize,
    /// Wave Net hidden dimension
    pub wave_net_hidden: usize,
    /// Fourier features
    pub fourier_features: usize,
    /// Stiffness texture resolution
    pub stiffness_resolution: usize,
    /// Physics loss weight
    pub physics_weight: f32,
    /// Data loss weight
    pub data_weight: f32,
    /// Snapshot resolution
    pub snapshot_resolution: usize,
}

impl Default for MreServiceConfig {
    fn default() -> Self {
        Self {
            wave_net_layers: 4,
            wave_net_hidden: 128,
            fourier_features: 32,
            stiffness_resolution: 64,
            physics_weight: 1.0,
            data_weight: 1.0,
            snapshot_resolution: 64,
        }
    }
}

impl MreServiceConfig {
    /// Creates an MRE config from service config
    fn to_mre_config(&self) -> MreConfig {
        MreConfig::default()
            .with_wave_net_layers(self.wave_net_layers)
            .with_wave_net_hidden(self.wave_net_hidden)
            .with_fourier_features(self.fourier_features)
            .with_stiffness_resolution(self.stiffness_resolution, self.stiffness_resolution)
            .with_loss_weights(self.physics_weight, self.data_weight)
    }

    /// Fast configuration for testing
    #[must_use]
    pub fn fast() -> Self {
        Self {
            wave_net_layers: 2,
            wave_net_hidden: 32,
            fourier_features: 8,
            stiffness_resolution: 32,
            physics_weight: 1.0,
            data_weight: 1.0,
            snapshot_resolution: 32,
        }
    }
}

/// MRE simulation handle
struct MreHandle {
    /// Solver
    solver: MreSolver,
    /// Whether training is in progress
    training: bool,
    /// Last step time in milliseconds
    last_step_ms: f32,
    /// Total training steps
    total_steps: usize,
    /// Start time
    start_time: Instant,
}

impl MreHandle {
    /// Create a new handle
    fn new(solver: MreSolver) -> Self {
        Self {
            solver,
            training: false,
            last_step_ms: 0.0,
            total_steps: 0,
            start_time: Instant::now(),
        }
    }

    /// Get status struct
    fn status(&self) -> MreStatus {
        let elapsed = self.start_time.elapsed().as_secs_f32();
        let steps_per_second = if elapsed > 0.0 {
            self.total_steps as f32 / elapsed
        } else {
            0.0
        };

        MreStatus {
            initialized: true,
            training: self.training,
            step: self.solver.current_step(),
            total_steps: self.total_steps,
            steps_per_second,
        }
    }
}

/// MRE Elastography service
///
/// This service manages the MRE inverse solver lifecycle,
/// handles training steps, and provides snapshots for visualization.
pub struct MreService {
    /// Service configuration
    config: MreServiceConfig,
    /// Active handle (if any)
    handle: Arc<RwLock<Option<MreHandle>>>,
}

impl MreService {
    /// Creates a new MRE service
    #[must_use]
    pub fn new(config: MreServiceConfig) -> Self {
        Self {
            config,
            handle: Arc::new(RwLock::new(None)),
        }
    }

    /// Creates a service with default configuration
    #[must_use]
    pub fn with_defaults() -> Self {
        Self::new(MreServiceConfig::default())
    }

    /// Initialize with wave field data
    ///
    /// # Errors
    ///
    /// Returns error if already initialized
    pub async fn initialize(&self, wave: WaveField) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;

        if handle_guard.is_some() {
            return Err(ServerError::AlreadyInitialized);
        }

        let mre_config = self.config.to_mre_config();
        let mut solver = MreSolver::new(mre_config)
            .map_err(|e| ServerError::internal(format!("Failed to create solver: {e}")))?;

        solver.set_measured_wave(wave);

        *handle_guard = Some(MreHandle::new(solver));
        Ok(())
    }

    /// Initialize with a phantom configuration
    ///
    /// Creates synthetic wave field data from phantom for testing
    pub async fn initialize_phantom(&self, phantom_type: PhantomType) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;

        if handle_guard.is_some() {
            return Err(ServerError::AlreadyInitialized);
        }

        let mre_config = self.config.to_mre_config();
        let phantom_config = match phantom_type {
            PhantomType::SingleTumor => PhantomConfig::single_tumor(),
            PhantomType::MultipleLesions => PhantomConfig::multiple_lesions(),
            PhantomType::Layered => PhantomConfig::single_tumor(), // Use single_tumor for now
        };

        let phantom = PhantomGenerator::new(phantom_config, mre_config.clone());
        let (_stiffness, wave) = phantom.generate();

        let mut solver = MreSolver::new(mre_config)
            .map_err(|e| ServerError::internal(format!("Failed to create solver: {e}")))?;

        solver.set_measured_wave(wave);

        *handle_guard = Some(MreHandle::new(solver));
        Ok(())
    }

    /// Reset the service
    pub async fn reset(&self) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;
        *handle_guard = None;
        Ok(())
    }

    /// Check if initialized
    pub async fn is_initialized(&self) -> bool {
        self.handle.read().await.is_some()
    }

    /// Run a single training step
    ///
    /// # Returns
    ///
    /// Loss record from the step
    pub async fn step(&self) -> ServerResult<LossRecord> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();
        handle.training = true;

        let loss = handle
            .solver
            .step()
            .map_err(|e| ServerError::training(format!("Training step failed: {e}")))?;

        handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
        handle.total_steps += 1;

        Ok(loss)
    }

    /// Run multiple training steps
    ///
    /// # Arguments
    ///
    /// * `num_steps` - Number of steps to run
    ///
    /// # Returns
    ///
    /// Vector of loss records
    pub async fn train(&self, num_steps: usize) -> ServerResult<Vec<LossRecord>> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();
        handle.training = true;

        let losses = handle
            .solver
            .train(num_steps)
            .map_err(|e| ServerError::training(format!("Training failed: {e}")))?;

        handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
        handle.total_steps += num_steps;

        Ok(losses)
    }

    /// Get a snapshot of current state
    ///
    /// Returns stiffness map, wave field, and residual for visualization
    pub async fn snapshot(&self) -> ServerResult<MreSnapshot> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        let res = self.config.snapshot_resolution;
        handle
            .solver
            .snapshot(res, res)
            .map_err(|e| ServerError::inference(format!("Snapshot failed: {e}")))
    }

    /// Get current status
    pub async fn status(&self) -> ServerResult<MreStatus> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(handle.status()),
            None => Ok(MreStatus::default()),
        }
    }

    /// Get loss history
    pub async fn loss_history(&self) -> ServerResult<Vec<LossRecord>> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.solver.loss_history().to_vec())
    }

    /// Get current step count
    pub async fn current_step(&self) -> ServerResult<usize> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.solver.current_step())
    }

    /// Get performance metrics
    pub async fn metrics(&self) -> ServerResult<MreMetrics> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(MreMetrics {
                step_time_ms: handle.last_step_ms,
                total_steps: handle.total_steps,
                elapsed_secs: handle.start_time.elapsed().as_secs_f32(),
            }),
            None => Ok(MreMetrics::default()),
        }
    }
}

/// Performance metrics for MRE service
#[derive(Debug, Clone, Default)]
pub struct MreMetrics {
    /// Last step time in milliseconds
    pub step_time_ms: f32,
    /// Total steps executed
    pub total_steps: usize,
    /// Total elapsed time in seconds
    pub elapsed_secs: f32,
}

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

    async fn create_test_service() -> MreService {
        MreService::new(MreServiceConfig::fast())
    }

    #[tokio::test]
    async fn test_service_creation() {
        let service = create_test_service().await;
        assert!(!service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_initialize_phantom() {
        let service = create_test_service().await;

        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();
        assert!(service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_double_initialize_fails() {
        let service = create_test_service().await;

        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();
        let result = service.initialize_phantom(PhantomType::SingleTumor).await;

        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ServerError::AlreadyInitialized
        ));
    }

    #[tokio::test]
    #[ignore = "Pre-existing MRE service step error"]
    async fn test_step() {
        let service = create_test_service().await;
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        let loss = service.step().await.unwrap();
        assert!(loss.total_loss.is_finite());
        assert_eq!(loss.step, 1);
    }

    #[tokio::test]
    #[ignore = "Pre-existing MRE service train error"]
    async fn test_train() {
        let service = create_test_service().await;
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        let losses = service.train(5).await.unwrap();
        assert_eq!(losses.len(), 5);
    }

    #[tokio::test]
    async fn test_snapshot() {
        let service = create_test_service().await;
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        let snapshot = service.snapshot().await.unwrap();
        assert!(!snapshot.stiffness.values.is_empty());
        assert!(!snapshot.wave_real.is_empty());
    }

    #[tokio::test]
    async fn test_reset() {
        let service = create_test_service().await;
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        service.reset().await.unwrap();
        assert!(!service.is_initialized().await);

        // Should be able to initialize again
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();
        assert!(service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_status() {
        let service = create_test_service().await;

        let status = service.status().await.unwrap();
        assert!(!status.initialized);

        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        let status = service.status().await.unwrap();
        assert!(status.initialized);
    }

    #[tokio::test]
    #[ignore = "Pre-existing MRE service loss history error"]
    async fn test_loss_history() {
        let service = create_test_service().await;
        service
            .initialize_phantom(PhantomType::SingleTumor)
            .await
            .unwrap();

        service.train(3).await.unwrap();

        let history = service.loss_history().await.unwrap();
        assert_eq!(history.len(), 3);
    }
}