use std::path::PathBuf;

use crate::errors;
use crate::facedet::{FaceDetectionConfig, FaceDetector, retinaface};
use crate::faceembed::facenet;
use crate::gui::app::{ComparisonResult, DetectionResult, ExecutorType};
use bounding_box::Aabb2;
use bounding_box::roi::MultiRoi as _;
use error_stack::ResultExt;
use fast_image_resize::ResizeOptions;
use ndarray::{Array1, Array2, Array3, Array4};
use ndarray_image::ImageToNdarray;
use ndarray_math::CosineSimilarity;
use ndarray_resize::NdFir;

const RETINAFACE_MODEL_MNN: &[u8] = include_bytes!("../../models/retinaface.mnn");
const FACENET_MODEL_MNN: &[u8] = include_bytes!("../../models/facenet.mnn");
const RETINAFACE_MODEL_ONNX: &[u8] = include_bytes!("../../models/retinaface.onnx");
const FACENET_MODEL_ONNX: &[u8] = include_bytes!("../../models/facenet.onnx");

pub struct FaceDetectionBridge;

impl FaceDetectionBridge {
    pub async fn detect_faces(
        image_path: PathBuf,
        output_path: Option<PathBuf>,
        threshold: f32,
        nms_threshold: f32,
        executor_type: ExecutorType,
    ) -> DetectionResult {
        let start_time = std::time::Instant::now();

        match Self::run_detection_internal(
            image_path.clone(),
            output_path,
            threshold,
            nms_threshold,
            executor_type,
        )
        .await
        {
            Ok((faces_count, processed_image)) => {
                let processing_time = start_time.elapsed().as_secs_f64();
                DetectionResult::Success {
                    image_path,
                    faces_count,
                    processed_image,
                    processing_time,
                }
            }
            Err(error) => DetectionResult::Error(error.to_string()),
        }
    }

    pub async fn compare_faces(
        image1_path: PathBuf,
        image2_path: PathBuf,
        threshold: f32,
        nms_threshold: f32,
        executor_type: ExecutorType,
    ) -> ComparisonResult {
        let start_time = std::time::Instant::now();

        match Self::run_comparison_internal(
            image1_path,
            image2_path,
            threshold,
            nms_threshold,
            executor_type,
        )
        .await
        {
            Ok((
                image1_faces,
                image2_faces,
                image1_face_rois,
                image2_face_rois,
                best_similarity,
            )) => {
                let processing_time = start_time.elapsed().as_secs_f64();
                ComparisonResult::Success {
                    image1_faces,
                    image2_faces,
                    image1_face_rois,
                    image2_face_rois,
                    best_similarity,
                    processing_time,
                }
            }
            Err(error) => ComparisonResult::Error(error.to_string()),
        }
    }

