mod cli;
use bounding_box::roi::MultiRoi;
use detector::*;
use detector::{database::FaceDatabase, facedet, facedet::FaceDetectionConfig, faceembed};
use errors::*;
use fast_image_resize::ResizeOptions;

use ndarray::*;
use ndarray_image::*;
use ndarray_resize::NdFir;
const RETINAFACE_MODEL_MNN: &[u8] = include_bytes!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/models/retinaface.mnn"
));
const FACENET_MODEL_MNN: &[u8] =
    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/models/facenet.mnn"));
const RETINAFACE_MODEL_ONNX: &[u8] = include_bytes!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/models/retinaface.onnx"
));
const FACENET_MODEL_ONNX: &[u8] =
    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/models/facenet.onnx"));
pub fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter("info")
        .with_thread_ids(true)
        .with_thread_names(true)
        .with_target(false)
        .init();
    let args = <cli::Cli as clap::Parser>::parse();
    match args.cmd {
        cli::SubCommand::Detect(detect) => {
            // Choose backend based on executor type (defaulting to MNN for backward compatibility)

            let executor = detect
                .mnn_forward_type
                .map(|f| cli::Executor::Mnn(f))
                .or_else(|| {
                    if detect.ort_execution_provider.is_empty() {
                        None
                    } else {
                        Some(cli::Executor::Ort(detect.ort_execution_provider.clone()))
                    }
                })
                .unwrap_or(cli::Executor::Mnn(mnn::ForwardType::CPU));

            match executor {
                cli::Executor::Mnn(forward) => {
                    let retinaface =
                        facedet::retinaface::mnn::FaceDetection::builder(RETINAFACE_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::mnn::EmbeddingGenerator::builder(FACENET_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_detection(detect, retinaface, facenet)?;
                }
                cli::Executor::Ort(ep) => {
                    let retinaface =
                        facedet::retinaface::ort::FaceDetection::builder(RETINAFACE_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(&ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::ort::EmbeddingGenerator::builder(FACENET_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_detection(detect, retinaface, facenet)?;
                }
            }
        }
        cli::SubCommand::DetectMulti(detect_multi) => {
            // Choose backend based on executor type (defaulting to MNN for backward compatibility)

            let executor = detect_multi
                .mnn_forward_type
                .map(|f| cli::Executor::Mnn(f))
                .or_else(|| {
                    if detect_multi.ort_execution_provider.is_empty() {
                        None
                    } else {
                        Some(cli::Executor::Ort(
                            detect_multi.ort_execution_provider.clone(),
                        ))
                    }
                })
                .unwrap_or(cli::Executor::Mnn(mnn::ForwardType::CPU));

            match executor {
                cli::Executor::Mnn(forward) => {
                    let retinaface =
                        facedet::retinaface::mnn::FaceDetection::builder(RETINAFACE_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::mnn::EmbeddingGenerator::builder(FACENET_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_multi_detection(detect_multi, retinaface, facenet)?;
                }
                cli::Executor::Ort(ep) => {
                    let retinaface =
                        facedet::retinaface::ort::FaceDetection::builder(RETINAFACE_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(&ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::ort::EmbeddingGenerator::builder(FACENET_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_multi_detection(detect_multi, retinaface, facenet)?;
                }
            }
        }
        cli::SubCommand::Query(query) => {
            run_query(query)?;
        }
        cli::SubCommand::Similar(similar) => {
            run_similar(similar)?;
        }
        cli::SubCommand::Stats(stats) => {
            run_stats(stats)?;
        }
        cli::SubCommand::Compare(compare) => {
            // Choose backend based on executor type (defaulting to MNN for backward compatibility)
            let executor = compare
                .mnn_forward_type
                .map(|f| cli::Executor::Mnn(f))
                .or_else(|| {
                    if compare.ort_execution_provider.is_empty() {
                        None
                    } else {
                        Some(cli::Executor::Ort(compare.ort_execution_provider.clone()))
                    }
                })
                .unwrap_or(cli::Executor::Mnn(mnn::ForwardType::CPU));

            match executor {
                cli::Executor::Mnn(forward) => {
                    let retinaface =
                        facedet::retinaface::mnn::FaceDetection::builder(RETINAFACE_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::mnn::EmbeddingGenerator::builder(FACENET_MODEL_MNN)
                            .change_context(Error)?
                            .with_forward_type(forward)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_compare(compare, retinaface, facenet)?;
                }
                cli::Executor::Ort(ep) => {
                    let retinaface =
                        facedet::retinaface::ort::FaceDetection::builder(RETINAFACE_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(&ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face detection model")?;
                    let facenet =
                        faceembed::facenet::ort::EmbeddingGenerator::builder(FACENET_MODEL_ONNX)
                            .change_context(Error)?
                            .with_execution_providers(ep)
                            .build()
                            .change_context(errors::Error)
                            .attach_printable("Failed to create face embedding model")?;

                    run_compare(compare, retinaface, facenet)?;
                }
            }
        }
        cli::SubCommand::Gui => {
            if let Err(e) = detector::gui::run() {
                eprintln!("GUI error: {}", e);
                std::process::exit(1);
            }
        }
        cli::SubCommand::Completions { shell } => {
            cli::Cli::completions(shell);
        }
    }
    Ok(())
}

fn run_detection<D, E>(detect: cli::Detect, mut retinaface: D, mut facenet: E) -> Result<()>
where
    D: facedet::FaceDetector,
    E: faceembed::FaceEmbedder,
{
    // Initialize database if requested
    let db = if detect.save_to_db {
        let db_path = detect
            .database
            .as_ref()
            .map(|p| p.as_path())
            .unwrap_or_else(|| std::path::Path::new("face_detections.db"));
        Some(FaceDatabase::new(db_path).change_context(Error)?)
    } else {
        None
    };
    let image = image::open(&detect.image)
        .change_context(Error)
        .attach_printable(detect.image.to_string_lossy().to_string())?;
    let image = image.into_rgb8();
    let (image_width, image_height) = image.dimensions();
    let mut array = image
        .into_ndarray()
        .change_context(errors::Error)
        .attach_printable("Failed to convert image to ndarray")?;
    let output = retinaface
        .detect_faces(
            array.view(),
            &FaceDetectionConfig::default()
                .with_threshold(detect.threshold)
                .with_nms_threshold(detect.nms_threshold),
        )
        .change_context(errors::Error)
        .attach_printable("Failed to detect faces")?;

    // Store image and face detections in database if requested
    let (_image_id, face_ids) = if let Some(ref database) = db {
        let image_path = detect.image.to_string_lossy();
        let img_id = database
            .store_image(&image_path, image_width, image_height)
            .change_context(Error)?;
        let face_ids = database
            .store_face_detections(img_id, &output)
            .change_context(Error)?;
        tracing::info!(
            "Stored image {} with {} faces in database",
            img_id,
            face_ids.len()
        );
        (Some(img_id), Some(face_ids))
    } else {
        (None, None)
    };

    for bbox in &output.bbox {
        tracing::info!("Detected face: {:?}", bbox);
        use bounding_box::draw::*;
        array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 1);
    }
    let face_rois = array
        .view()
        .multi_roi(&output.bbox)
        .change_context(Error)?
        .into_iter()
        // .inspect(|f| {
        //     tracing::info!("Face ROI shape before resize: {:?}", f.dim());
        // })
        .map(|roi| {
            roi.as_standard_layout()
                .fast_resize(320, 320, &ResizeOptions::default())
                .change_context(Error)
        })
        // .inspect(|f| {
        //     f.as_ref().inspect(|f| {
        //         tracing::info!("Face ROI shape after resize: {:?}", f.dim());
        //     });
        // })
        .collect::<Result<Vec<_>>>()?;
    let face_roi_views = face_rois.iter().map(|roi| roi.view()).collect::<Vec<_>>();

    let chunk_size = detect.batch_size;
    let embeddings = face_roi_views
        .chunks(chunk_size)
        .map(|chunk| {
            tracing::info!("Processing chunk of size: {}", chunk.len());

            if chunk.len() < chunk_size {
                tracing::warn!("Chunk size is less than 8, padding with zeros");
                let zeros = Array3::zeros((320, 320, 3));
                let chunk: Vec<_> = chunk
                    .iter()
                    .map(|arr| arr.reborrow())
                    .chain(core::iter::repeat(zeros.view()))
                    .take(chunk_size)
                    .collect();
                let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk.as_slice())
                    .change_context(errors::Error)
                    .attach_printable("Failed to stack rois together")?;
                let output = facenet.run_models(face_rois.view()).change_context(Error)?;
                Ok(output)
            } else {
                let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk)
                    .change_context(errors::Error)
                    .attach_printable("Failed to stack rois together")?;
                let output = facenet.run_models(face_rois.view()).change_context(Error)?;
                Ok(output)
            }
        })
        .collect::<Result<Vec<Array2<f32>>>>()?;

    // Store embeddings in database if requested
    if let (Some(database), Some(face_ids)) = (&db, &face_ids) {
        let embedding_ids = database
            .store_embeddings(face_ids, &embeddings, &detect.model_name)
            .change_context(Error)?;
        tracing::info!("Stored {} embeddings in database", embedding_ids.len());

        // Print database statistics
        let (num_images, num_faces, num_landmarks, num_embeddings) =
            database.get_stats().change_context(Error)?;
        tracing::info!(
            "Database stats - Images: {}, Faces: {}, Landmarks: {}, Embeddings: {}",
            num_images,
            num_faces,
            num_landmarks,
            num_embeddings
        );
    }

    let v = array.view();
    if let Some(output) = detect.output {
        let image: image::RgbImage = v
            .to_image()
            .change_context(errors::Error)
            .attach_printable("Failed to convert ndarray to image")?;
        image
            .save(output)
            .change_context(errors::Error)
            .attach_printable("Failed to save output image")?;
    }

    Ok(())
}

fn run_query(query: cli::Query) -> Result<()> {
    let db = FaceDatabase::new(&query.database).change_context(Error)?;

    if let Some(image_id) = query.image_id {
        if let Some(image) = db.get_image(image_id).change_context(Error)? {
            println!("Image: {}", image.file_path);
            println!("Dimensions: {}x{}", image.width, image.height);
            println!("Created: {}", image.created_at);

            let faces = db.get_faces_for_image(image_id).change_context(Error)?;
            println!("Faces found: {}", faces.len());

            for face in faces {
                println!(
                    "  Face ID {}: bbox({:.1}, {:.1}, {:.1}, {:.1}), confidence: {:.3}",
                    face.id,
                    face.bbox_x1,
                    face.bbox_y1,
                    face.bbox_x2,
                    face.bbox_y2,
                    face.confidence
                );

                if query.show_landmarks {
                    if let Some(landmarks) = db.get_landmarks(face.id).change_context(Error)? {
                        println!(
                            "    Landmarks: left_eye({:.1}, {:.1}), right_eye({:.1}, {:.1}), nose({:.1}, {:.1})",
                            landmarks.left_eye_x,
                            landmarks.left_eye_y,
                            landmarks.right_eye_x,
                            landmarks.right_eye_y,
                            landmarks.nose_x,
                            landmarks.nose_y
                        );
                    }
                }

                if query.show_embeddings {
                    let embeddings = db.get_embeddings(face.id).change_context(Error)?;
                    for embedding in embeddings {
                        println!(
                            "    Embedding ({}): {} dims, model: {}",
                            embedding.id,
                            embedding.embedding.len(),
                            embedding.model_name
                        );
                    }
                }
            }
        } else {
            println!("Image with ID {} not found", image_id);
        }
    }

    if let Some(face_id) = query.face_id {
        if let Some(landmarks) = db.get_landmarks(face_id).change_context(Error)? {
            println!(
                "Landmarks for face {}: left_eye({:.1}, {:.1}), right_eye({:.1}, {:.1}), nose({:.1}, {:.1})",
                face_id,
                landmarks.left_eye_x,
                landmarks.left_eye_y,
                landmarks.right_eye_x,
                landmarks.right_eye_y,
                landmarks.nose_x,
                landmarks.nose_y
            );
        } else {
            println!("No landmarks found for face {}", face_id);
        }

        let embeddings = db.get_embeddings(face_id).change_context(Error)?;
        println!(
            "Embeddings for face {}: {} found",
            face_id,
            embeddings.len()
        );
        for embedding in embeddings {
            println!(
                "  Embedding {}: {} dims, model: {}, created: {}",
                embedding.id,
                embedding.embedding.len(),
                embedding.model_name,
                embedding.created_at
            );
            // if query.show_embeddings {
            //     println!("    Values: {:?}", &embedding.embedding);
            // }
        }
    }

    Ok(())
}

fn run_compare<D, E>(compare: cli::Compare, mut retinaface: D, mut facenet: E) -> Result<()>
where
    D: facedet::FaceDetector,
    E: faceembed::FaceEmbedder,
{
    // Helper function to detect faces and compute embeddings for an image
    fn process_image<D, E>(
        image_path: &std::path::Path,
        retinaface: &mut D,
        facenet: &mut E,
        config: &FaceDetectionConfig,
        batch_size: usize,
    ) -> Result<(Vec<Array1<f32>>, usize)>
    where
        D: facedet::FaceDetector,
        E: faceembed::FaceEmbedder,
    {
        let image = image::open(image_path)
            .change_context(Error)
            .attach_printable(image_path.to_string_lossy().to_string())?;
        let image = image.into_rgb8();
        let array = image
            .into_ndarray()
            .change_context(errors::Error)
            .attach_printable("Failed to convert image to ndarray")?;

        let output = retinaface
            .detect_faces(array.view(), config)
            .change_context(errors::Error)
            .attach_printable("Failed to detect faces")?;

        tracing::info!(
            "Detected {} faces in {}",
            output.bbox.len(),
            image_path.display()
        );

        if output.bbox.is_empty() {
            return Ok((Vec::new(), 0));
        }

        let face_rois = array
            .view()
            .multi_roi(&output.bbox)
            .change_context(Error)?
            .into_iter()
            .map(|roi| {
                roi.as_standard_layout()
                    .fast_resize(320, 320, &ResizeOptions::default())
                    .change_context(Error)
            })
            .collect::<Result<Vec<_>>>()?;

        let face_roi_views = face_rois.iter().map(|roi| roi.view()).collect::<Vec<_>>();
        let chunk_size = batch_size;
        let embeddings = face_roi_views
            .chunks(chunk_size)
            .map(|chunk| {
                if chunk.len() < chunk_size {
                    let zeros = Array3::zeros((320, 320, 3));
                    let chunk: Vec<_> = chunk
                        .iter()
                        .map(|arr| arr.reborrow())
                        .chain(core::iter::repeat(zeros.view()))
                        .take(chunk_size)
                        .collect();
                    let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk.as_slice())
                        .change_context(errors::Error)
                        .attach_printable("Failed to stack rois together")?;
                    facenet.run_models(face_rois.view()).change_context(Error)
                } else {
                    let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk)
                        .change_context(errors::Error)
                        .attach_printable("Failed to stack rois together")?;
                    facenet.run_models(face_rois.view()).change_context(Error)
                }
            })
            .collect::<Result<Vec<Array2<f32>>>>()?;

        // Flatten embeddings into individual face embeddings
        let mut face_embeddings = Vec::new();
        for embedding_batch in embeddings {
            for i in 0..output.bbox.len().min(embedding_batch.nrows()) {
                face_embeddings.push(embedding_batch.row(i).to_owned());
            }
        }

        Ok((face_embeddings, output.bbox.len()))
    }

    // Helper function to compute cosine similarity between two embeddings
    fn cosine_similarity(a: &Array1<f32>, b: &Array1<f32>) -> f32 {
        let dot_product = a.dot(b);
        let norm_a = a.dot(a).sqrt();
        let norm_b = b.dot(b).sqrt();
        dot_product / (norm_a * norm_b)
    }

    let config = FaceDetectionConfig::default()
        .with_threshold(compare.threshold)
        .with_nms_threshold(compare.nms_threshold);

    // Process both images
    let (embeddings1, face_count1) = process_image(
        &compare.image1,
        &mut retinaface,
        &mut facenet,
        &config,
        compare.batch_size,
    )?;

    let (embeddings2, face_count2) = process_image(
        &compare.image2,
        &mut retinaface,
        &mut facenet,
        &config,
        compare.batch_size,
    )?;

    println!(
        "Image 1 ({}): {} faces detected",
        compare.image1.display(),
        face_count1
    );
    println!(
        "Image 2 ({}): {} faces detected",
        compare.image2.display(),
        face_count2
    );

    if embeddings1.is_empty() && embeddings2.is_empty() {
        println!("No faces detected in either image");
        return Ok(());
    }

    if embeddings1.is_empty() {
        println!("No faces detected in image 1");
        return Ok(());
    }

    if embeddings2.is_empty() {
        println!("No faces detected in image 2");
        return Ok(());
    }

    // Compare all faces between the two images
    println!("\nFace comparison results:");
    println!("========================");

    let mut max_similarity = f32::NEG_INFINITY;
    let mut best_match = (0, 0);

    for (i, emb1) in embeddings1.iter().enumerate() {
        for (j, emb2) in embeddings2.iter().enumerate() {
            let similarity = cosine_similarity(emb1, emb2);
            println!(
                "Face {} (image 1) vs Face {} (image 2): {:.4}",
                i + 1,
                j + 1,
                similarity
            );

            if similarity > max_similarity {
                max_similarity = similarity;
                best_match = (i + 1, j + 1);
            }
        }
    }

    println!(
        "\nBest match: Face {} (image 1) vs Face {} (image 2) with similarity: {:.4}",
        best_match.0, best_match.1, max_similarity
    );

    // Interpretation of similarity score
    if max_similarity > 0.8 {
        println!("Interpretation: Very likely the same person");
    } else if max_similarity > 0.6 {
        println!("Interpretation: Possibly the same person");
    } else if max_similarity > 0.4 {
        println!("Interpretation: Unlikely to be the same person");
    } else {
        println!("Interpretation: Very unlikely to be the same person");
    }

    Ok(())
}

fn run_multi_detection<D, E>(
    detect_multi: cli::DetectMulti,
    mut retinaface: D,
    mut facenet: E,
) -> Result<()>
where
    D: facedet::FaceDetector,
    E: faceembed::FaceEmbedder,
{
    use std::fs;

    // Initialize database - always save to database for multi-detection
    let db = FaceDatabase::new(&detect_multi.database).change_context(Error)?;

    // Parse supported extensions
    let extensions: std::collections::HashSet<String> = detect_multi
        .extensions
        .split(',')
        .map(|ext| ext.trim().to_lowercase())
        .collect();

    // Create output directory if specified
    if let Some(ref output_dir) = detect_multi.output_dir {
        fs::create_dir_all(output_dir)
            .change_context(Error)
            .attach_printable("Failed to create output directory")?;
    }

    // Read directory and filter image files
    let entries = fs::read_dir(&detect_multi.input_dir)
        .change_context(Error)
        .attach_printable("Failed to read input directory")?;

    let mut image_paths = Vec::new();
    for entry in entries {
        let entry = entry.change_context(Error)?;
        let path = entry.path();

        if path.is_file() {
            if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
                if extensions.contains(&ext.to_lowercase()) {
                    image_paths.push(path);
                }
            }
        }
    }

    if image_paths.is_empty() {
        tracing::warn!(
            "No image files found in directory: {:?}",
            detect_multi.input_dir
        );
        return Ok(());
    }

    tracing::info!("Found {} image files to process", image_paths.len());

    let mut total_faces = 0;
    let mut processed_images = 0;

    // Process each image
    for (idx, image_path) in image_paths.iter().enumerate() {
        tracing::info!(
            "Processing image {}/{}: {:?}",
            idx + 1,
            image_paths.len(),
            image_path
        );

        // Load and process image
        let image = match image::open(image_path) {
            Ok(img) => img.into_rgb8(),
            Err(e) => {
                tracing::error!("Failed to load image {:?}: {}", image_path, e);
                continue;
            }
        };

        let (image_width, image_height) = image.dimensions();
        let mut array = match image.into_ndarray().change_context(errors::Error) {
            Ok(arr) => arr,
            Err(e) => {
                tracing::error!("Failed to convert image to ndarray: {:?}", e);
                continue;
            }
        };

        let config = FaceDetectionConfig::default()
            .with_threshold(detect_multi.threshold)
            .with_nms_threshold(detect_multi.nms_threshold);
        // Detect faces
        let output = match retinaface.detect_faces(array.view(), &config) {
            Ok(output) => output,
            Err(e) => {
                tracing::error!("Failed to detect faces in {:?}: {:?}", image_path, e);
                continue;
            }
        };

        let num_faces = output.bbox.len();
        total_faces += num_faces;

        if num_faces == 0 {
            tracing::info!("No faces detected in {:?}", image_path);
        } else {
            tracing::info!("Detected {} faces in {:?}", num_faces, image_path);
        }

        // Store image and detections in database
        let image_path_str = image_path.to_string_lossy();
        let img_id = match db.store_image(&image_path_str, image_width, image_height) {
            Ok(id) => id,
            Err(e) => {
                tracing::error!("Failed to store image in database: {:?}", e);
                continue;
            }
        };

        let face_ids = match db.store_face_detections(img_id, &output) {
            Ok(ids) => ids,
            Err(e) => {
                tracing::error!("Failed to store face detections in database: {:?}", e);
                continue;
            }
        };

        // Draw bounding boxes if output directory is specified
        if detect_multi.output_dir.is_some() {
            for bbox in &output.bbox {
                use bounding_box::draw::*;
                array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 1);
            }
        }

        // Process face embeddings if faces were detected
        if !face_ids.is_empty() {
            let face_rois = match array.view().multi_roi(&output.bbox).change_context(Error) {
                Ok(rois) => rois,
                Err(e) => {
                    tracing::error!("Failed to extract face ROIs: {:?}", e);
                    continue;
                }
            };

            let face_rois: Result<Vec<_>> = face_rois
                .into_iter()
                .map(|roi| {
                    roi.as_standard_layout()
                        .fast_resize(320, 320, &ResizeOptions::default())
                        .change_context(Error)
                })
                .collect();

            let face_rois = match face_rois {
                Ok(rois) => rois,
                Err(e) => {
                    tracing::error!("Failed to resize face ROIs: {:?}", e);
                    continue;
                }
            };

            let face_roi_views = face_rois.iter().map(|roi| roi.view()).collect::<Vec<_>>();

            let chunk_size = detect_multi.batch_size;
            let embeddings: Result<Vec<Array2<f32>>> = face_roi_views
                .chunks(chunk_size)
                .map(|chunk| {
                    if chunk.len() < chunk_size {
                        let zeros = Array3::zeros((320, 320, 3));
                        let chunk: Vec<_> = chunk
                            .iter()
                            .map(|arr| arr.reborrow())
                            .chain(core::iter::repeat(zeros.view()))
                            .take(chunk_size)
                            .collect();
                        let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk.as_slice())
                            .change_context(errors::Error)
                            .attach_printable("Failed to stack rois together")?;
                        facenet.run_models(face_rois.view()).change_context(Error)
                    } else {
                        let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk)
                            .change_context(errors::Error)
                            .attach_printable("Failed to stack rois together")?;
                        facenet.run_models(face_rois.view()).change_context(Error)
                    }
                })
                .collect();

            let embeddings = match embeddings {
                Ok(emb) => emb,
                Err(e) => {
                    tracing::error!("Failed to generate embeddings: {:?}", e);
                    continue;
                }
            };

            // Store embeddings in database
            if let Err(e) = db.store_embeddings(&face_ids, &embeddings, &detect_multi.model_name) {
                tracing::error!("Failed to store embeddings in database: {:?}", e);
                continue;
            }
        }

        // Save output image if directory specified
        if let Some(ref output_dir) = detect_multi.output_dir {
            let output_filename = format!(
                "detected_{}",
                image_path.file_name().unwrap().to_string_lossy()
            );
            let output_path = output_dir.join(output_filename);

            let v = array.view();
            let output_image: image::RgbImage = match v.to_image().change_context(errors::Error) {
                Ok(img) => img,
                Err(e) => {
                    tracing::error!("Failed to convert ndarray to image: {:?}", e);
                    continue;
                }
            };

            if let Err(e) = output_image.save(&output_path) {
                tracing::error!("Failed to save output image to {:?}: {}", output_path, e);
                continue;
            }

            tracing::info!("Saved output image to {:?}", output_path);
        }

        processed_images += 1;
    }

    // Print final statistics
    tracing::info!(
        "Processing complete: {}/{} images processed successfully, {} total faces detected",
        processed_images,
        image_paths.len(),
        total_faces
    );

    let (num_images, num_faces, num_landmarks, num_embeddings) =
        db.get_stats().change_context(Error)?;
    tracing::info!(
        "Database stats - Images: {}, Faces: {}, Landmarks: {}, Embeddings: {}",
        num_images,
        num_faces,
        num_landmarks,
        num_embeddings
    );

    Ok(())
}

fn run_similar(similar: cli::Similar) -> Result<()> {
    let db = FaceDatabase::new(&similar.database).change_context(Error)?;

    let embeddings = db.get_embeddings(similar.face_id).change_context(Error)?;
    if embeddings.is_empty() {
        println!("No embeddings found for face {}", similar.face_id);
        return Ok(());
    }

    let query_embedding = &embeddings[0].embedding;
    let similar_faces = db
        .find_similar_faces(query_embedding, similar.threshold, similar.limit)
        .change_context(Error)?;
    // Get image information for the similar faces
    println!(
        "Found {} similar faces (threshold: {:.3}):",
        similar_faces.len(),
        similar.threshold
    );
    for (face_id, similarity) in &similar_faces {
        if let Some(image_info) = db.get_image_for_face(*face_id).change_context(Error)? {
            println!(
                "  Face {}: similarity {:.3}, image: {}",
                face_id, similarity, image_info.file_path
            );
        }
    }

    Ok(())
}

fn run_stats(stats: cli::Stats) -> Result<()> {
    let db = FaceDatabase::new(&stats.database).change_context(Error)?;
    let (images, faces, landmarks, embeddings) = db.get_stats().change_context(Error)?;

    println!("Database Statistics:");
    println!("  Images: {}", images);
    println!("  Faces: {}", faces);
    println!("  Landmarks: {}", landmarks);
    println!("  Embeddings: {}", embeddings);

    Ok(())
}
