//! Tests for the synthesis engine

// Shape guards tests - active tests
use crate::aot::shape_guards::{
    GuardCheckResult, GuardFailure, ShapeDim, ShapeGuard, ShapeGuardConfig, ShapeGuardManager,
    ShapeSignature,
};

#[test]
fn test_shape_dim_concrete() {
    let dim = ShapeDim::concrete(256);
    assert!(dim.matches(256));
    assert!(!dim.matches(128));
    assert!(dim.is_concrete());
}

#[test]
fn test_shape_dim_symbolic() {
    let dim = ShapeDim::symbolic("batch_size");
    assert!(dim.matches(1));
    assert!(dim.matches(256));
    assert!(dim.matches(10000));
    assert!(!dim.is_concrete());
    assert_eq!(dim.symbolic_name(), Some("batch_size"));
}

#[test]
fn test_shape_dim_bounded() {
    let dim = ShapeDim::bounded("seq_len", 1, 512);
    assert!(dim.matches(1));
    assert!(dim.matches(256));
    assert!(dim.matches(512));
    assert!(!dim.matches(0));
    assert!(!dim.matches(513));
}

#[test]
fn test_shape_guard_check_pass() {
    let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);

    match guard.check(&[32, 256, 256], Some("fp32")) {
        GuardCheckResult::Passed => {}
        GuardCheckResult::Failed(f) => panic!("Guard should pass: {:?}", f),
    }
}

#[test]
fn test_shape_guard_check_rank_mismatch() {
    let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);

    match guard.check(&[32, 256], None) {
        GuardCheckResult::Failed(GuardFailure::RankMismatch { .. }) => {}
        other => panic!("Expected RankMismatch, got {:?}", other),
    }
}

#[test]
fn test_shape_guard_check_dim_mismatch() {
    let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);

    match guard.check(&[64, 256, 256], None) {
        GuardCheckResult::Failed(GuardFailure::DimensionMismatch { dim: 0, .. }) => {}
        other => panic!("Expected DimensionMismatch at dim 0, got {:?}", other),
    }
}

#[test]
fn test_shape_guard_symbolic_bindings() {
    let guard = ShapeGuard::new(
        "input_0",
        vec![
            ShapeDim::symbolic("batch"),
            ShapeDim::concrete(256),
            ShapeDim::symbolic("seq_len"),
        ],
    );

    let bindings = guard.extract_bindings(&[32, 256, 512]);
    assert_eq!(bindings.get("batch"), Some(&32));
    assert_eq!(bindings.get("seq_len"), Some(&512));
    assert_eq!(bindings.len(), 2);
}

#[test]
fn test_shape_signature_hash() {
    let inputs1 = vec![("input_0".to_string(), vec![32, 256], "fp32".to_string())];
    let inputs2 = vec![("input_0".to_string(), vec![32, 256], "fp32".to_string())];
    let inputs3 = vec![("input_0".to_string(), vec![64, 256], "fp32".to_string())];

    let sig1 = ShapeSignature::from_inputs(&inputs1, 12345);
    let sig2 = ShapeSignature::from_inputs(&inputs2, 12345);
    let sig3 = ShapeSignature::from_inputs(&inputs3, 12345);

    assert_eq!(sig1.to_cache_key(), sig2.to_cache_key());
    assert_ne!(sig1.to_cache_key(), sig3.to_cache_key());
}

#[test]
fn test_shape_guard_manager_stats() {
    let manager = ShapeGuardManager::default_manager();
    assert_eq!(manager.stats().total_checks, 0);
    assert_eq!(manager.stats().cache_hits, 0);
}

#[cfg(feature = "disabled_tests")]
mod integration_tests {
    use super::*;
    use crate::{
        SynthesisEngine,
        aot::{CompiledGraph, ExecutionPlan},
        templates,
    };

