//! Test to verify correct Device usage in rtx-transformers
//! RED phase: This test should fail with compilation errors first
//!
//! NOTE: Disabled until transformer API is fully implemented

#![cfg(feature = "disabled_tests")]

use rtx_tensor::{Device, Tensor};
use rtx_transformers::architectures::TransformerModel;
use rtx_transformers::architectures::bert::BertModel;

#[test]
fn test_bert_model_device_creation() {
    // RED: This test verifies that BertModel can be created with proper device
    let device = Device::cuda(0).unwrap_or(Device::default());

    // This should compile and run without issues
    let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device);
    assert!(model.is_ok(), "BertModel should be created successfully");

    let model = model.unwrap();
    // Verify the model is on the correct device
    assert_eq!(model.device(), &device);
}

#[test]
fn test_bert_model_forward_pass() {
    // RED: Test forward pass with proper device handling
    let device = Device::cuda(0).unwrap_or(Device::default());

    let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();

    // Create input tensors on the same device
    let input_ids = Tensor::zeros(&[2, 128], &device).unwrap();
    let attention_mask = Tensor::ones(&[2, 128], &device).unwrap();

    // Forward pass should work
    let output = model.forward(&input_ids, Some(&attention_mask));
    assert!(output.is_ok(), "Forward pass should succeed");

    let output = output.unwrap();
    assert_eq!(output.device(), &device, "Output should be on same device");
}

#[test]
fn test_bert_model_parameters_on_device() {
    // RED: Test that all parameters are on the correct device
    let device = Device::cuda(0).unwrap_or(Device::default());

    let model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();

    let params = model.parameters();
    assert!(!params.is_empty(), "Model should have parameters");

    // All parameters should be on the same device
    for param in params {
        assert_eq!(
            param.device(),
            &device,
            "All parameters should be on same device"
        );
    }
}

#[test]
fn test_bert_model_gradient_flow() {
    // RED: Test gradient flow with device consistency
    let device = Device::cuda(0).unwrap_or(Device::default());

    let mut model = BertModel::new(768, 12, 12, 512, 30522, 512, &device).unwrap();
    model.train(true);

    // Create input with requires_grad
    let mut input_ids = Tensor::zeros(&[1, 64], &device).unwrap();
    input_ids.set_requires_grad(true).unwrap();

    let output = model.forward(&input_ids, None).unwrap();

    // Output should track gradients
    assert!(
        output.requires_grad(),
        "Output should require gradients in training mode"
    );
}

#[test]
fn test_device_transfer() {
    // RED: Test transferring model between devices
    let initial_device = Device::cuda(0).unwrap_or(Device::default());

    let mut model = BertModel::new(768, 12, 12, 512, 30522, 512, &initial_device).unwrap();

    // Try to transfer to same device (should work)
    let result = model.to_device(&initial_device);
    assert!(result.is_ok(), "Transfer to same device should work");

    // Verify model is still on correct device
    assert_eq!(model.device(), &initial_device);
}