//! Lightweight, local subject analysis for artist-hero composition.
//!
//! Person boxes come from the bundled YOLOX-Nano model through the Windows
//! ML runtime already present in Windows 10/11. Faces come from the native
//! Windows face detector, so no second model or inference runtime is shipped.
//! All failures are recoverable: callers fall back to a deterministic crop.

use crate::hero_crop::{FaceBox, SubjectBox};

#[cfg(windows)]
mod windows_detector {
    use std::sync::OnceLock;

    use image::{DynamicImage, GenericImageView as _, ImageEncoder as _};
    use tauri::Manager as _;
    use windows::{
        core::{Interface as _, HSTRING},
        Graphics::Imaging::{BitmapDecoder, BitmapPixelFormat, BitmapSize, SoftwareBitmap},
        Media::FaceAnalysis::FaceDetector,
        Storage::Streams::{DataWriter, InMemoryRandomAccessStream},
        Win32::System::WinRT::{RoInitialize, RoUninitialize, RO_INIT_MULTITHREADED},
        AI::MachineLearning::{
            LearningModel, LearningModelBinding, LearningModelDevice, LearningModelDeviceKind,
            LearningModelSession, TensorFloat,
        },
    };

    use super::{FaceBox, SubjectBox};

    const YOLOX_EDGE: u32 = 416;
    const YOLOX_CANDIDATES: usize = 3549;
    const YOLOX_VALUES: usize = 85;
    const PERSON_SCORE_MIN: f32 = 0.25;
    const PERSON_NMS_IOU: f32 = 0.45;
    const FACE_ANALYSIS_MAX_EDGE: u32 = 960;

    static PERSON_SESSION: OnceLock<Result<LearningModelSession, String>> = OnceLock::new();
    static FACE_DETECTOR: OnceLock<Result<FaceDetector, String>> = OnceLock::new();

    struct WinRtGuard;

    impl Drop for WinRtGuard {
        fn drop(&mut self) {
            unsafe { RoUninitialize() };
        }
    }

    fn initialize_winrt() -> Result<Option<WinRtGuard>, String> {
        match unsafe { RoInitialize(RO_INIT_MULTITHREADED) } {
            Ok(()) => Ok(Some(WinRtGuard)),
            Err(error) if error.code().0 as u32 == 0x8001_0106 => {
                // The current thread already has a different apartment. WinRT
                // remains usable there; it simply must not be uninitialized by us.
                Ok(None)
            }
            Err(error) => Err(format!("Windows Runtime initialization failed: {error}")),
        }
    }

    fn model_path(app: &tauri::AppHandle) -> Result<std::path::PathBuf, String> {
        let bundled = app
            .path()
            .resolve(
                "resources/yolox_nano.onnx",
                tauri::path::BaseDirectory::Resource,
            )
            .map_err(|error| error.to_string())?;
        if bundled.is_file() {
            return Ok(bundled);
        }
        let development = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("resources")
            .join("yolox_nano.onnx");
        if development.is_file() {
            return Ok(development);
        }
        Err("bundled YOLOX-Nano model is missing".into())
    }

    fn person_session_from_path(path: &std::path::Path) -> Result<LearningModelSession, String> {
        PERSON_SESSION
            .get_or_init(|| {
                let model = LearningModel::LoadFromFilePath(&HSTRING::from(
                    path.to_string_lossy().as_ref(),
                ))
                .map_err(|error| format!("could not load YOLOX-Nano: {error}"))?;
                let device = LearningModelDevice::Create(LearningModelDeviceKind::Cpu)
                    .map_err(|error| format!("could not create WinML CPU device: {error}"))?;
                LearningModelSession::CreateFromModelOnDevice(&model, &device)
                    .map_err(|error| format!("could not create YOLOX session: {error}"))
            })
            .clone()
    }

    fn person_session(app: &tauri::AppHandle) -> Result<LearningModelSession, String> {
        person_session_from_path(&model_path(app)?)
    }

    fn face_detector() -> Result<FaceDetector, String> {
        FACE_DETECTOR
            .get_or_init(|| {
                if !FaceDetector::IsSupported()
                    .map_err(|error| format!("could not query face detector: {error}"))?
                {
                    return Err("Windows face detection is unavailable".into());
                }
                let detector = FaceDetector::CreateAsync()
                    .map_err(|error| format!("could not create face detector: {error}"))?
                    .get()
                    .map_err(|error| format!("could not initialize face detector: {error}"))?;
                // Small faces matter in group press photos. Analysis is capped
                // at 960 px, so 18 px still avoids tiny texture false positives.
                let _ = detector.SetMinDetectableFaceSize(BitmapSize {
                    Width: 18,
                    Height: 18,
                });
                Ok(detector)
            })
            .clone()
    }