    #[tokio::test]
    async fn test_synthesis_engine_creation() {
        let result = SynthesisEngine::new("sm_120");
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_synthesis_engine_initialization() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
        let result = engine.initialize().await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_json_ir_parsing() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        // Test simple JSON operation
        let json_ir = r#"{
            "id": "test_op",
            "op_type": "gemm",
            "attributes": {
                "M": 256,
                "N": 256,
                "K": 256,
                "transpose_a": false,
                "transpose_b": true
            },
            "inputs": ["input_a", "input_b"],
            "outputs": ["output_c"]
        }"#;

        let operations = engine
            .parse_graph_ir(json_ir)
            .expect("Should parse JSON successfully");
        assert_eq!(operations.len(), 1);

        let op = &operations[0];
        assert_eq!(op.id, "test_op");
        assert_eq!(op.inputs, vec!["input_a", "input_b"]);
        assert_eq!(op.outputs, vec!["output_c"]);

        if let templates::KernelOperation::Gemm {
            m,
            n,
            k,
            transpose_a,
            transpose_b,
        } = op.operation
        {
            assert_eq!(m, 256);
            assert_eq!(n, 256);
            assert_eq!(k, 256);
            assert!(!transpose_a);
            assert!(transpose_b);
        } else {
            panic!("Expected GEMM operation");
        }
    }

    #[tokio::test]
    async fn test_json_array_ir_parsing() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        // Test array of operations
        let json_ir = r#"[
            {
                "id": "conv_op",
                "op_type": "conv2d",
                "attributes": {
                    "batch_size": 1,
                    "in_channels": 64,
                    "out_channels": 128,
                    "height": 224,
                    "width": 224,
                    "kernel_size": 3
                }
            },
            {
                "id": "relu_op",
                "op_type": "relu",
                "attributes": {
                    "size": 200704
                }
            }
        ]"#;

        let operations = engine
            .parse_graph_ir(json_ir)
            .expect("Should parse JSON array successfully");
        assert_eq!(operations.len(), 2);

        // Check convolution
        if let templates::KernelOperation::Convolution {
            batch_size,
            in_channels,
            out_channels,
            height,
            width,
            kernel_size,
        } = operations[0].operation
        {
            assert_eq!(batch_size, 1);
            assert_eq!(in_channels, 64);
            assert_eq!(out_channels, 128);
            assert_eq!(height, 224);
            assert_eq!(width, 224);
            assert_eq!(kernel_size, 3);
        } else {
            panic!("Expected Convolution operation");
        }

        // Check ReLU
        if let templates::KernelOperation::Elementwise { operation, size } = operations[1].operation
        {
            assert_eq!(size, 200704);
            if let templates::ElementwiseOp::ReLU = operation {
                // Correct
            } else {
                panic!("Expected ReLU elementwise operation");
            }
        } else {
            panic!("Expected Elementwise operation");
        }
    }

    #[tokio::test]
    async fn test_custom_dsl_parsing() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        let dsl_ir = r#"
# Test DSL format
op matrix_mult gemm
  inputs: A, B
  outputs: C
  fusable: true

op activation add
  inputs: C, bias
  outputs: result
  fusable: true
