use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelCategory {
    ImageClassification,
    ObjectDetection,
    Segmentation,
    TextGeneration,
    SpeechRecognition,
    NLP,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelInfo {
    pub id: String,
    pub name: String,
    pub description: String,
    pub category: ModelCategory,
    pub size_mb: f64,
    pub parameters: u64,
    pub accuracy_metric: String,
    pub download_url: String,
    pub license: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelStatus {
    NotDownloaded,
    Downloading,
    Downloaded,
    Loading,
    Ready,
    Error,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelZooConfig {
    pub selected_model: Option<String>,
    pub download_path: String,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct InferenceRequest {
    pub model_id: String,
    pub input_data: String,
    pub input_type: String,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct InferenceResult {
    pub model_id: String,
    pub output: String,
    pub inference_time_ms: f64,
    pub device: String,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ModelZooRequest {
    ListModels,
    FilterByCategory { category: ModelCategory },
    GetModelInfo { model_id: String },
    GetModelStatus { model_id: String },
    DownloadModel { model_id: String },
    LoadModel { model_id: String },
    UnloadModel { model_id: String },
    RunInference { request: InferenceRequest },
    GetConfig,
    UpdateConfig { config: ModelZooConfig },
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ModelZooResponse {
    ModelList { models: Vec<ModelInfo> },
    ModelInfo { info: ModelInfo },
    ModelStatus { status: ModelStatus },
    DownloadStarted { model_id: String },
    ModelLoaded { model_id: String },
    ModelUnloaded { model_id: String },
    InferenceResult { result: InferenceResult },
    Config { config: ModelZooConfig },
    ConfigUpdated,
    Error { message: String },
}

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

    #[test]
    fn test_model_category_serialization() {
        let category = ModelCategory::ImageClassification;
        let json = serde_json::to_string(&category).expect("Failed to serialize");
        let deserialized: ModelCategory =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(category, deserialized);
    }

    #[test]
    fn test_model_category_equality() {
        assert_eq!(
            ModelCategory::ImageClassification,
            ModelCategory::ImageClassification
        );
        assert_ne!(
            ModelCategory::ImageClassification,
            ModelCategory::ObjectDetection
        );
    }

    #[test]
    fn test_model_category_debug() {
        let category = ModelCategory::ImageClassification;
        let debug_str = format!("{category:?}");
        assert_eq!(debug_str, "ImageClassification");
    }

    #[test]
    fn test_model_category_all_variants() {
        let categories = vec![
            ModelCategory::ImageClassification,
            ModelCategory::ObjectDetection,
            ModelCategory::Segmentation,
            ModelCategory::TextGeneration,
            ModelCategory::SpeechRecognition,
            ModelCategory::NLP,
        ];
        assert_eq!(categories.len(), 6);
    }

    #[test]
    fn test_model_info_creation() {
        let info = ModelInfo {
            id: "resnet50".to_string(),
            name: "ResNet-50".to_string(),
            description: "50-layer residual network".to_string(),
            category: ModelCategory::ImageClassification,
            size_mb: 97.8,
            parameters: 25_600_000,
            accuracy_metric: "Top-1: 76.1%".to_string(),
            download_url: "https://example.com/resnet50.pth".to_string(),
            license: "MIT".to_string(),
        };
        assert_eq!(info.id, "resnet50");
        assert_eq!(info.name, "ResNet-50");
        assert_eq!(info.category, ModelCategory::ImageClassification);
        assert_eq!(info.size_mb, 97.8);
        assert_eq!(info.parameters, 25_600_000);
    }

    #[test]
    fn test_model_info_serialization() {
        let info = ModelInfo {
            id: "resnet50".to_string(),
            name: "ResNet-50".to_string(),
            description: "50-layer residual network".to_string(),
            category: ModelCategory::ImageClassification,
            size_mb: 97.8,
            parameters: 25_600_000,
            accuracy_metric: "Top-1: 76.1%".to_string(),
            download_url: "https://example.com/resnet50.pth".to_string(),
            license: "MIT".to_string(),
        };
        let json = serde_json::to_string(&info).expect("Failed to serialize");
        let deserialized: ModelInfo = serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(info, deserialized);
    }

    #[test]
    fn test_model_info_clone() {
        let info = ModelInfo {
            id: "resnet50".to_string(),
            name: "ResNet-50".to_string(),
            description: "50-layer residual network".to_string(),
            category: ModelCategory::ImageClassification,
            size_mb: 97.8,
            parameters: 25_600_000,
            accuracy_metric: "Top-1: 76.1%".to_string(),
            download_url: "https://example.com/resnet50.pth".to_string(),
            license: "MIT".to_string(),
        };
        let cloned = info.clone();
        assert_eq!(info, cloned);
    }

    #[test]
    fn test_model_status_all_variants() {
        let statuses = vec![
            ModelStatus::NotDownloaded,
            ModelStatus::Downloading,
            ModelStatus::Downloaded,
            ModelStatus::Loading,
            ModelStatus::Ready,
            ModelStatus::Error,
        ];
        assert_eq!(statuses.len(), 6);
    }

    #[test]
    fn test_model_status_equality() {
        assert_eq!(ModelStatus::Ready, ModelStatus::Ready);
        assert_ne!(ModelStatus::Ready, ModelStatus::Loading);
    }

    #[test]
    fn test_model_status_serialization() {
        let status = ModelStatus::Ready;
        let json = serde_json::to_string(&status).expect("Failed to serialize");
        let deserialized: ModelStatus = serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(status, deserialized);
    }

    #[test]
    fn test_model_status_debug() {
        let status = ModelStatus::Ready;
        let debug_str = format!("{status:?}");
        assert_eq!(debug_str, "Ready");
    }

    #[test]
    fn test_model_zoo_config_creation() {
        let config = ModelZooConfig {
            selected_model: Some("resnet50".to_string()),
            download_path: "/tmp/models".to_string(),
        };
        assert_eq!(config.selected_model, Some("resnet50".to_string()));
        assert_eq!(config.download_path, "/tmp/models");
    }

    #[test]
    fn test_model_zoo_config_no_selection() {
        let config = ModelZooConfig {
            selected_model: None,
            download_path: "/tmp/models".to_string(),
        };
        assert_eq!(config.selected_model, None);
    }

    #[test]
    fn test_model_zoo_config_serialization() {
        let config = ModelZooConfig {
            selected_model: Some("resnet50".to_string()),
            download_path: "/tmp/models".to_string(),
        };
        let json = serde_json::to_string(&config).expect("Failed to serialize");
        let deserialized: ModelZooConfig =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(config, deserialized);
    }

    #[test]
    fn test_inference_request_creation() {
        let request = InferenceRequest {
            model_id: "resnet50".to_string(),
            input_data: "base64encodedimage".to_string(),
            input_type: "image/jpeg".to_string(),
        };
        assert_eq!(request.model_id, "resnet50");
        assert_eq!(request.input_data, "base64encodedimage");
        assert_eq!(request.input_type, "image/jpeg");
    }

    #[test]
    fn test_inference_request_serialization() {
        let request = InferenceRequest {
            model_id: "resnet50".to_string(),
            input_data: "base64encodedimage".to_string(),
            input_type: "image/jpeg".to_string(),
        };
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: InferenceRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_inference_result_creation() {
        let result = InferenceResult {
            model_id: "resnet50".to_string(),
            output: "class: cat, confidence: 0.95".to_string(),
            inference_time_ms: 15.5,
            device: "CPU".to_string(),
        };
        assert_eq!(result.model_id, "resnet50");
        assert_eq!(result.output, "class: cat, confidence: 0.95");
        assert_eq!(result.inference_time_ms, 15.5);
        assert_eq!(result.device, "CPU");
    }

    #[test]
    fn test_inference_result_serialization() {
        let result = InferenceResult {
            model_id: "resnet50".to_string(),
            output: "class: cat, confidence: 0.95".to_string(),
            inference_time_ms: 15.5,
            device: "CPU".to_string(),
        };
        let json = serde_json::to_string(&result).expect("Failed to serialize");
        let deserialized: InferenceResult =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(result, deserialized);
    }

    #[test]
    fn test_model_zoo_request_list_models() {
        let request = ModelZooRequest::ListModels;
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: ModelZooRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_model_zoo_request_filter_by_category() {
        let request = ModelZooRequest::FilterByCategory {
            category: ModelCategory::ImageClassification,
        };
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: ModelZooRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_model_zoo_request_get_model_info() {
        let request = ModelZooRequest::GetModelInfo {
            model_id: "resnet50".to_string(),
        };
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: ModelZooRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_model_zoo_request_download_model() {
        let request = ModelZooRequest::DownloadModel {
            model_id: "resnet50".to_string(),
        };
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: ModelZooRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_model_zoo_request_run_inference() {
        let inference_request = InferenceRequest {
            model_id: "resnet50".to_string(),
            input_data: "base64encodedimage".to_string(),
            input_type: "image/jpeg".to_string(),
        };
        let request = ModelZooRequest::RunInference {
            request: inference_request,
        };
        let json = serde_json::to_string(&request).expect("Failed to serialize");
        let deserialized: ModelZooRequest =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_model_zoo_response_model_list() {
        let model_info = ModelInfo {
            id: "resnet50".to_string(),
            name: "ResNet-50".to_string(),
            description: "50-layer residual network".to_string(),
            category: ModelCategory::ImageClassification,
            size_mb: 97.8,
            parameters: 25_600_000,
            accuracy_metric: "Top-1: 76.1%".to_string(),
            download_url: "https://example.com/resnet50.pth".to_string(),
            license: "MIT".to_string(),
        };
        let response = ModelZooResponse::ModelList {
            models: vec![model_info],
        };
        let json = serde_json::to_string(&response).expect("Failed to serialize");
        let deserialized: ModelZooResponse =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(response, deserialized);
    }

    #[test]
    fn test_model_zoo_response_model_status() {
        let response = ModelZooResponse::ModelStatus {
            status: ModelStatus::Ready,
        };
        let json = serde_json::to_string(&response).expect("Failed to serialize");
        let deserialized: ModelZooResponse =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(response, deserialized);
    }

    #[test]
    fn test_model_zoo_response_inference_result() {
        let inference_result = InferenceResult {
            model_id: "resnet50".to_string(),
            output: "class: cat, confidence: 0.95".to_string(),
            inference_time_ms: 15.5,
            device: "CPU".to_string(),
        };
        let response = ModelZooResponse::InferenceResult {
            result: inference_result,
        };
        let json = serde_json::to_string(&response).expect("Failed to serialize");
        let deserialized: ModelZooResponse =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(response, deserialized);
    }

    #[test]
    fn test_model_zoo_response_error() {
        let response = ModelZooResponse::Error {
            message: "Model not found".to_string(),
        };
        let json = serde_json::to_string(&response).expect("Failed to serialize");
        let deserialized: ModelZooResponse =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(response, deserialized);
    }

    #[test]
    fn test_all_request_variants() {
        let requests = vec![
            ModelZooRequest::ListModels,
            ModelZooRequest::FilterByCategory {
                category: ModelCategory::ImageClassification,
            },
            ModelZooRequest::GetModelInfo {
                model_id: "test".to_string(),
            },
            ModelZooRequest::GetModelStatus {
                model_id: "test".to_string(),
            },
            ModelZooRequest::DownloadModel {
                model_id: "test".to_string(),
            },
            ModelZooRequest::LoadModel {
                model_id: "test".to_string(),
            },
            ModelZooRequest::UnloadModel {
                model_id: "test".to_string(),
            },
            ModelZooRequest::RunInference {
                request: InferenceRequest {
                    model_id: "test".to_string(),
                    input_data: "data".to_string(),
                    input_type: "type".to_string(),
                },
            },
            ModelZooRequest::GetConfig,
            ModelZooRequest::UpdateConfig {
                config: ModelZooConfig {
                    selected_model: None,
                    download_path: "/tmp".to_string(),
                },
            },
        ];
        assert_eq!(requests.len(), 10);
    }

    #[test]
    fn test_all_response_variants() {
        let model_info = ModelInfo {
            id: "test".to_string(),
            name: "Test".to_string(),
            description: "Test model".to_string(),
            category: ModelCategory::ImageClassification,
            size_mb: 1.0,
            parameters: 1000,
            accuracy_metric: "N/A".to_string(),
            download_url: "https://example.com".to_string(),
            license: "MIT".to_string(),
        };
        let responses = vec![
            ModelZooResponse::ModelList {
                models: vec![model_info.clone()],
            },
            ModelZooResponse::ModelInfo { info: model_info },
            ModelZooResponse::ModelStatus {
                status: ModelStatus::Ready,
            },
            ModelZooResponse::DownloadStarted {
                model_id: "test".to_string(),
            },
            ModelZooResponse::ModelLoaded {
                model_id: "test".to_string(),
            },
            ModelZooResponse::ModelUnloaded {
                model_id: "test".to_string(),
            },
            ModelZooResponse::InferenceResult {
                result: InferenceResult {
                    model_id: "test".to_string(),
                    output: "output".to_string(),
                    inference_time_ms: 1.0,
                    device: "CPU".to_string(),
                },
            },
            ModelZooResponse::Config {
                config: ModelZooConfig {
                    selected_model: None,
                    download_path: "/tmp".to_string(),
                },
            },
            ModelZooResponse::ConfigUpdated,
            ModelZooResponse::Error {
                message: "error".to_string(),
            },
        ];
        assert_eq!(responses.len(), 10);
    }
}