    pub(super) fn detect(
        app: &tauri::AppHandle,
        image: &DynamicImage,
    ) -> Result<(Vec<SubjectBox>, Vec<FaceBox>), String> {
        let _winrt = initialize_winrt()?;
        let subjects = detect_people(app, image)?;
        // Face refinement is optional. A native-detector failure must not throw
        // away valid person boxes from WinML.
        let faces = match detect_faces(image) {
            Ok(faces) => faces,
            Err(error) => {
                log::warn!("artist hero face refinement unavailable: {error}");
                Vec::new()
            }
        };
        Ok((subjects, faces))
    }

    pub(super) fn warm_up(app: tauri::AppHandle) {
        std::thread::Builder::new()
            .name("hero-vision-warmup".into())
            .spawn(move || {
                let Ok(_winrt) = initialize_winrt() else {
                    return;
                };
                if let Err(error) = person_session(&app) {
                    log::warn!("artist hero person detector warm-up failed: {error}");
                }
                if let Err(error) = face_detector() {
                    log::warn!("artist hero face detector warm-up failed: {error}");
                }
            })
            .ok();
    }

    fn detect_people(
        app: &tauri::AppHandle,
        image: &DynamicImage,
    ) -> Result<Vec<SubjectBox>, String> {
        detect_people_with_model_path(&model_path(app)?, image)
    }

    fn detect_people_with_model_path(
        model_path: &std::path::Path,
        image: &DynamicImage,
    ) -> Result<Vec<SubjectBox>, String> {
        let (source_width, source_height) = image.dimensions();
        if source_width == 0 || source_height == 0 {
            return Ok(Vec::new());
        }

        let scale =
            (YOLOX_EDGE as f32 / source_width as f32).min(YOLOX_EDGE as f32 / source_height as f32);
        let resized_width = ((source_width as f32 * scale).floor() as u32).clamp(1, YOLOX_EDGE);
        let resized_height = ((source_height as f32 * scale).floor() as u32).clamp(1, YOLOX_EDGE);
        let rgb = image.to_rgb8();
        let resized = image::imageops::resize(
            &rgb,
            resized_width,
            resized_height,
            image::imageops::FilterType::Triangle,
        );

        let plane_len = (YOLOX_EDGE * YOLOX_EDGE) as usize;
        let mut tensor_data = vec![114.0f32; plane_len * 3];
        for y in 0..resized_height {
            for x in 0..resized_width {
                let pixel = resized.get_pixel(x, y).0;
                let offset = (y * YOLOX_EDGE + x) as usize;
                // Official YOLOX preprocessing consumes OpenCV BGR without
                // normalization, transposed directly into NCHW planes.
                tensor_data[offset] = f32::from(pixel[2]);
                tensor_data[plane_len + offset] = f32::from(pixel[1]);
                tensor_data[plane_len * 2 + offset] = f32::from(pixel[0]);
            }
        }

        let input = TensorFloat::CreateFromShapeArrayAndDataArray(
            &[1, 3, i64::from(YOLOX_EDGE), i64::from(YOLOX_EDGE)],
            &tensor_data,
        )
        .map_err(|error| format!("could not create YOLOX input: {error}"))?;
        let session = person_session_from_path(model_path)?;
        let binding = LearningModelBinding::CreateFromSession(&session)
            .map_err(|error| format!("could not create YOLOX binding: {error}"))?;
        binding
            .Bind(&HSTRING::from("images"), &input)
            .map_err(|error| format!("could not bind YOLOX input: {error}"))?;
        let result = session
            .Evaluate(&binding, &HSTRING::from("artist-hero"))
            .map_err(|error| format!("YOLOX inference failed: {error}"))?;
        if !result
            .Succeeded()
            .map_err(|error| format!("could not read YOLOX status: {error}"))?
        {
            return Err(format!(
                "YOLOX evaluation failed with status {}",
                result.ErrorStatus().unwrap_or_default()
            ));
        }
        let output = result
            .Outputs()
            .and_then(|outputs| outputs.Lookup(&HSTRING::from("output")))
            .map_err(|error| format!("YOLOX output is missing: {error}"))?;
        let tensor: TensorFloat = output
            .cast()
            .map_err(|error| format!("YOLOX output has the wrong type: {error}"))?;
        let view = tensor
            .GetAsVectorView()
            .map_err(|error| format!("could not read YOLOX output: {error}"))?;
        let values: Vec<f32> = view.into_iter().collect();
        if values.len() != YOLOX_CANDIDATES * YOLOX_VALUES {
            return Err(format!("unexpected YOLOX output length: {}", values.len()));
        }

        let mut candidates = Vec::new();
        let mut candidate_index = 0usize;
        for stride in [8u32, 16, 32] {
            let grid = YOLOX_EDGE / stride;
            for gy in 0..grid {
                for gx in 0..grid {
                    let row = &values
                        [candidate_index * YOLOX_VALUES..(candidate_index + 1) * YOLOX_VALUES];
                    candidate_index += 1;
                    let confidence = row[4] * row[5];
                    if !confidence.is_finite() || confidence < PERSON_SCORE_MIN {
                        continue;
                    }
                    let center_x = (row[0] + gx as f32) * stride as f32 / scale;
                    let center_y = (row[1] + gy as f32) * stride as f32 / scale;
                    let width = row[2].clamp(-10.0, 10.0).exp() * stride as f32 / scale;
                    let height = row[3].clamp(-10.0, 10.0).exp() * stride as f32 / scale;
                    if !center_x.is_finite()
                        || !center_y.is_finite()
                        || !width.is_finite()
                        || !height.is_finite()
                    {
                        continue;
                    }
                    let left = (center_x - width * 0.5).clamp(0.0, source_width as f32);
                    let top = (center_y - height * 0.5).clamp(0.0, source_height as f32);
                    let right = (center_x + width * 0.5).clamp(0.0, source_width as f32);
                    let bottom = (center_y + height * 0.5).clamp(0.0, source_height as f32);
                    if right - left < 8.0 || bottom - top < 16.0 {
                        continue;
                    }
                    candidates.push(SubjectBox {
                        x: left,
                        y: top,
                        width: right - left,
                        height: bottom - top,
                        confidence,
                    });
                }
            }
        }
        debug_assert_eq!(candidate_index, YOLOX_CANDIDATES);
        Ok(non_maximum_suppression(candidates, PERSON_NMS_IOU, 24))
    }

