//! Configuration types for inference profiling.

use serde::{Deserialize, Serialize};

/// Model architecture type for profiling.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelType {
    /// ResNet-18 (11M parameters)
    #[serde(rename = "resnet18")]
    ResNet18,
    /// ResNet-50 (25M parameters)
    #[serde(rename = "resnet50")]
    ResNet50,
    /// Vision Transformer Base/16 (86M parameters)
    #[serde(rename = "vit-b16")]
    ViTB16,
    /// Vision Transformer Large/16 (304M parameters)
    #[serde(rename = "vit-l16")]
    ViTL16,
    /// `ConvNeXt` Tiny (29M parameters)
    #[serde(rename = "convnext-tiny")]
    ConvNeXtTiny,
    /// `ConvNeXt` Base (89M parameters)
    #[serde(rename = "convnext-base")]
    ConvNeXtBase,
    /// Custom user-provided model
    #[serde(rename = "custom")]
    Custom,
}

/// Compute device type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum DeviceType {
    /// CPU execution
    CPU,
    /// CUDA GPU execution
    CUDA,
    /// Metal GPU execution (macOS)
    Metal,
}

/// Input shape for profiling (NCHW format).
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct InputShape {
    /// Number of channels
    pub channels: usize,
    /// Height
    pub height: usize,
    /// Width
    pub width: usize,
}

/// Profile configuration.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProfileConfig {
    /// Model to profile
    pub model_type: ModelType,
    /// Device to run on
    pub device: DeviceType,
    /// Batch sizes to test
    pub batch_sizes: Vec<usize>,
    /// Warmup iterations before measurement
    pub warmup_iterations: usize,
    /// Benchmark iterations for measurement
    pub benchmark_iterations: usize,
    /// Input shape (excluding batch dimension)
    pub input_shape: InputShape,
}

impl ProfileConfig {
    /// Creates a new profile configuration.
    ///
    /// # Errors
    /// Returns an error if validation fails.
    pub fn new(
        model_type: ModelType,
        device: DeviceType,
        batch_sizes: Vec<usize>,
        warmup_iterations: usize,
        benchmark_iterations: usize,
        input_shape: InputShape,
    ) -> Result<Self, String> {
        let config = Self {
            model_type,
            device,
            batch_sizes,
            warmup_iterations,
            benchmark_iterations,
            input_shape,
        };
        config.validate()?;
        Ok(config)
    }

    /// Validates the configuration.
    ///
    /// # Errors
    /// Returns an error if any validation constraint fails.
    pub fn validate(&self) -> Result<(), String> {
        if self.batch_sizes.is_empty() {
            return Err("batch_sizes cannot be empty".to_string());
        }

        for &batch_size in &self.batch_sizes {
            if batch_size == 0 {
                return Err("batch_size cannot be zero".to_string());
            }
            if batch_size > 1024 {
                return Err(format!("batch_size {batch_size} exceeds maximum of 1024"));
            }
        }

        if self.warmup_iterations > 10000 {
            return Err(format!(
                "warmup_iterations {} exceeds maximum of 10000",
                self.warmup_iterations
            ));
        }

        if self.benchmark_iterations == 0 {
            return Err("benchmark_iterations cannot be zero".to_string());
        }

        if self.benchmark_iterations > 100000 {
            return Err(format!(
                "benchmark_iterations {} exceeds maximum of 100000",
                self.benchmark_iterations
            ));
        }

        if self.input_shape.channels == 0 {
            return Err("input_shape.channels cannot be zero".to_string());
        }

        if self.input_shape.height == 0 {
            return Err("input_shape.height cannot be zero".to_string());
        }

        if self.input_shape.width == 0 {
            return Err("input_shape.width cannot be zero".to_string());
        }

        Ok(())
    }
}

impl Default for InputShape {
    fn default() -> Self {
        Self {
            channels: 3,
            height: 224,
            width: 224,
        }
    }
}