    async fn run_detection_internal(
        image_path: PathBuf,
        output_path: Option<PathBuf>,
        threshold: f32,
        nms_threshold: f32,
        executor_type: ExecutorType,
    ) -> Result<(usize, Option<Vec<u8>>), Box<dyn std::error::Error + Send + Sync>> {
        // Load the image
        let img = image::open(&image_path)?;
        let img_rgb = img.to_rgb8();

        // Convert to ndarray format
        let image_array = img_rgb.as_ndarray()?;

        // Create detection configuration
        let config = FaceDetectionConfig::default()
            .with_threshold(threshold)
            .with_nms_threshold(nms_threshold)
            .with_input_width(1024)
            .with_input_height(1024);

        // Create detector and detect faces
        let faces = match executor_type {
            ExecutorType::MnnCpu | ExecutorType::MnnMetal | ExecutorType::MnnCoreML => {
                let forward_type = match executor_type {
                    ExecutorType::MnnCpu => mnn::ForwardType::CPU,
                    ExecutorType::MnnMetal => mnn::ForwardType::Metal,
                    ExecutorType::MnnCoreML => mnn::ForwardType::CoreML,
                    _ => unreachable!(),
                };

                let mut detector = retinaface::mnn::FaceDetection::builder(RETINAFACE_MODEL_MNN)
                    .map_err(|e| format!("Failed to create MNN detector: {}", e))?
                    .with_forward_type(forward_type)
                    .build()
                    .map_err(|e| format!("Failed to build MNN detector: {}", e))?;

                detector
                    .detect_faces(image_array.view(), &config)
                    .map_err(|e| format!("Detection failed: {}", e))?
            }
            ExecutorType::OnnxCpu => {
                let mut detector = retinaface::ort::FaceDetection::builder(RETINAFACE_MODEL_ONNX)
                    .map_err(|e| format!("Failed to create ONNX detector: {}", e))?
                    .build()
                    .map_err(|e| format!("Failed to build ONNX detector: {}", e))?;

                detector
                    .detect_faces(image_array.view(), &config)
                    .map_err(|e| format!("Detection failed: {}", e))?
            }
        };

        let faces_count = faces.bbox.len();

        // Generate output image with bounding boxes if requested
        let processed_image = if output_path.is_some() || true {
            // Always generate for GUI display
            let mut output_img = img.to_rgb8();

            for bbox in &faces.bbox {
                let min_point = bbox.min_vertex();
                let size = bbox.size();
                let rect = imageproc::rect::Rect::at(min_point.x as i32, min_point.y as i32)
                    .of_size(size.x as u32, size.y as u32);
                imageproc::drawing::draw_hollow_rect_mut(
                    &mut output_img,
                    rect,
                    image::Rgb([255, 0, 0]),
                );
            }

            // Convert to bytes for GUI display
            let mut buffer = Vec::new();
            let mut cursor = std::io::Cursor::new(&mut buffer);
            image::DynamicImage::ImageRgb8(output_img.clone())
                .write_to(&mut cursor, image::ImageFormat::Png)?;

            // Save to file if output path is specified
            if let Some(ref output_path) = output_path {
                output_img.save(output_path)?;
            }

            Some(buffer)
        } else {
            None
        };

        Ok((faces_count, processed_image))
    }

    async fn run_comparison_internal(
        image1_path: PathBuf,
        image2_path: PathBuf,
        threshold: f32,
        nms_threshold: f32,
        executor_type: ExecutorType,
    ) -> Result<
        (usize, usize, Vec<Array3<u8>>, Vec<Array3<u8>>, f32),
        Box<dyn std::error::Error + Send + Sync>,
    > {
        // Create detector and embedder, detect faces and generate embeddings
        let (image1_faces, image2_faces, image1_rois, image2_rois, best_similarity) =
            match executor_type {
                ExecutorType::MnnCpu | ExecutorType::MnnMetal | ExecutorType::MnnCoreML => {
                    let forward_type = match executor_type {
                        ExecutorType::MnnCpu => mnn::ForwardType::CPU,
                        ExecutorType::MnnMetal => mnn::ForwardType::Metal,
                        ExecutorType::MnnCoreML => mnn::ForwardType::CoreML,
                        _ => unreachable!(),
                    };

                    let mut detector =
                        retinaface::mnn::FaceDetection::builder(RETINAFACE_MODEL_MNN)
                            .map_err(|e| format!("Failed to create MNN detector: {}", e))?
                            .with_forward_type(forward_type.clone())
                            .build()
                            .map_err(|e| format!("Failed to build MNN detector: {}", e))?;

                    let mut embedder = facenet::mnn::EmbeddingGenerator::builder(FACENET_MODEL_MNN)
                        .map_err(|e| format!("Failed to create MNN embedder: {}", e))?
                        .with_forward_type(forward_type)
                        .build()
                        .map_err(|e| format!("Failed to build MNN embedder: {}", e))?;

                    let img_1 = run_detection(
                        image1_path,
                        &mut detector,
                        &mut embedder,
                        threshold,
                        nms_threshold,
                        2,
                    )?;
                    let img_2 = run_detection(
                        image2_path,
                        &mut detector,
                        &mut embedder,
                        threshold,
                        nms_threshold,
                        2,
                    )?;

                    let image1_rois = img_1.rois;
                    let image2_rois = img_2.rois;
                    let image1_bbox_len = img_1.bbox.len();
                    let image2_bbox_len = img_2.bbox.len();
                    let best_similarity = compare_faces(&img_1.embeddings, &img_2.embeddings)?;

                    (
                        image1_bbox_len,
                        image2_bbox_len,
                        image1_rois,
                        image2_rois,
                        best_similarity,
                    )
                }
                ExecutorType::OnnxCpu => unimplemented!(),
            };

        Ok((
            image1_faces,
            image2_faces,
            image1_rois,
            image2_rois,
            best_similarity,
        ))
    }
}