    fn detect_faces(image: &DynamicImage) -> Result<Vec<FaceBox>, String> {
        let (source_width, source_height) = image.dimensions();
        if source_width == 0 || source_height == 0 {
            return Ok(Vec::new());
        }
        let detection = image.thumbnail(FACE_ANALYSIS_MAX_EDGE, FACE_ANALYSIS_MAX_EDGE);
        let gray = detection.to_luma8();
        let (width, height) = gray.dimensions();
        let mut png = Vec::new();
        image::codecs::png::PngEncoder::new(&mut png)
            .write_image(gray.as_raw(), width, height, image::ExtendedColorType::L8)
            .map_err(|error| format!("could not encode face-analysis raster: {error}"))?;

        let stream = InMemoryRandomAccessStream::new()
            .map_err(|error| format!("could not allocate face-analysis stream: {error}"))?;
        let writer = DataWriter::CreateDataWriter(&stream)
            .map_err(|error| format!("could not create face-analysis writer: {error}"))?;
        writer
            .WriteBytes(&png)
            .map_err(|error| format!("could not write face-analysis raster: {error}"))?;
        writer
            .StoreAsync()
            .and_then(|operation| operation.get())
            .map_err(|error| format!("could not store face-analysis raster: {error}"))?;
        writer
            .FlushAsync()
            .and_then(|operation| operation.get())
            .map_err(|error| format!("could not flush face-analysis raster: {error}"))?;
        let _ = writer.DetachStream();
        stream
            .Seek(0)
            .map_err(|error| format!("could not rewind face-analysis raster: {error}"))?;
        let decoder = BitmapDecoder::CreateAsync(&stream)
            .and_then(|operation| operation.get())
            .map_err(|error| format!("could not decode face-analysis raster: {error}"))?;
        let decoded_bitmap = decoder
            .GetSoftwareBitmapAsync()
            .and_then(|operation| operation.get())
            .map_err(|error| format!("could not prepare face-analysis bitmap: {error}"))?;
        let bitmap = SoftwareBitmap::Convert(&decoded_bitmap, BitmapPixelFormat::Gray8)
            .map_err(|error| format!("could not convert face-analysis bitmap: {error}"))?;
        let found = face_detector()?
            .DetectFacesAsync(&bitmap)
            .and_then(|operation| operation.get())
            .map_err(|error| format!("face detection failed: {error}"))?;

        let scale_x = source_width as f32 / width as f32;
        let scale_y = source_height as f32 / height as f32;
        let mut faces = Vec::new();
        for detected in &found {
            let bounds = detected
                .FaceBox()
                .map_err(|error| format!("could not read face bounds: {error}"))?;
            faces.push(FaceBox {
                x: bounds.X as f32 * scale_x,
                y: bounds.Y as f32 * scale_y,
                width: bounds.Width as f32 * scale_x,
                height: bounds.Height as f32 * scale_y,
                confidence: 1.0,
            });
        }
        Ok(faces)
    }

