dspy-rs 0.7.3

use anyhow::Result;
use dspy_rs::data::dataloader::DataLoader;
use rstest::rstest;

#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_hf_awesome_chatgpt_prompts() -> Result<()> {
    // Load the HuggingFace dataset
    let input_keys = vec!["events".to_string(), "inputs".to_string()];
    let output_keys = vec!["output".to_string()];

    let examples = DataLoader::load_hf(
        "zed-industries/zeta",
        input_keys.clone(),
        output_keys.clone(),
        "",      // No specific subset
        "train", // Split to load
        true,    // Not verbose
    )?;

    // Verify we got some data
    assert!(
        !examples.is_empty(),
        "Should have loaded some examples from HuggingFace dataset"
    );

    // Check the first example has the expected structure
    let first_example = &examples[0];

    // Print available keys to debug

    // Verify input and output keys are set correctly
    assert_eq!(first_example.input_keys, input_keys);
    assert_eq!(first_example.output_keys, output_keys);

    // Check what fields are actually present
    let has_act = first_example.data.contains_key("act");
    let has_prompt = first_example.data.contains_key("prompt");

    // Verify the data contains the expected fields (this will now provide better error info)
    assert!(
        has_act || !first_example.keys().is_empty(),
        "Example should contain 'act' field or have some data. Available fields: {:?}",
        first_example.keys()
    );
    assert!(
        has_prompt || !first_example.keys().is_empty(),
        "Example should contain 'prompt' field or have some data. Available fields: {:?}",
        first_example.keys()
    );

    // If expected fields exist, verify they're not empty
    if has_act && has_prompt {
        let act_value = first_example.get("act", None);
        let prompt_value = first_example.get("prompt", None);
        assert!(!act_value.is_null(), "act field should not be null");
        assert!(!prompt_value.is_null(), "prompt field should not be null");

        // Convert to string for display
        let act_str = act_value.as_str().unwrap_or("");
        let prompt_str = prompt_value.as_str().unwrap_or("");
        assert!(!act_str.is_empty(), "act field should not be empty");
        assert!(!prompt_str.is_empty(), "prompt field should not be empty");
    }

    Ok(())
}

// Test loading CSV from URL: snakes_count_10.csv
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_csv_from_url() -> Result<()> {
    let url = "https://people.sc.fsu.edu/~jburkardt/data/csv/snakes_count_10.csv";
    let input_keys = vec!["Game Number".to_string()];
    let output_keys = vec!["Game Length".to_string()];

    let examples = DataLoader::load_csv(
        url,
        ',', // delimiter
        input_keys.clone(),
        output_keys.clone(),
        true, // has headers
    )?;

    // Verify we got some data
    assert!(
        !examples.is_empty(),
        "Should have loaded some examples from CSV"
    );
    assert_eq!(
        examples.len(),
        10,
        "Should have loaded exactly 10 game records"
    );

    // Check the first example
    let first_example = &examples[0];

    // Verify input and output keys are set correctly
    assert_eq!(first_example.input_keys, input_keys);
    assert_eq!(first_example.output_keys, output_keys);

    // Verify we have data (columns should be indexed as 0, 1, etc for CSV without named headers)
    assert!(
        !first_example.data.is_empty(),
        "Example should contain data"
    );

    Ok(())
}

