foxhunt/ml/src/deployment/endpoints.rs

//! gRPC endpoints for model deployment management
//! 
//! This module provides gRPC service implementations for managing model deployments,
//! including deployment operations, monitoring, and administration.

use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use tonic::{Request, Response, Status, Code};
use uuid::Uuid;
use serde::{Serialize, Deserialize};
use prost::Message;

use crate::types::{MLResult, MLError};
use crate::traits::MLModel;
use super::{
    ModelDeploymentRegistry, DeploymentConfig, DeploymentStrategy,
    versioning::ModelVersion,
    ab_testing::ABTestConfig,
    monitoring::MonitoringConfig,
};

// Protocol buffer definitions (would normally be in a .proto file)
#[derive(Clone, PartialEq, Message)]
pub struct DeployModelRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
    #[prost(string, tag = "2")]
    pub model_type: String,
    #[prost(bytes, tag = "3")]
    pub model_data: Vec<u8>,
    #[prost(string, tag = "4")]
    pub version: String,
    #[prost(message, optional, tag = "5")]
    pub config: Option<DeploymentConfigProto>,
}

#[derive(Clone, PartialEq, Message)]
pub struct DeployModelResponse {
    #[prost(string, tag = "1")]
    pub deployment_id: String,
    #[prost(bool, tag = "2")]
    pub success: bool,
    #[prost(string, tag = "3")]
    pub message: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetDeploymentRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetDeploymentResponse {
    #[prost(message, optional, tag = "1")]
    pub deployment: Option<DeploymentInfoProto>,
    #[prost(bool, tag = "2")]
    pub found: bool,
}

#[derive(Clone, PartialEq, Message)]
pub struct ListDeploymentsRequest {
    #[prost(string, optional, tag = "1")]
    pub filter: Option<String>,
    #[prost(int32, tag = "2")]
    pub limit: i32,
    #[prost(int32, tag = "3")]
    pub offset: i32,
}

#[derive(Clone, PartialEq, Message)]
pub struct ListDeploymentsResponse {
    #[prost(message, repeated, tag = "1")]
    pub deployments: Vec<DeploymentInfoProto>,
    #[prost(int32, tag = "2")]
    pub total_count: i32,
}

#[derive(Clone, PartialEq, Message)]
pub struct RollbackModelRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
    #[prost(string, optional, tag = "2")]
    pub target_version: Option<String>,
}