    fn non_maximum_suppression(
        mut candidates: Vec<SubjectBox>,
        threshold: f32,
        limit: usize,
    ) -> Vec<SubjectBox> {
        candidates.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
        let mut kept = Vec::new();
        'candidate: for candidate in candidates {
            for accepted in &kept {
                if intersection_over_union(candidate, *accepted) > threshold {
                    continue 'candidate;
                }
            }
            kept.push(candidate);
            if kept.len() == limit {
                break;
            }
        }
        kept
    }

    fn intersection_over_union(a: SubjectBox, b: SubjectBox) -> f32 {
        let left = a.x.max(b.x);
        let top = a.y.max(b.y);
        let right = (a.x + a.width).min(b.x + b.width);
        let bottom = (a.y + a.height).min(b.y + b.height);
        let intersection = (right - left).max(0.0) * (bottom - top).max(0.0);
        let union = a.width * a.height + b.width * b.height - intersection;
        if union > 0.0 {
            intersection / union
        } else {
            0.0
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        #[test]
        fn nms_keeps_the_best_overlapping_person() {
            let kept = non_maximum_suppression(
                vec![
                    SubjectBox {
                        x: 0.0,
                        y: 0.0,
                        width: 100.0,
                        height: 200.0,
                        confidence: 0.9,
                    },
                    SubjectBox {
                        x: 2.0,
                        y: 4.0,
                        width: 100.0,
                        height: 200.0,
                        confidence: 0.8,
                    },
                    SubjectBox {
                        x: 300.0,
                        y: 0.0,
                        width: 100.0,
                        height: 200.0,
                        confidence: 0.7,
                    },
                ],
                0.45,
                24,
            );
            assert_eq!(kept.len(), 2);
            assert_eq!(kept[0].confidence, 0.9);
        }

        /// Manual, reusable integration probe. It is ignored in normal CI;
        /// provide an image through SOULSTREAM_HERO_FIXTURE when validating a
        /// new detector/model or an especially difficult artist photo.
        #[test]
        #[ignore]
        fn fixture_detection_and_crop_probe() {
            let fixture = std::env::var("SOULSTREAM_HERO_FIXTURE")
                .expect("set SOULSTREAM_HERO_FIXTURE to a source image");
            let image = image::open(fixture).expect("fixture must decode");
            let _winrt = initialize_winrt().expect("WinRT must initialize");
            let model = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
                .join("resources")
                .join("yolox_nano.onnx");
            let subjects = detect_people_with_model_path(&model, &image)
                .expect("person inference must succeed");
            let faces = detect_faces(&image).expect("face detection must succeed");
            let plan =
                crate::hero_crop::plan_hero_crop(image.width(), image.height(), &subjects, &faces)
                    .expect("fixture must plan");
            eprintln!("subjects={subjects:#?}");
            eprintln!("faces={faces:#?}");
            eprintln!("plan={plan:#?}");
            assert!(plan.crop.width > 0 && plan.crop.height > 0);
        }
    }
}

/// Detects people and faces in source-image coordinates. This is deliberately
/// a fallible enhancement: hero processing must continue when vision APIs are
/// unavailable (including non-Windows development builds).
pub(crate) fn detect_subjects(
    app: &tauri::AppHandle,
    image: &image::DynamicImage,
) -> Result<(Vec<SubjectBox>, Vec<FaceBox>), String> {
    #[cfg(windows)]
    {
        windows_detector::detect(app, image)
    }
    #[cfg(not(windows))]
    {
        let _ = (app, image);
        Err("content-aware hero detection is only available on Windows".into())
    }
}

pub(crate) fn warm_up(app: tauri::AppHandle) {
    #[cfg(windows)]
    windows_detector::warm_up(app);
    #[cfg(not(windows))]
    let _ = app;
}