// RTX Segmentation Demo - Mock segmentation backend
//
// This crate provides mock implementations for semantic segmentation inference,
// including realistic region-based mask generation and colormap visualization.

mod colormap;
mod error;
mod mock_segmentation;

pub use colormap::{apply_colormap_with_transparency, blend_with_image, mask_to_rgb};
pub use error::{Result, SegmentationError};
pub use mock_segmentation::{
    generate_for_model, generate_mock_segmentation, generate_with_distribution,
};

// Re-export shared types for convenience
pub use segmentation_shared::{
    ADE20K_CLASSES, CITYSCAPES_CLASSES, ClassInfo, PASCAL_VOC_CLASSES, SegmentRequest,
    SegmentResponse, SegmentationConfig, SegmentationMask, SegmentationModel, SegmentationResult,
};

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

    #[test]
    fn test_library_exports_error_types() {
        let _err: SegmentationError = SegmentationError::InvalidDimensions("test".to_string());
        let _result: Result<()> = Ok(());
    }

    #[test]
    fn test_library_exports_colormap_functions() {
        let mask = SegmentationMask::new(1, 1, vec![0]);
        let classes = &[ClassInfo::new(0, "bg", "#000000")];

        let _rgb = mask_to_rgb(&mask, classes).unwrap();
        let _rgba = apply_colormap_with_transparency(&mask, classes, true).unwrap();

        let original = vec![100, 100, 100];
        let _blended = blend_with_image(&original, &mask, classes, 0.5).unwrap();
    }

    #[test]
    fn test_library_exports_mock_generation() {
        let config =
            SegmentationConfig::new(SegmentationModel::DeepLabV3, "cpu".to_string(), 21, 512);
        let _result = generate_mock_segmentation(&config, 10, 10, Some(42)).unwrap();
        let _result2 = generate_for_model(SegmentationModel::UNet, 10, 10, Some(42)).unwrap();
        let _result3 = generate_with_distribution(10, 10, &[1.0, 1.0], Some(42)).unwrap();
    }

    #[test]
    fn test_library_exports_shared_types() {
        let _model = SegmentationModel::DeepLabV3;
        let _config = SegmentationConfig::new(_model, "cpu".to_string(), 21, 512);
        let _mask = SegmentationMask::new(1, 1, vec![0]);
        let _result = SegmentationResult::new(_mask.clone(), 100.0);
        let _class = ClassInfo::new(0, "bg", "#000000");
    }

    #[test]
    fn test_library_exports_ipc_types() {
        let _request = SegmentRequest::GetModels;
        let _response = SegmentResponse::error("test".to_string());
    }

    #[test]
    fn test_library_exports_class_sets() {
        assert!(!PASCAL_VOC_CLASSES.is_empty());
        assert!(!ADE20K_CLASSES.is_empty());
        assert!(!CITYSCAPES_CLASSES.is_empty());
    }

    #[test]
    fn test_end_to_end_workflow() {
        let config =
            SegmentationConfig::new(SegmentationModel::DeepLabV3, "cpu".to_string(), 21, 512);

        let result = generate_mock_segmentation(&config, 50, 50, Some(42)).unwrap();

        assert!(result.validate().is_ok());

        let rgb = mask_to_rgb(&result.mask, PASCAL_VOC_CLASSES).unwrap();
        assert_eq!(rgb.len(), 50 * 50 * 3);

        let rgba =
            apply_colormap_with_transparency(&result.mask, PASCAL_VOC_CLASSES, true).unwrap();
        assert_eq!(rgba.len(), 50 * 50 * 4);
    }

    #[test]
    fn test_end_to_end_with_blending() {
        let config = SegmentationConfig::new(SegmentationModel::UNet, "cpu".to_string(), 21, 512);

        let result = generate_mock_segmentation(&config, 20, 20, Some(99)).unwrap();

        let original_image = vec![128u8; 20 * 20 * 3];

        let blended =
            blend_with_image(&original_image, &result.mask, PASCAL_VOC_CLASSES, 0.6).unwrap();
        assert_eq!(blended.len(), original_image.len());

        // All values are valid u8 by type (no need to check bounds)
    }

    #[test]
    fn test_multiple_models_generate_different_class_counts() {
        let models = [
            (SegmentationModel::DeepLabV3, 21),
            (SegmentationModel::SegFormer, 150),
            (SegmentationModel::UNet, 21),
            (SegmentationModel::FCN, 21),
        ];

        for (model, expected_classes) in models {
            let result = generate_for_model(model, 30, 30, Some(42)).unwrap();
            assert_eq!(
                result.class_confidences.len(),
                expected_classes,
                "model {:?} should have {} classes",
                model,
                expected_classes
            );
        }
    }
}