#[derive(Clone, PartialEq, Message)]
pub struct RollbackModelResponse {
    #[prost(bool, tag = "1")]
    pub success: bool,
    #[prost(string, tag = "2")]
    pub message: String,
    #[prost(string, tag = "3")]
    pub new_version: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct UndeployModelRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct UndeployModelResponse {
    #[prost(bool, tag = "1")]
    pub success: bool,
    #[prost(string, tag = "2")]
    pub message: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetModelMetricsRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
    #[prost(int64, optional, tag = "2")]
    pub start_time: Option<i64>,
    #[prost(int64, optional, tag = "3")]
    pub end_time: Option<i64>,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetModelMetricsResponse {
    #[prost(message, optional, tag = "1")]
    pub metrics: Option<ModelMetricsProto>,
}

#[derive(Clone, PartialEq, Message)]
pub struct StartABTestRequest {
    #[prost(string, tag = "1")]
    pub model_id: String,
    #[prost(string, tag = "2")]
    pub treatment_model_type: String,
    #[prost(bytes, tag = "3")]
    pub treatment_model_data: Vec<u8>,
    #[prost(string, tag = "4")]
    pub treatment_version: String,
    #[prost(message, optional, tag = "5")]
    pub ab_test_config: Option<ABTestConfigProto>,
}

#[derive(Clone, PartialEq, Message)]
pub struct StartABTestResponse {
    #[prost(string, tag = "1")]
    pub experiment_id: String,
    #[prost(bool, tag = "2")]
    pub success: bool,
    #[prost(string, tag = "3")]
    pub message: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetABTestStatusRequest {
    #[prost(string, tag = "1")]
    pub experiment_id: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct GetABTestStatusResponse {
    #[prost(message, optional, tag = "1")]
    pub status: Option<ABTestStatusProto>,
}

// Supporting message types
#[derive(Clone, PartialEq, Message)]
pub struct DeploymentConfigProto {
    #[prost(bool, tag = "1")]
    pub enable_ab_testing: bool,
    #[prost(message, optional, tag = "2")]
    pub ab_test_config: Option<ABTestConfigProto>,
    #[prost(bool, tag = "3")]
    pub validation_required: bool,
    #[prost(message, optional, tag = "4")]
    pub monitoring_config: Option<MonitoringConfigProto>,
    #[prost(bool, tag = "5")]
    pub auto_rollback: bool,
    #[prost(enumeration = "DeploymentStrategyProto", tag = "6")]
    pub deployment_strategy: i32,
}

#[derive(Clone, PartialEq, Message)]
pub struct DeploymentInfoProto {
    #[prost(string, tag = "1")]
    pub deployment_id: String,
    #[prost(string, tag = "2")]
    pub model_id: String,
    #[prost(string, tag = "3")]
    pub model_type: String,
    #[prost(string, tag = "4")]
    pub version: String,
    #[prost(enumeration = "DeploymentStatusProto", tag = "5")]
    pub status: i32,
    #[prost(int64, tag = "6")]
    pub deployed_at: i64,
    #[prost(int64, tag = "7")]
    pub last_updated: i64,
    #[prost(message, repeated, tag = "8")]
    pub deployment_history: Vec<DeploymentEventProto>,
}

#[derive(Clone, PartialEq, Message)]
pub struct DeploymentEventProto {
    #[prost(string, tag = "1")]
    pub event_id: String,
    #[prost(enumeration = "DeploymentEventTypeProto", tag = "2")]
    pub event_type: i32,
    #[prost(int64, tag = "3")]
    pub timestamp: i64,
    #[prost(string, tag = "4")]
    pub version: String,
    #[prost(string, tag = "5")]
    pub details: String,
    #[prost(string, optional, tag = "6")]
    pub user_id: Option<String>,
}

#[derive(Clone, PartialEq, Message)]
pub struct ABTestConfigProto {
    #[prost(string, tag = "1")]
    pub name: String,
    #[prost(string, optional, tag = "2")]
    pub description: Option<String>,
    #[prost(double, tag = "3")]
    pub control_traffic_percentage: f64,
    #[prost(double, tag = "4")]
    pub treatment_traffic_percentage: f64,
    #[prost(int64, tag = "5")]
    pub duration_seconds: i64,
    #[prost(string, repeated, tag = "6")]
    pub success_criteria: Vec<String>,
    #[prost(bool, tag = "7")]
    pub auto_promote: bool,
    #[prost(bool, tag = "8")]
    pub auto_rollback: bool,
}

#[derive(Clone, PartialEq, Message)]
pub struct MonitoringConfigProto {
    #[prost(bool, tag = "1")]
    pub enable_metrics_collection: bool,
    #[prost(int64, tag = "2")]
    pub metrics_interval_seconds: i64,
    #[prost(bool, tag = "3")]
    pub enable_alerting: bool,
    #[prost(message, repeated, tag = "4")]
    pub sla_thresholds: Vec<SlaThresholdProto>,
}

#[derive(Clone, PartialEq, Message)]
pub struct SlaThresholdProto {
    #[prost(string, tag = "1")]
    pub metric_name: String,
    #[prost(double, tag = "2")]
    pub threshold_value: f64,
    #[prost(enumeration = "ThresholdTypeProto", tag = "3")]
    pub threshold_type: i32,
}

#[derive(Clone, PartialEq, Message)]
pub struct ModelMetricsProto {
    #[prost(double, tag = "1")]
    pub avg_latency_ms: f64,
    #[prost(double, tag = "2")]
    pub p95_latency_ms: f64,
    #[prost(double, tag = "3")]
    pub p99_latency_ms: f64,
    #[prost(double, tag = "4")]
    pub error_rate: f64,
    #[prost(int64, tag = "5")]
    pub total_requests: i64,
    #[prost(double, tag = "6")]
    pub cpu_usage_percent: f64,
    #[prost(double, tag = "7")]
    pub memory_usage_mb: f64,
}

#[derive(Clone, PartialEq, Message)]
pub struct ABTestStatusProto {
    #[prost(string, tag = "1")]
    pub experiment_id: String,
    #[prost(enumeration = "ABTestStatusEnum", tag = "2")]
    pub status: i32,
    #[prost(double, tag = "3")]
    pub progress_percentage: f64,
    #[prost(message, optional, tag = "4")]
    pub control_metrics: Option<ModelMetricsProto>,
    #[prost(message, optional, tag = "5")]
    pub treatment_metrics: Option<ModelMetricsProto>,
    #[prost(bool, tag = "6")]
    pub is_significant: bool,
    #[prost(string, optional, tag = "7")]
    pub recommendation: Option<String>,
}

// Enums
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(i32)]
pub enum DeploymentStrategyProto {
    Immediate = 0,
    BlueGreen = 1,
    Canary = 2,
    AbTest = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(i32)]
pub enum DeploymentStatusProto {
    Pending = 0,
    Deploying = 1,
    Active = 2,
    Failed = 3,
    Archived = 4,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(i32)]
pub enum DeploymentEventTypeProto {
    Deployed = 0,
    Updated = 1,
    HotSwapped = 2,
    AbTestStarted = 3,
    AbTestCompleted = 4,
    RolledBack = 5,
    ValidationFailed = 6,
    PerformanceDegraded = 7,
    Archived = 8,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(i32)]
pub enum ThresholdTypeProto {
    LessThan = 0,
    GreaterThan = 1,
    Equals = 2,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(i32)]
pub enum ABTestStatusEnum {
    Running = 0,
    Completed = 1,
    Failed = 2,
    Cancelled = 3,
}

/// gRPC service trait for model deployment management
#[tonic::async_trait]
pub trait ModelDeploymentService {
    /// Deploy a new model or update an existing one
    async fn deploy_model(
        &self,
        request: Request<DeployModelRequest>,
    ) -> Result<Response<DeployModelResponse>, Status>;

    /// Get deployment information for a specific model
    async fn get_deployment(
        &self,
        request: Request<GetDeploymentRequest>,
    ) -> Result<Response<GetDeploymentResponse>, Status>;

    /// List all deployments with optional filtering
    async fn list_deployments(
        &self,
        request: Request<ListDeploymentsRequest>,
    ) -> Result<Response<ListDeploymentsResponse>, Status>;

    /// Rollback a model to a previous version
    async fn rollback_model(
        &self,
        request: Request<RollbackModelRequest>,
    ) -> Result<Response<RollbackModelResponse>, Status>;

    /// Undeploy a model completely
    async fn undeploy_model(
        &self,
        request: Request<UndeployModelRequest>,
    ) -> Result<Response<UndeployModelResponse>, Status>;

    /// Get performance metrics for a model
    async fn get_model_metrics(
        &self,
        request: Request<GetModelMetricsRequest>,
    ) -> Result<Response<GetModelMetricsResponse>, Status>;

    /// Start an A/B test experiment
    async fn start_ab_test(
        &self,
        request: Request<StartABTestRequest>,
    ) -> Result<Response<StartABTestResponse>, Status>;

    /// Get A/B test status and results
    async fn get_ab_test_status(
        &self,
        request: Request<GetABTestStatusRequest>,
    ) -> Result<Response<GetABTestStatusResponse>, Status>;
}

/// Implementation of the gRPC model deployment service
pub struct ModelDeploymentServiceImpl {
    registry: Arc<ModelDeploymentRegistry>,
    model_factory: Arc<dyn ModelFactory>,
}

/// Factory trait for creating models from serialized data
#[tonic::async_trait]
pub trait ModelFactory: Send + Sync {
    async fn create_model(
        &self,
        model_type: &str,
        model_data: &[u8],
        version: &ModelVersion,
    ) -> MLResult<Arc<dyn MLModel>>;
}

impl ModelDeploymentServiceImpl {
    /// Create a new deployment service
    pub fn new(
        registry: Arc<ModelDeploymentRegistry>,
        model_factory: Arc<dyn ModelFactory>,
    ) -> Self {
        Self {
            registry,
            model_factory,
        }
    }
}

#[tonic::async_trait]
impl ModelDeploymentService for ModelDeploymentServiceImpl {
    async fn deploy_model(
        &self,
        request: Request<DeployModelRequest>,
    ) -> Result<Response<DeployModelResponse>, Status> {
        let req = request.into_inner();
        
        // Parse version
        let version = ModelVersion::parse(&req.version)
            .map_err(|e| Status::new(Code::InvalidArgument, format!("Invalid version: {}", e)))?;

        // Create model from factory
        let model = self.model_factory
            .create_model(&req.model_type, &req.model_data, &version)
            .await
            .map_err(|e| Status::new(Code::Internal, format!("Failed to create model: {}", e)))?;

        // Convert config
        let config = req.config
            .map(|c| convert_deployment_config(c))
            .unwrap_or_default();

        // Deploy model
        match self.registry.deploy_model(req.model_id, model, version, config).await {
            Ok(deployment_id) => {
                let response = DeployModelResponse {
                    deployment_id: deployment_id.to_string(),
                    success: true,
                    message: "Model deployed successfully".to_string(),
                };
                Ok(Response::new(response))
            },
            Err(e) => {
                let response = DeployModelResponse {
                    deployment_id: String::new(),
                    success: false,
                    message: format!("Deployment failed: {}", e),
                };
                Ok(Response::new(response))
            }
        }
    }

    async fn get_deployment(
        &self,
        request: Request<GetDeploymentRequest>,
    ) -> Result<Response<GetDeploymentResponse>, Status> {
        let req = request.into_inner();
        
        match self.registry.get_deployment(&req.model_id).await {
            Ok(Some(entry)) => {
                let deployment_info = convert_to_deployment_info_proto(&entry);
                let response = GetDeploymentResponse {
                    deployment: Some(deployment_info),
                    found: true,
                };
                Ok(Response::new(response))
            },
            Ok(None) => {
                let response = GetDeploymentResponse {
                    deployment: None,
                    found: false,
                };
                Ok(Response::new(response))
            },
            Err(e) => Err(Status::new(Code::Internal, format!("Failed to get deployment: {}", e))),
        }
    }

    async fn list_deployments(
        &self,
        request: Request<ListDeploymentsRequest>,
    ) -> Result<Response<ListDeploymentsResponse>, Status> {
        let req = request.into_inner();
        
        match self.registry.list_deployments().await {
            Ok(entries) => {
                let mut deployments: Vec<_> = entries
                    .into_iter()
                    .map(|entry| convert_to_deployment_info_proto(&entry))
                    .collect();

                // Apply filtering if provided
                if let Some(filter) = req.filter {
                    deployments.retain(|d| d.model_id.contains(&filter) || d.model_type.contains(&filter));
                }

                let total_count = deployments.len() as i32;

                // Apply pagination
                let start = req.offset as usize;
                let end = if req.limit > 0 {
                    std::cmp::min(start + req.limit as usize, deployments.len())
                } else {
                    deployments.len()
                };

                if start < deployments.len() {
                    deployments = deployments[start..end].to_vec();
                } else {
                    deployments.clear();
                }

                let response = ListDeploymentsResponse {
                    deployments,
                    total_count,
                };
                Ok(Response::new(response))
            },
            Err(e) => Err(Status::new(Code::Internal, format!("Failed to list deployments: {}", e))),
        }
    }

    async fn rollback_model(
        &self,
        request: Request<RollbackModelRequest>,
    ) -> Result<Response<RollbackModelResponse>, Status> {
        let req = request.into_inner();
        
        match self.registry.rollback_model(&req.model_id).await {
            Ok(()) => {
                // Get the new version after rollback
                let new_version = if let Ok(Some(entry)) = self.registry.get_deployment(&req.model_id).await {
                    entry.metadata.version.to_string()
                } else {
                    "unknown".to_string()
                };

                let response = RollbackModelResponse {
                    success: true,
                    message: "Model rolled back successfully".to_string(),
                    new_version,
                };
                Ok(Response::new(response))
            },
            Err(e) => {
                let response = RollbackModelResponse {
                    success: false,
                    message: format!("Rollback failed: {}", e),
                    new_version: String::new(),
                };
                Ok(Response::new(response))
            }
        }
    }

    async fn undeploy_model(
        &self,
        request: Request<UndeployModelRequest>,
    ) -> Result<Response<UndeployModelResponse>, Status> {
        let req = request.into_inner();
        
        match self.registry.undeploy_model(&req.model_id).await {
            Ok(()) => {
                let response = UndeployModelResponse {
                    success: true,
                    message: "Model undeployed successfully".to_string(),
                };
                Ok(Response::new(response))
            },
            Err(e) => {
                let response = UndeployModelResponse {
                    success: false,
                    message: format!("Undeploy failed: {}", e),
                };
                Ok(Response::new(response))
            }
        }
    }

    async fn get_model_metrics(
        &self,
        request: Request<GetModelMetricsRequest>,
    ) -> Result<Response<GetModelMetricsResponse>, Status> {
        let req = request.into_inner();
        
        match self.registry.get_deployment(&req.model_id).await {
            Ok(Some(entry)) => {
                // Get metrics from the performance monitor
                match entry.monitor.get_current_metrics().await {
                    Ok(metrics) => {
                        let metrics_proto = ModelMetricsProto {
                            avg_latency_ms: metrics.avg_latency.as_secs_f64() * 1000.0,
                            p95_latency_ms: metrics.p95_latency.as_secs_f64() * 1000.0,
                            p99_latency_ms: metrics.p99_latency.as_secs_f64() * 1000.0,
                            error_rate: metrics.error_rate,
                            total_requests: metrics.total_requests as i64,
                            cpu_usage_percent: metrics.cpu_usage * 100.0,
                            memory_usage_mb: metrics.memory_usage / 1024.0 / 1024.0,
                        };

                        let response = GetModelMetricsResponse {
                            metrics: Some(metrics_proto),
                        };
                        Ok(Response::new(response))
                    },
                    Err(e) => Err(Status::new(Code::Internal, format!("Failed to get metrics: {}", e))),
                }
            },
            Ok(None) => Err(Status::new(Code::NotFound, "Model not found")),
            Err(e) => Err(Status::new(Code::Internal, format!("Failed to get deployment: {}", e))),
        }
    }

    async fn start_ab_test(
        &self,
        request: Request<StartABTestRequest>,
    ) -> Result<Response<StartABTestResponse>, Status> {
        let req = request.into_inner();
        
        // Parse treatment version
        let version = ModelVersion::parse(&req.treatment_version)
            .map_err(|e| Status::new(Code::InvalidArgument, format!("Invalid version: {}", e)))?;

        // Create treatment model
        let treatment_model = self.model_factory
            .create_model(&req.treatment_model_type, &req.treatment_model_data, &version)
            .await
            .map_err(|e| Status::new(Code::Internal, format!("Failed to create treatment model: {}", e)))?;

        // Convert A/B test config
        let ab_config = req.ab_test_config
            .map(|c| convert_ab_test_config(c))
            .unwrap_or_default();

        // Create deployment config for A/B test
        let deployment_config = DeploymentConfig {
            enable_ab_testing: true,
            ab_test_config: Some(ab_config.clone()),
            validation_required: true,
            monitoring_config: None,
            auto_rollback: true,
            deployment_strategy: DeploymentStrategy::ABTest { config: ab_config },
        };

        // Start A/B test deployment
        match self.registry.deploy_model(req.model_id, treatment_model, version, deployment_config).await {
            Ok(deployment_id) => {
                let response = StartABTestResponse {
                    experiment_id: deployment_id.to_string(),
                    success: true,
                    message: "A/B test started successfully".to_string(),
                };
                Ok(Response::new(response))
            },
            Err(e) => {
                let response = StartABTestResponse {
                    experiment_id: String::new(),
                    success: false,
                    message: format!("Failed to start A/B test: {}", e),
                };
                Ok(Response::new(response))
            }
        }
    }

    async fn get_ab_test_status(
        &self,
        request: Request<GetABTestStatusRequest>,
    ) -> Result<Response<GetABTestStatusResponse>, Status> {
        let _req = request.into_inner();
        
        // For now, return a placeholder response
        // In a full implementation, this would query the A/B test manager
        let status = ABTestStatusProto {
            experiment_id: "placeholder".to_string(),
            status: ABTestStatusEnum::Running as i32,
            progress_percentage: 50.0,
            control_metrics: None,
            treatment_metrics: None,
            is_significant: false,
            recommendation: None,
        };

        let response = GetABTestStatusResponse {
            status: Some(status),
        };
        Ok(Response::new(response))
    }
}

// Helper functions for converting between internal types and protobuf types

fn convert_deployment_config(proto: DeploymentConfigProto) -> DeploymentConfig {
    let deployment_strategy = match proto.deployment_strategy {
        0 => DeploymentStrategy::Immediate,
        1 => DeploymentStrategy::BlueGreen,
        2 => DeploymentStrategy::Canary { percentage: 10.0 }, // Default 10%
        3 => DeploymentStrategy::ABTest { 
            config: proto.ab_test_config
                .map(convert_ab_test_config)
                .unwrap_or_default()
        },
        _ => DeploymentStrategy::Immediate,
    };

    DeploymentConfig {
        enable_ab_testing: proto.enable_ab_testing,
        ab_test_config: proto.ab_test_config.map(convert_ab_test_config),
        validation_required: proto.validation_required,
        monitoring_config: proto.monitoring_config.map(convert_monitoring_config),
        auto_rollback: proto.auto_rollback,
        deployment_strategy,
    }
}

fn convert_ab_test_config(proto: ABTestConfigProto) -> ABTestConfig {
    ABTestConfig {
        name: proto.name,
        description: proto.description,
        control_traffic_percentage: proto.control_traffic_percentage,
        treatment_traffic_percentage: proto.treatment_traffic_percentage,
        duration: std::time::Duration::from_secs(proto.duration_seconds as u64),
        success_criteria: proto.success_criteria,
        auto_promote: proto.auto_promote,
        auto_rollback: proto.auto_rollback,
    }
}

fn convert_monitoring_config(proto: MonitoringConfigProto) -> MonitoringConfig {
    MonitoringConfig {
        enable_metrics_collection: proto.enable_metrics_collection,
        metrics_interval: std::time::Duration::from_secs(proto.metrics_interval_seconds as u64),
        enable_alerting: proto.enable_alerting,
        sla_thresholds: proto.sla_thresholds
            .into_iter()
            .map(|threshold| {
                use super::monitoring::{SlaThreshold, ThresholdType};
                SlaThreshold {
                    metric_name: threshold.metric_name,
                    threshold_value: threshold.threshold_value,
                    threshold_type: match threshold.threshold_type {
                        0 => ThresholdType::LessThan,
                        1 => ThresholdType::GreaterThan,
                        2 => ThresholdType::Equals,
                        _ => ThresholdType::LessThan,
                    },
                }
            })
            .collect(),
        alert_channels: vec![], // Would be populated from additional proto fields
    }
}

fn convert_to_deployment_info_proto(entry: &super::RegistryEntry) -> DeploymentInfoProto {
    let deployed_at = entry.metadata.deployed_at
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64;
    
    let last_updated = entry.last_updated
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64;

    let deployment_history = entry.deployment_history
        .iter()
        .map(|event| {
            let timestamp = event.timestamp
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs() as i64;

            DeploymentEventProto {
                event_id: event.event_id.to_string(),
                event_type: match event.event_type {
                    super::DeploymentEventType::Deployed => DeploymentEventTypeProto::Deployed as i32,
                    super::DeploymentEventType::Updated => DeploymentEventTypeProto::Updated as i32,
                    super::DeploymentEventType::HotSwapped => DeploymentEventTypeProto::HotSwapped as i32,
                    super::DeploymentEventType::ABTestStarted => DeploymentEventTypeProto::AbTestStarted as i32,
                    super::DeploymentEventType::ABTestCompleted => DeploymentEventTypeProto::AbTestCompleted as i32,
                    super::DeploymentEventType::RolledBack => DeploymentEventTypeProto::RolledBack as i32,
                    super::DeploymentEventType::ValidationFailed => DeploymentEventTypeProto::ValidationFailed as i32,
                    super::DeploymentEventType::PerformanceDegraded => DeploymentEventTypeProto::PerformanceDegraded as i32,
                    super::DeploymentEventType::Archived => DeploymentEventTypeProto::Archived as i32,
                },
                timestamp,
                version: event.version.to_string(),
                details: event.details.clone(),
                user_id: event.user_id.clone(),
            }
        })
        .collect();

    DeploymentInfoProto {
        deployment_id: entry.deployment_id.to_string(),
        model_id: entry.model_id.clone(),
        model_type: format!("{:?}", entry.metadata.model_type),
        version: entry.metadata.version.to_string(),
        status: match entry.metadata.status {
            super::DeploymentStatus::Pending => DeploymentStatusProto::Pending as i32,
            super::DeploymentStatus::Deploying => DeploymentStatusProto::Deploying as i32,
            super::DeploymentStatus::Active => DeploymentStatusProto::Active as i32,
            super::DeploymentStatus::Failed => DeploymentStatusProto::Failed as i32,
            super::DeploymentStatus::Archived => DeploymentStatusProto::Archived as i32,
        },
        deployed_at,
        last_updated,
        deployment_history,
    }
}

impl Default for DeploymentConfig {
    fn default() -> Self {
        Self {
            enable_ab_testing: false,
            ab_test_config: None,
            validation_required: true,
            monitoring_config: None,
            auto_rollback: true,
            deployment_strategy: DeploymentStrategy::Immediate,
        }
    }
}

impl Default for ABTestConfig {
    fn default() -> Self {
        Self {
            name: "default_ab_test".to_string(),
            description: None,
            control_traffic_percentage: 50.0,
            treatment_traffic_percentage: 50.0,
            duration: std::time::Duration::from_secs(3600), // 1 hour
            success_criteria: vec![],
            auto_promote: false,
            auto_rollback: true,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use crate::models::MockMLModel;

    struct MockModelFactory;

    #[tonic::async_trait]
    impl ModelFactory for MockModelFactory {
        async fn create_model(
            &self,
            _model_type: &str,
            _model_data: &[u8],
            _version: &ModelVersion,
        ) -> MLResult<Arc<dyn MLModel>> {
            Ok(Arc::new(MockMLModel::new()))
        }
    }

    #[tokio::test]
    async fn test_deploy_model_endpoint() {
        let registry = Arc::new(
            ModelDeploymentRegistry::new(super::super::RegistryConfig::default())
                .await
                .unwrap()
        );
        let factory = Arc::new(MockModelFactory);
        let service = ModelDeploymentServiceImpl::new(registry, factory);

        let request = Request::new(DeployModelRequest {
            model_id: "test_model".to_string(),
            model_type: "mock".to_string(),
            model_data: vec![1, 2, 3, 4],
            version: "1.0.0".to_string(),
            config: Some(DeploymentConfigProto {
                enable_ab_testing: false,
                ab_test_config: None,
                validation_required: false,
                monitoring_config: None,
                auto_rollback: true,
                deployment_strategy: DeploymentStrategyProto::Immediate as i32,
            }),
        });

        let response = service.deploy_model(request).await.unwrap();
        let response = response.into_inner();

        assert!(response.success);
        assert!(!response.deployment_id.is_empty());
    }

    #[tokio::test]
    async fn test_get_deployment_endpoint() {
        let registry = Arc::new(
            ModelDeploymentRegistry::new(super::super::RegistryConfig::default())
                .await
                .unwrap()
        );
        let factory = Arc::new(MockModelFactory);
        let service = ModelDeploymentServiceImpl::new(registry.clone(), factory);

        // First deploy a model
        let model = Arc::new(MockMLModel::new());
        let version = ModelVersion::new(1, 0, 0);
        registry.deploy_model(
            "test_model".to_string(),
            model,
            version.clone(),
            DeploymentConfig::default(),
        ).await.unwrap();

        // Then get deployment info
        let request = Request::new(GetDeploymentRequest {
            model_id: "test_model".to_string(),
        });

        let response = service.get_deployment(request).await.unwrap();
        let response = response.into_inner();

        assert!(response.found);
        assert!(response.deployment.is_some());
        
        let deployment = response.deployment.unwrap();
        assert_eq!(deployment.model_id, "test_model");
        assert_eq!(deployment.version, "1.0.0");
    }
}