edgefirst-client 2.9.5

// SPDX-License-Identifier: Apache-2.0
// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.

//! Test the samples.populate2 API with a real image and bounding box
//! annotation.
//!
//! This example creates a 640x480 image with a red circle in the top-left
//! quadrant and uploads it with a bounding box annotation around the circle.

use edgefirst_client::{Annotation, Box2d, Client, Error, Sample, SampleFile};
use std::env;

#[tokio::main]
async fn main() -> Result<(), Error> {
    env_logger::init();

    // Get command line arguments
    let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        eprintln!("Usage: {} <project_name> <dataset_name>", args[0]);
        eprintln!("\nExample:");
        eprintln!("  {} 'Unit Testing' 'Test Labels'", args[0]);
        std::process::exit(1);
    }

    let project_name = &args[1];
    let dataset_name = &args[2];

    // Get credentials from environment
    let server = env::var("STUDIO_SERVER").unwrap_or_else(|_| "test".to_string());
    let username =
        env::var("STUDIO_USERNAME").expect("STUDIO_USERNAME environment variable not set");
    let password =
        env::var("STUDIO_PASSWORD").expect("STUDIO_PASSWORD environment variable not set");

    println!("Connecting to EdgeFirst Studio ({} server)...", server);
    let client = Client::new()?
        .with_server(&server)?
        .with_login(&username, &password)
        .await?;
    client.verify_token().await?;
    println!("✓ Connected successfully");

    // Find the project
    println!("\nFinding project '{}'...", project_name);
    let projects = client.projects(Some(project_name)).await?;
    if projects.is_empty() {
        eprintln!("Error: Project '{}' not found", project_name);
        std::process::exit(1);
    }

    let project = projects
        .iter()
        .find(|p| p.name() == project_name)
        .or_else(|| projects.first())
        .ok_or_else(|| Error::InvalidParameters(format!("Project '{}' not found", project_name)))?;

    println!("✓ Found project: {} (ID: {})", project.name(), project.id());

    // Find the dataset
    println!("\nFinding dataset '{}'...", dataset_name);
    let datasets = client.datasets(project.id(), Some(dataset_name)).await?;

    let dataset = datasets.first().ok_or_else(|| {
        Error::InvalidParameters(format!(
            "Dataset '{}' not found in project '{}'",
            dataset_name,
            project.name()
        ))
    })?;
    println!("✓ Found dataset: {} (ID: {})", dataset.name(), dataset.id());

    // Get annotation sets for the dataset
    println!("\nFetching annotation sets...");
    let annotation_sets = client.annotation_sets(dataset.id()).await?;
    if annotation_sets.is_empty() {
        eprintln!(
            "Error: No annotation sets found for dataset '{}'",
            dataset_name
        );
        eprintln!("Please create an annotation set in Studio first.");
        std::process::exit(1);
    }

    let annotation_set = annotation_sets.first().unwrap();
    println!(
        "✓ Using annotation set: {} (ID: {})",
        annotation_set.name(),
        annotation_set.id()
    );

    // Generate image with circle
    println!("\nGenerating 640x480 image with red circle...");
    let (image_data, circle_bbox) = generate_image_with_circle();

    // Create temporary file
    let timestamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs();
    let temp_dir = std::env::temp_dir();
    let test_image_path = temp_dir.join(format!("test_circle_{}.png", timestamp));

    std::fs::write(&test_image_path, &image_data).expect("Failed to write test image");

    // Also save a copy to the current directory for inspection
    let local_copy = format!("test_circle_{}.png", timestamp);
    std::fs::write(&local_copy, &image_data).expect("Failed to write local copy");

    println!("  ✓ Created image: {:?}", test_image_path);
    println!("  ✓ Local copy saved: {}", local_copy);
    println!(
        "  ✓ Circle bounding box: x={:.1}, y={:.1}, w={:.1}, h={:.1}",
        circle_bbox.0, circle_bbox.1, circle_bbox.2, circle_bbox.3
    );

    // Create sample with annotation
    let mut sample = Sample::new();
    let img_width = 640.0;
    let img_height = 480.0;
    sample.width = Some(img_width as u32);
    sample.height = Some(img_height as u32);
    // UUID will be auto-generated by populate_samples()
    sample.group = Some("train".to_string());

    // Add bounding box annotation around the circle
    // IMPORTANT: Box2d coordinates are ALWAYS normalized (0.0-1.0 range)
    // Box2d::new() takes (x, y, width, height) all in normalized coordinates
    let mut annotation = Annotation::new();
    annotation.set_label(Some("circle".to_string())); // Server will auto-create label if needed
    annotation.set_object_id(Some("circle-object-1".to_string())); // Optional: for tracking objects across frames

    // Normalize coordinates: divide pixel values by image dimensions
    let normalized_x = circle_bbox.0 / img_width;
    let normalized_y = circle_bbox.1 / img_height;
    let normalized_w = circle_bbox.2 / img_width;
    let normalized_h = circle_bbox.3 / img_height;

    let bbox = Box2d::new(normalized_x, normalized_y, normalized_w, normalized_h);
    annotation.set_box2d(Some(bbox));
    sample.annotations = vec![annotation];

    // Add the image file
    let image_file = SampleFile::with_filename(
        "image".to_string(),
        test_image_path.to_str().unwrap().to_string(),
    );
    sample.files = vec![image_file];

    println!("\n=== Sample Details ===");
    println!("UUID: <will be auto-generated>");
    println!(
        "Dimensions: {}x{}",
        sample.width.unwrap(),
        sample.height.unwrap()
    );
    println!("Group: {}", sample.group().as_ref().unwrap());
    println!("Annotations: {}", sample.annotations.len());
    println!(
        "Bounding Box (pixel): ({:.1}, {:.1}) - {:.1}x{:.1}",
        circle_bbox.0, circle_bbox.1, circle_bbox.2, circle_bbox.3
    );
    println!(
        "Bounding Box (normalized): ({:.3}, {:.3}) - {:.3}x{:.3}",
        normalized_x, normalized_y, normalized_w, normalized_h
    );

    // Show JSON with annotation
    println!("\n=== JSON Preview ===");
    let sample_json = serde_json::to_string_pretty(&sample).unwrap();
    println!("{}", sample_json);

    // Call populate API with progress tracking
    println!("\nCalling samples.populate2 API with annotation_set_id...");

    let (tx, mut rx) = tokio::sync::mpsc::channel::<edgefirst_client::Progress>(1);

    // Spawn task to print progress
    let progress_task = tokio::spawn(async move {
        while let Some(progress) = rx.recv().await {
            println!("Upload progress: {}/{}", progress.current, progress.total);
        }
    });

    match client
        .populate_samples(
            dataset.id(),
            Some(annotation_set.id()),
            vec![sample],
            Some(tx),
        )
        .await
    {
        Ok(results) => {
            println!("✓ API call successful!");
            println!(
                "\nResponse: {} sample(s) populated and uploaded",
                results.len()
            );

            for (idx, result) in results.iter().enumerate() {
                println!("\n  Sample #{} (UUID: {})", idx + 1, result.uuid);
                println!("  Files uploaded: {}", result.urls.len());
            }
        }
        Err(e) => {
            eprintln!("✗ API call failed!");
            eprintln!("\nError: {:?}", e);
            std::process::exit(1);
        }
    }

    // Wait for progress task to complete
    progress_task.await.unwrap();

    // Clean up
    let _ = std::fs::remove_file(&test_image_path);

    println!("\n✓ Test completed successfully!");
    println!("\n=== VERIFICATION ===");
    println!("Local Image: test_circle_{}.png", timestamp);
    println!("  - You can open this image locally to verify it looks correct");
    println!("  - Should show a 640x480 image with a red circle in the top-left quadrant");
    println!("\nStudio UUID: test-circle-{}", timestamp);
    println!("  - Go to EdgeFirst Studio to find this sample");
    println!("  - Verify the image appears correctly in the dataset");
    println!("  - Verify the bounding box annotation appears around the circle");
    println!("\n=== IMAGE DETAILS ===");
    println!("  - Size: 640x480 pixels");
    println!("  - Circle: Red, centered in top-left quadrant");
    println!("  - Circle center: ({:.0}, {:.0})", 160.0, 120.0);
    println!("  - Circle radius: 50 pixels");
    println!("\n=== ANNOTATION DETAILS ===");
    println!("  - Label: circle");
    println!("  - Object Reference: circle-object-1");
    println!(
        "  - Bounding box (pixel): x={:.1}, y={:.1}, w={:.1}, h={:.1}",
        circle_bbox.0, circle_bbox.1, circle_bbox.2, circle_bbox.3
    );
    println!(
        "  - Bounding box (normalized): x={:.3}, y={:.3}, w={:.3}, h={:.3}",
        normalized_x, normalized_y, normalized_w, normalized_h
    );

    Ok(())
}