impl Default for ProfileConfig {
    fn default() -> Self {
        Self {
            model_type: ModelType::ResNet18,
            device: DeviceType::CPU,
            batch_sizes: vec![1, 2, 4, 8],
            warmup_iterations: 10,
            benchmark_iterations: 100,
            input_shape: InputShape::default(),
        }
    }
}

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

    #[test]
    fn test_model_type_serialization() {
        let model = ModelType::ResNet18;
        let json = serde_json::to_string(&model).expect("serialization failed");
        assert_eq!(json, "\"resnet18\"");
    }

    #[test]
    fn test_model_type_deserialization() {
        let json = "\"vit-b16\"";
        let model: ModelType = serde_json::from_str(json).expect("deserialization failed");
        assert_eq!(model, ModelType::ViTB16);
    }

    #[test]
    fn test_model_type_roundtrip() {
        let models = vec![
            ModelType::ResNet18,
            ModelType::ResNet50,
            ModelType::ViTB16,
            ModelType::ViTL16,
            ModelType::ConvNeXtTiny,
            ModelType::ConvNeXtBase,
            ModelType::Custom,
        ];

        for model in models {
            let json = serde_json::to_string(&model).expect("serialization failed");
            let decoded: ModelType = serde_json::from_str(&json).expect("deserialization failed");
            assert_eq!(model, decoded);
        }
    }

    #[test]
    fn test_device_type_serialization() {
        let device = DeviceType::CUDA;
        let json = serde_json::to_string(&device).expect("serialization failed");
        assert_eq!(json, "\"CUDA\"");
    }

    #[test]
    fn test_device_type_deserialization() {
        let json = "\"CPU\"";
        let device: DeviceType = serde_json::from_str(json).expect("deserialization failed");
        assert_eq!(device, DeviceType::CPU);
    }

    #[test]
    fn test_device_type_roundtrip() {
        let devices = vec![DeviceType::CPU, DeviceType::CUDA, DeviceType::Metal];

        for device in devices {
            let json = serde_json::to_string(&device).expect("serialization failed");
            let decoded: DeviceType = serde_json::from_str(&json).expect("deserialization failed");
            assert_eq!(device, decoded);
        }
    }

    #[test]
    fn test_profile_config_validation_success() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2, 4],
            10,
            100,
            InputShape {
                channels: 3,
                height: 224,
                width: 224,
            },
        );

        assert!(config.is_ok());
    }

    #[test]
    fn test_profile_config_validation_empty_batch_sizes() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![],
            10,
            100,
            InputShape::default(),
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "batch_sizes cannot be empty");
    }

    #[test]
    fn test_profile_config_validation_zero_batch_size() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 0, 4],
            10,
            100,
            InputShape::default(),
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "batch_size cannot be zero");
    }

    #[test]
    fn test_profile_config_validation_batch_size_too_large() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2048],
            10,
            100,
            InputShape::default(),
        );

        assert!(config.is_err());
        assert!(config.unwrap_err().contains("exceeds maximum of 1024"));
    }

    #[test]
    fn test_profile_config_validation_warmup_iterations_too_large() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            20000,
            100,
            InputShape::default(),
        );

        assert!(config.is_err());
        let err = config.unwrap_err();
        assert!(err.contains("warmup_iterations"));
        assert!(err.contains("exceeds maximum of 10000"));
    }

    #[test]
    fn test_profile_config_validation_zero_benchmark_iterations() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            10,
            0,
            InputShape::default(),
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "benchmark_iterations cannot be zero");
    }

    #[test]
    fn test_profile_config_validation_benchmark_iterations_too_large() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            10,
            200000,
            InputShape::default(),
        );

        assert!(config.is_err());
        let err = config.unwrap_err();
        assert!(err.contains("benchmark_iterations"));
        assert!(err.contains("exceeds maximum of 100000"));
    }

    #[test]
    fn test_profile_config_validation_zero_channels() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            10,
            100,
            InputShape {
                channels: 0,
                height: 224,
                width: 224,
            },
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "input_shape.channels cannot be zero");
    }

    #[test]
    fn test_profile_config_validation_zero_height() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            10,
            100,
            InputShape {
                channels: 3,
                height: 0,
                width: 224,
            },
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "input_shape.height cannot be zero");
    }

    #[test]
    fn test_profile_config_validation_zero_width() {
        let config = ProfileConfig::new(
            ModelType::ResNet18,
            DeviceType::CPU,
            vec![1, 2],
            10,
            100,
            InputShape {
                channels: 3,
                height: 224,
                width: 0,
            },
        );

        assert!(config.is_err());
        assert_eq!(config.unwrap_err(), "input_shape.width cannot be zero");
    }

    #[test]
    fn test_profile_config_serialization() {
        let config = ProfileConfig::default();
        let json = serde_json::to_string(&config).expect("serialization failed");
        assert!(json.contains("\"model_type\""));
        assert!(json.contains("\"resnet18\""));
    }

    #[test]
    fn test_profile_config_deserialization() {
        let json = r#"{
            "model_type": "vit-b16",
            "device": "CUDA",
            "batch_sizes": [1, 2, 4],
            "warmup_iterations": 10,
            "benchmark_iterations": 100,
            "input_shape": {
                "channels": 3,
                "height": 224,
                "width": 224
            }
        }"#;

        let config: ProfileConfig = serde_json::from_str(json).expect("deserialization failed");
        assert_eq!(config.model_type, ModelType::ViTB16);
        assert_eq!(config.device, DeviceType::CUDA);
        assert_eq!(config.batch_sizes, vec![1, 2, 4]);
    }

    #[test]
    fn test_input_shape_default() {
        let shape = InputShape::default();
        assert_eq!(shape.channels, 3);
        assert_eq!(shape.height, 224);
        assert_eq!(shape.width, 224);
    }

    #[test]
    fn test_profile_config_default() {
        let config = ProfileConfig::default();
        assert_eq!(config.model_type, ModelType::ResNet18);
        assert_eq!(config.device, DeviceType::CPU);
        assert_eq!(config.batch_sizes, vec![1, 2, 4, 8]);
        assert_eq!(config.warmup_iterations, 10);
        assert_eq!(config.benchmark_iterations, 100);
    }
}