use crate::errors::Error;
pub fn compare_faces(
    faces_1: &[Array1<f32>],
    faces_2: &[Array1<f32>],
) -> Result<f32, error_stack::Report<crate::errors::Error>> {
    use error_stack::Report;

    if faces_1.is_empty() || faces_2.is_empty() {
        Err(Report::new(crate::errors::Error))
            .attach_printable("One or both images have no detected faces")?;
    }
    if faces_1.len() != faces_2.len() {
        Err(Report::new(crate::errors::Error))
            .attach_printable("Face count mismatch between images")?;
    }
    Ok(faces_1
        .iter()
        .zip(faces_2)
        .flat_map(|(face_1, face_2)| face_1.cosine_similarity(face_2))
        .inspect(|v| tracing::info!("Cosine similarity: {}", v))
        .map(|v| ordered_float::OrderedFloat(v))
        .max()
        .map(|v| v.0)
        .ok_or(Report::new(Error))?)
}

#[derive(Debug)]
pub struct DetectionOutput {
    bbox: Vec<Aabb2<usize>>,
    rois: Vec<ndarray::Array3<u8>>,
    embeddings: Vec<Array1<f32>>,
}

fn run_detection<D, E>(
    image: impl AsRef<std::path::Path>,
    retinaface: &mut D,
    facenet: &mut E,
    threshold: f32,
    nms_threshold: f32,
    chunk_size: usize,
) -> crate::errors::Result<DetectionOutput>
where
    D: crate::facedet::FaceDetector,
    E: crate::faceembed::FaceEmbedder,
{
    use errors::*;
    // Initialize database if requested
    let image = image.as_ref();
    let image = image::open(image)
        .change_context(Error)
        .attach_printable(image.to_string_lossy().to_string())?;
    let image = image.into_rgb8();
    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(threshold)
                .with_nms_threshold(nms_threshold),
        )
        .change_context(errors::Error)
        .attach_printable("Failed to detect faces")?;

    let bboxes = output
        .bbox
        .iter()
        .inspect(|bbox| tracing::info!("Raw bbox: {:?}", bbox))
        .map(|bbox| bbox.as_::<f32>().scale_uniform(1.30).as_::<usize>())
        .inspect(|bbox| tracing::info!("Padded bbox: {:?}", bbox))
        .collect_vec();
    for bbox in &bboxes {
        tracing::info!("Detected face: {:?}", bbox);
        use bounding_box::draw::*;
        array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 1);
    }
    use itertools::Itertools;
    let face_rois = array
        .view()
        .multi_roi(&bboxes)
        .change_context(Error)?
        .into_iter()
        .map(|roi| {
            roi.as_standard_layout()
                .fast_resize(224, 224, &ResizeOptions::default())
                .change_context(Error)
        })
        .collect::<Result<Vec<_>>>()?;
    let face_roi_views = face_rois.iter().map(|roi| roi.view()).collect::<Vec<_>>();

    let embeddings: Vec<Array1<f32>> = face_roi_views
        .chunks(chunk_size)
        .map(|chunk| {
            tracing::info!("Processing chunk of size: {}", chunk.len());

            let og_size = chunk.len();
            if chunk.len() < chunk_size {
                tracing::warn!("Chunk size is less than 8, padding with zeros");
                let zeros = Array3::zeros((224, 224, 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(ndarray::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, og_size))
            } else {
                let face_rois: Array4<u8> = ndarray::stack(ndarray::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, og_size))
            }
        })
        .collect::<Result<Vec<(Array2<f32>, usize)>>>()?
        .into_iter()
        .map(|(chunk, size): (Array2<f32>, usize)| {
            use itertools::Itertools;
            chunk
                .rows()
                .into_iter()
                .take(size)
                .map(|row| row.to_owned())
                .collect_vec()
                .into_iter()
        })
        .flatten()
        .collect::<Vec<Array1<f32>>>();

    Ok(DetectionOutput {
        bbox: bboxes,
        rois: face_rois,
        embeddings,
    })
}