// Test loading JSON from URL: grok-2 config.json
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_json_from_url() -> Result<()> {
    let url = "https://huggingface.co/xai-org/grok-2/raw/main/config.json";
    let input_keys = vec!["vocab_size".to_string(), "hidden_size".to_string()];
    let output_keys = vec![]; // No output keys for this config file

    // This is a single JSON object, not JSON lines
    let examples = DataLoader::load_json(
        url,
        false, // not JSON lines
        input_keys.clone(),
        output_keys.clone(),
    )?;

    // For a single JSON object, we expect it to be parsed as a single Example
    // or as an array of Examples depending on the structure
    assert!(!examples.is_empty(), "Should have loaded data from JSON");

    // Get the first (and likely only) example
    let config_example = &examples[0];

    // Verify the data contains the expected fields
    assert!(
        config_example.data.contains_key("vocab_size"),
        "Config should contain 'vocab_size' field"
    );
    assert!(
        config_example.data.contains_key("hidden_size"),
        "Config should contain 'hidden_size' field"
    );

    // Get and verify the values
    let vocab_size = config_example.get("vocab_size", None);
    let hidden_size = config_example.get("hidden_size", None);

    assert!(!vocab_size.is_null(), "vocab_size should not be null");
    assert!(!hidden_size.is_null(), "hidden_size should not be null");

    Ok(())
}

// Additional test: Load JSON with specific structure verification
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_json_grok2_with_multiple_fields() -> Result<()> {
    let url = "https://huggingface.co/xai-org/grok-2/raw/main/config.json";

    // Test loading with more comprehensive input keys
    let input_keys = vec![
        "vocab_size".to_string(),
        "hidden_size".to_string(),
        "intermediate_size".to_string(),
        "num_hidden_layers".to_string(),
    ];
    let output_keys = vec![];

    let examples = DataLoader::load_json(url, false, input_keys.clone(), output_keys.clone())?;

    assert!(!examples.is_empty(), "Should have loaded data from JSON");

    let config = &examples[0];

    // Verify all requested input fields exist
    for key in &input_keys {
        assert!(
            config.data.contains_key(key),
            "Config should contain '{key}' field"
        );
        let value = config.get(key, None);
        assert!(!value.is_null(), "{key} should not be null");
    }

    Ok(())
}

// Test CSV with headers parsing
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_csv_verify_columns() -> Result<()> {
    // First, let's load without specifying input/output keys to see all columns
    let url = "https://people.sc.fsu.edu/~jburkardt/data/csv/snakes_count_10.csv";
    let examples = DataLoader::load_csv(
        url,
        ',',
        vec![], // No specific input keys
        vec![], // No specific output keys
        true,   // has headers
    )?;

    assert!(!examples.is_empty(), "Should have loaded examples");

    // Examine the structure of the data
    let first_example = &examples[0];
    let keys = first_example.keys();

    // Verify we have exactly 10 rows (games)
    assert_eq!(examples.len(), 10, "Should have 10 game records");

    // Verify each example has the same structure
    for (i, example) in examples.iter().enumerate() {
        assert_eq!(
            example.keys().len(),
            keys.len(),
            "Row {i} should have same number of columns"
        );
    }

    Ok(())
}

// Test error handling for invalid URLs
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_invalid_url_handling() {
    let invalid_url = "https://invalid-url-that-does-not-exist.com/data.csv";

    let result = DataLoader::load_csv(
        invalid_url,
        ',',
        vec!["col1".to_string()],
        vec!["col2".to_string()],
        true,
    );

    assert!(result.is_err(), "Should fail when loading from invalid URL");
}

// Test HuggingFace dataset with specific split
#[rstest]
#[cfg_attr(miri, ignore = "MIRI has issues with network operations")]
fn test_load_hf_with_verbose() -> Result<()> {
    let input_keys = vec!["events".to_string(), "inputs".to_string()];
    let output_keys = vec!["output".to_string()];

    // Load with verbose output to see what files are being processed
    let examples = DataLoader::load_hf(
        "zed-industries/zeta",
        input_keys.clone(),
        output_keys.clone(),
        "",      // No specific subset
        "train", // Split
        true,    // Verbose - will print loading information
    )?;

    assert!(!examples.is_empty(), "Should have loaded examples");

    // Verify data integrity
    for example in examples.iter().take(3) {
        // Verify structure
        assert_eq!(example.input_keys, input_keys);
        assert_eq!(example.output_keys, output_keys);
    }

    Ok(())
}