"#;

        let operations = engine
            .parse_graph_ir(dsl_ir)
            .expect("Should parse DSL successfully");
        assert_eq!(operations.len(), 2);

        // Check first operation
        let gemm_op = &operations[0];
        assert_eq!(gemm_op.id, "matrix_mult");
        assert_eq!(gemm_op.inputs, vec!["A", "B"]);
        assert_eq!(gemm_op.outputs, vec!["C"]);
        assert!(gemm_op.fusable);

        // Check second operation
        let add_op = &operations[1];
        assert_eq!(add_op.id, "activation");
        assert_eq!(add_op.inputs, vec!["C", "bias"]);
        assert_eq!(add_op.outputs, vec!["result"]);
        assert!(add_op.fusable);
    }

    #[tokio::test]
    async fn test_graph_object_parsing() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        let graph_json = r#"{
            "graph": {
                "nodes": [
                    {
                        "id": "attention_op",
                        "op_type": "attention",
                        "attributes": {
                            "sequence_length": 1024,
                            "head_dim": 64,
                            "num_heads": 16
                        }
                    }
                ],
                "edges": [
                    {"from": "input", "to": "attention_op"}
                ]
            }
        }"#;

        let operations = engine
            .parse_graph_ir(graph_json)
            .expect("Should parse graph JSON successfully");
        assert_eq!(operations.len(), 1);

        if let templates::KernelOperation::Attention {
            sequence_length,
            head_dim,
            num_heads,
        } = operations[0].operation
        {
            assert_eq!(sequence_length, 1024);
            assert_eq!(head_dim, 64);
            assert_eq!(num_heads, 16);
        } else {
            panic!("Expected Attention operation");
        }
    }

    #[test]
    fn test_flops_calculation() {
        let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        // Test GEMM FLOPS
        let gemm_op = templates::KernelOperation::Gemm {
            m: 128,
            n: 128,
            k: 128,
            transpose_a: false,
            transpose_b: false,
        };
        let flops = engine.calculate_theoretical_flops(&gemm_op);
        let expected = 2.0 * 128.0 * 128.0 * 128.0; // 2 * M * N * K
        assert!(
            (flops - expected).abs() < 1.0,
            "GEMM FLOPS calculation incorrect"
        );

        // Test Convolution FLOPS
        let conv_op = templates::KernelOperation::Convolution {
            batch_size: 1,
            in_channels: 64,
            out_channels: 128,
            height: 224,
            width: 224,
            kernel_size: 3,
        };
        let conv_flops = engine.calculate_theoretical_flops(&conv_op);
        let expected_conv = 2.0 * 1.0 * 128.0 * 224.0 * 224.0 * 64.0 * 9.0; // 2 * B * OC * H * W * IC * K^2
        assert!(
            (conv_flops - expected_conv).abs() < 1.0,
            "Convolution FLOPS calculation incorrect"
        );

        // Test Elementwise FLOPS
        let elem_op = templates::KernelOperation::Elementwise {
            operation: templates::ElementwiseOp::Add,
            size: 1024,
        };
        let elem_flops = engine.calculate_theoretical_flops(&elem_op);
        assert_eq!(elem_flops, 1024.0, "Elementwise FLOPS should equal size");
    }

    #[test]
    fn test_memory_usage_estimation() {
        let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        // Test GEMM memory usage
        let gemm_op = templates::KernelOperation::Gemm {
            m: 128,
            n: 128,
            k: 128,
            transpose_a: false,
            transpose_b: false,
        };
        let memory = engine.estimate_memory_usage(&gemm_op);
        let expected = 4.0 * ((128 * 128) + (128 * 128) + (128 * 128)) as f64; // A + B + C in FP32
        assert!(
            (memory - expected).abs() < 1.0,
            "GEMM memory estimation incorrect"
        );

        // Test Elementwise memory usage
        let elem_op = templates::KernelOperation::Elementwise {
            operation: templates::ElementwiseOp::Add,
            size: 1024,
        };
        let elem_memory = engine.estimate_memory_usage(&elem_op);
        assert_eq!(
            elem_memory,
            8.0 * 1024.0,
            "Elementwise memory should be 2 * size * sizeof(f32)"
        );
    }

    #[tokio::test]
    async fn test_parameter_space_creation() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
        engine.initialize().await.expect("Should initialize");

        let template = templates::KernelTemplate::new("test_template".to_string());
        let gemm_op = templates::KernelOperation::Gemm {
            m: 256,
            n: 256,
            k: 256,
            transpose_a: false,
            transpose_b: false,
        };

        let param_space = engine
            .create_parameter_space(&template, &gemm_op)
            .expect("Should create parameter space");

        // Check that GEMM-specific parameters are included
        assert!(param_space.dimensions.contains_key("tile_m"));
        assert!(param_space.dimensions.contains_key("tile_n"));
        assert!(param_space.dimensions.contains_key("tile_k"));
        assert!(param_space.dimensions.contains_key("block_size_x"));
        assert!(param_space.dimensions.contains_key("block_size_y"));

        // Check constraints
        assert!(!param_space.constraints.is_empty());
        assert!(
            param_space
                .constraints
                .iter()
                .any(|c| c.contains("block_size_x * block_size_y"))
        );
    }

    #[test]
    fn test_parameter_conversion() {
        let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        let mut int_params = std::collections::HashMap::new();
        int_params.insert("block_size_x".to_string(), 128);
        int_params.insert("block_size_y".to_string(), 8);
        int_params.insert("tile_m".to_string(), 64);

        let template_params = engine
            .convert_to_template_params(&int_params)
            .expect("Should convert parameters");

        assert_eq!(template_params["block_size_x"], "128");
        assert_eq!(template_params["block_size_y"], "8");
        assert_eq!(template_params["tile_m"], "64");
    }

    #[tokio::test]
    async fn test_invalid_ir_formats() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        // Test invalid JSON
        let invalid_json = r#"{ invalid json }"#;
        let result = engine.parse_graph_ir(invalid_json);
        assert!(result.is_err(), "Should reject invalid JSON");

        // Test unknown format
        let unknown_format = "this is not a known IR format";
        let result = engine.parse_graph_ir(unknown_format);
        assert!(result.is_err(), "Should reject unknown format");
    }

    #[test]
    fn test_compiled_kernel_structure() {
        use crate::CompiledKernel;

        let compiled_graph = CompiledGraph {
            kernel_sequence: vec![],
            memory_plan: std::collections::HashMap::new(),
            execution_plan: ExecutionPlan {
                kernel_launches: vec![],
                sync_points: vec![],
            },
        };

        let kernel = CompiledKernel {
            code: "test kernel code".to_string(),
            compiled_graph,
        };

        assert_eq!(kernel.code, "test kernel code");
        assert!(kernel.compiled_graph.kernel_sequence.is_empty());
    }

    #[tokio::test]
    async fn test_tuning_config_creation() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
        engine.initialize().await.expect("Should initialize");

        // Test GEMM config (should have more iterations)
        let gemm_op = templates::KernelOperation::Gemm {
            m: 128,
            n: 128,
            k: 128,
            transpose_a: false,
            transpose_b: false,
        };
        let gemm_config = engine
            .create_tuning_config(&gemm_op)
            .expect("Should create GEMM config");
        assert_eq!(gemm_config.max_iterations, 100);

        // Test Attention config
        let attention_op = templates::KernelOperation::Attention {
            sequence_length: 512,
            head_dim: 64,
            num_heads: 8,
        };
        let attention_config = engine
            .create_tuning_config(&attention_op)
            .expect("Should create Attention config");
        assert_eq!(attention_config.max_iterations, 80);

        // Test default config for other operations
        let elem_op = templates::KernelOperation::Elementwise {
            operation: templates::ElementwiseOp::Add,
            size: 1024,
        };
        let elem_config = engine
            .create_tuning_config(&elem_op)
            .expect("Should create default config");
        assert_eq!(elem_config.max_iterations, 50);
    }

    #[test]
    fn test_execution_time_estimation() {
        let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");

        let mock_graph = CompiledGraph {
            kernel_sequence: vec![],
            memory_plan: std::collections::HashMap::new(),
            execution_plan: ExecutionPlan {
                kernel_launches: vec![],
                sync_points: vec![],
            },
        };

        let gemm_op = templates::KernelOperation::Gemm {
            m: 1024,
            n: 1024,
            k: 1024,
            transpose_a: false,
            transpose_b: false,
        };
        let exec_time = engine.estimate_execution_time(&gemm_op, &mock_graph);

        // Should be reasonable execution time (not zero, not too large)
        assert!(exec_time.as_secs_f64() > 0.0);
        assert!(exec_time.as_secs_f64() < 10.0); // Should complete within 10 seconds for reasonable sizes
    }
}