/// Generate a 640x480 PNG image with a red circle in the top-left quadrant.
/// Returns (image_data, bounding_box) where bounding_box is (x, y, width,
/// height).
fn generate_image_with_circle() -> (Vec<u8>, (f32, f32, f32, f32)) {
    // Image dimensions
    let width = 640u32;
    let height = 480u32;

    // Circle parameters - center in top-left quadrant (0-320, 0-240)
    let circle_center_x = 160.0f32; // Middle of left half
    let circle_center_y = 120.0f32; // Middle of top half
    let circle_radius = 50.0f32;

    // Bounding box around the circle (with small padding)
    let padding = 2.0f32;
    let bbox_x = circle_center_x - circle_radius - padding;
    let bbox_y = circle_center_y - circle_radius - padding;
    let bbox_width = (circle_radius * 2.0) + (padding * 2.0);
    let bbox_height = (circle_radius * 2.0) + (padding * 2.0);

    // Create RGB image buffer (white background)
    let mut image_data = vec![255u8; (width * height * 3) as usize];

    // Draw red circle
    for y in 0..height {
        for x in 0..width {
            let dx = (x as f32) - circle_center_x;
            let dy = (y as f32) - circle_center_y;
            let distance = (dx * dx + dy * dy).sqrt();

            if distance <= circle_radius {
                let idx = ((y * width + x) * 3) as usize;
                image_data[idx] = 255; // R
                image_data[idx + 1] = 0; // G
                image_data[idx + 2] = 0; // B
            }
        }
    }

    // Encode as PNG
    let mut png_data = Vec::new();
    {
        let mut encoder = png::Encoder::new(&mut png_data, width, height);
        encoder.set_color(png::ColorType::Rgb);
        encoder.set_depth(png::BitDepth::Eight);

        let mut writer = encoder.write_header().expect("Failed to write PNG header");
        writer
            .write_image_data(&image_data)
            .expect("Failed to write PNG data");
    }

    (png_data, (bbox_x, bbox_y, bbox_width, bbox_height))
}