chromaframe-sdk 0.1.0

use crate::color::srgb_to_lab;
use crate::image::{
    ManualRegionSamples, MeasureError, MeasurementEngine, MeasurementInput, MeasurementReport,
};
use crate::quality::{NormalizedImage, capture_quality_report, decode_image};
use crate::types::{CandidateColor, GoalVector};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fmt;
use thiserror::Error;

#[derive(
    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum RegionKind {
    Skin,
    Brow,
    Iris,
    Sclera,
    Lip,
    Hair,
    Beard,
    Clothing,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum RegionSource {
    MediapipeLandmarks,
    ParserMask,
    Manual,
    Approximation,
    NotMeasured,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum RegionStatus {
    Measured,
    Approximate,
    LowEvidence,
    NotMeasured { reason: String },
}

#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "encoding", rename_all = "snake_case")]
pub enum RegionMask {
    Polygon {
        polygons: Vec<Vec<Point>>,
    },
    RleBitmap {
        width: u32,
        height: u32,
        counts: Vec<u32>,
    },
}

impl fmt::Debug for RegionMask {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Polygon { polygons } => formatter
                .debug_struct("RegionMask::Polygon")
                .field("polygon_count", &polygons.len())
                .finish(),
            Self::RleBitmap {
                width,
                height,
                counts,
            } => formatter
                .debug_struct("RegionMask::RleBitmap")
                .field("dimensions", &(*width, *height))
                .field("run_count", &counts.len())
                .finish(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct Point {
    pub x: f32,
    pub y: f32,
}

#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct RegionObservation {
    pub kind: RegionKind,
    pub status: RegionStatus,
    pub source: RegionSource,
    pub confidence: f32,
    pub mask: Option<RegionMask>,
    pub sample_hint: Option<usize>,
    pub approximate_reason: Option<String>,
    pub not_measured_reason: Option<String>,
}

impl fmt::Debug for RegionObservation {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RegionObservation")
            .field("kind", &self.kind)
            .field("status", &self.status)
            .field("source", &self.source)
            .field("confidence", &self.confidence)
            .field("mask", &self.mask)
            .field("sample_hint", &self.sample_hint)
            .finish()
    }
}

#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct RegionSet {
    pub image_width: u32,
    pub image_height: u32,
    pub regions: Vec<RegionObservation>,
    pub extraction_quality: ExtractionQualityReport,
    pub warnings: Vec<String>,
}

impl fmt::Debug for RegionSet {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RegionSet")
            .field("dimensions", &(self.image_width, self.image_height))
            .field("region_count", &self.regions.len())
            .field("warning_count", &self.warnings.len())
            .finish()
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct ExtractionQualityReport {
    pub faces_detected: u32,
    pub selected_face_index: Option<u32>,
    pub backend: String,
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct VisionReadinessReport {
    pub backend_available: bool,
    pub python_version: Option<String>,
    pub missing_packages: Vec<String>,
    pub missing_models: Vec<String>,
    pub warnings: Vec<String>,
}

#[derive(Clone, Serialize, Deserialize, JsonSchema)]
pub struct ExtractionRequest {
    pub requested_regions: Vec<RegionKind>,
    pub min_region_pixels: usize,
    pub erode_pixels: u32,
    pub max_faces: u32,
    pub overlay_request: Option<crate::vision::OverlayRequest>,
    pub keep_debug_artifacts: bool,
}

impl Default for ExtractionRequest {
    fn default() -> Self {
        Self {
            requested_regions: vec![
                RegionKind::Skin,
                RegionKind::Brow,
                RegionKind::Iris,
                RegionKind::Sclera,
                RegionKind::Lip,
                RegionKind::Hair,
                RegionKind::Beard,
            ],
            min_region_pixels: 20,
            erode_pixels: 2,
            max_faces: 1,
            overlay_request: None,
            keep_debug_artifacts: false,
        }
    }
}

impl fmt::Debug for ExtractionRequest {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ExtractionRequest")
            .field("requested_regions", &self.requested_regions)
            .field("min_region_pixels", &self.min_region_pixels)
            .field("erode_pixels", &self.erode_pixels)
            .field("max_faces", &self.max_faces)
            .field(
                "overlay_request",
                &self.overlay_request.as_ref().map(|_| "[PRESENT_REDACTED]"),
            )
            .field("keep_debug_artifacts", &self.keep_debug_artifacts)
            .finish()
    }
}

#[derive(Debug, Error)]
pub enum VisionError {
    #[error("vision backend unavailable")]
    BackendUnavailable,
    #[error("vision package missing: {0}")]
    PackageMissing(String),
    #[error("vision model missing: {0}")]
    ModelMissing(String),
    #[error("vision model checksum mismatch")]
    ModelChecksumMismatch,
    #[error("vision model license unsupported")]
    UnsupportedModelLicense,
    #[error("vision backend version mismatch")]
    BackendVersionMismatch,
    #[error("vision helper io error")]
    HelperIo,
    #[error("vision helper protocol error: {0}")]
    HelperProtocolError(String),
    #[error("vision helper timed out")]
    HelperTimeout,
    #[error("region mask is invalid: {0}")]
    InvalidMask(String),
    #[error("region has insufficient samples: {kind:?}")]
    InsufficientSamples { kind: RegionKind },
    #[error(transparent)]
    Image(#[from] crate::quality::ImageError),
    #[error(transparent)]
    Measure(#[from] MeasureError),
}

#[async_trait]
pub trait RegionExtractor: Send + Sync {
    async fn extract_regions(
        &self,
        image: &NormalizedImage,
        request: &ExtractionRequest,
    ) -> Result<RegionSet, VisionError>;
}

pub trait DebugOverlayRenderer: Send + Sync {
    fn render_overlay(
        &self,
        image: &NormalizedImage,
        regions: &RegionSet,
        request: &crate::vision::OverlayRequest,
    ) -> Result<crate::vision::OverlayArtifactSummary, VisionError>;
}

impl RegionSet {
    pub fn parse(self) -> Result<Self, VisionError> {
        if self.image_width == 0 || self.image_height == 0 {
            return Err(VisionError::InvalidMask(
                "image dimensions must be non-zero".to_string(),
            ));
        }
        for region in &self.regions {
            validate_region(region, self.image_width, self.image_height)?;
        }
        Ok(self)
    }

    pub fn to_manual_region_samples(
        &self,
        image: &NormalizedImage,
    ) -> Result<ManualRegionSamples, VisionError> {
        if self.image_width != image.width || self.image_height != image.height {
            return Err(VisionError::InvalidMask(
                "region/image dimension mismatch".to_string(),
            ));
        }
        let mut grouped: BTreeMap<RegionKind, Vec<[u8; 3]>> = BTreeMap::new();
        for region in &self.regions {
            if !matches!(
                region.status,
                RegionStatus::Measured | RegionStatus::Approximate
            ) {
                continue;
            }
            let Some(mask) = &region.mask else {
                continue;
            };
            let samples = sample_mask(image, mask)?;
            let minimum = minimum_samples(region.kind);
            if samples.len() < minimum {
                return Err(VisionError::InsufficientSamples { kind: region.kind });
            }
            grouped.entry(region.kind).or_default().extend(samples);
        }
        Ok(ManualRegionSamples {
            skin: grouped.remove(&RegionKind::Skin),
            brow: grouped.remove(&RegionKind::Brow),
            iris: grouped.remove(&RegionKind::Iris),
            sclera: grouped.remove(&RegionKind::Sclera),
            lip: grouped.remove(&RegionKind::Lip),
            hair: grouped.remove(&RegionKind::Hair),
            beard: grouped.remove(&RegionKind::Beard),
            clothing: grouped.remove(&RegionKind::Clothing),
        })
    }
}

pub struct VisionMeasurementEngine<E> {
    pub extractor: E,
}

impl<E> VisionMeasurementEngine<E> {
    pub fn new(extractor: E) -> Self {
        Self { extractor }
    }
}

impl<E: RegionExtractor> VisionMeasurementEngine<E> {
    pub async fn analyze_bytes(
        &self,
        image_bytes: &[u8],
        extraction: &ExtractionRequest,
        goal_vector: GoalVector,
        candidates: Vec<CandidateColor>,
    ) -> Result<MeasurementReport, VisionError> {
        let image = decode_image(image_bytes, false)?;
        let quality = capture_quality_report(&image)?;
        let regions = self
            .extractor
            .extract_regions(&image, extraction)
            .await?
            .parse()?;
        let samples = regions.to_manual_region_samples(&image)?;
        Ok(MeasurementEngine::measure(&MeasurementInput {
            quality,
            mode: image.measurement_mode,
            goal_vector,
            candidates,
            samples,
        })?)
    }
}

fn validate_region(region: &RegionObservation, width: u32, height: u32) -> Result<(), VisionError> {
    if !region.confidence.is_finite() || !(0.0..=1.0).contains(&region.confidence) {
        return Err(VisionError::InvalidMask(
            "confidence must be finite in [0,1]".to_string(),
        ));
    }
    match (&region.status, &region.mask) {
        (
            RegionStatus::Measured | RegionStatus::Approximate | RegionStatus::LowEvidence,
            Some(mask),
        ) => validate_mask(mask, width, height),
        (RegionStatus::NotMeasured { .. }, _) => Ok(()),
        (_, None) => Err(VisionError::InvalidMask(
            "measured region requires mask".to_string(),
        )),
    }
}

pub fn validate_mask(mask: &RegionMask, width: u32, height: u32) -> Result<(), VisionError> {
    match mask {
        RegionMask::Polygon { polygons } => {
            if polygons.is_empty() {
                return Err(VisionError::InvalidMask(
                    "polygon list is empty".to_string(),
                ));
            }
            for polygon in polygons {
                if polygon.len() < 3 {
                    return Err(VisionError::InvalidMask(
                        "polygon needs at least three points".to_string(),
                    ));
                }
                for point in polygon {
                    if !point.x.is_finite() || !point.y.is_finite() {
                        return Err(VisionError::InvalidMask("point must be finite".to_string()));
                    }
                    if point.x < 0.0
                        || point.y < 0.0
                        || point.x >= width as f32
                        || point.y >= height as f32
                    {
                        return Err(VisionError::InvalidMask("point out of bounds".to_string()));
                    }
                }
            }
            Ok(())
        }
        RegionMask::RleBitmap {
            width: mask_width,
            height: mask_height,
            counts,
        } => {
            if *mask_width != width || *mask_height != height {
                return Err(VisionError::InvalidMask(
                    "rle dimensions mismatch".to_string(),
                ));
            }
            let total: u64 = counts.iter().map(|count| u64::from(*count)).sum();
            if total != u64::from(width) * u64::from(height) {
                return Err(VisionError::InvalidMask(
                    "rle count total mismatch".to_string(),
                ));
            }
            Ok(())
        }
    }
}

pub fn shrink_polygon(polygon: &[Point], pixels: f32) -> Result<Vec<Point>, VisionError> {
    if polygon.len() < 3 {
        return Err(VisionError::InvalidMask(
            "polygon needs at least three points".to_string(),
        ));
    }
    if !pixels.is_finite() || pixels < 0.0 {
        return Err(VisionError::InvalidMask(
            "shrink pixels must be finite and non-negative".to_string(),
        ));
    }
    let centroid = polygon_centroid(polygon)?;
    Ok(polygon
        .iter()
        .map(|point| shrink_point(*point, centroid, pixels))
        .collect())
}

fn sample_mask(image: &NormalizedImage, mask: &RegionMask) -> Result<Vec<[u8; 3]>, VisionError> {
    match mask {
        RegionMask::Polygon { polygons } => sample_polygons(image, polygons),
        RegionMask::RleBitmap { counts, .. } => sample_rle(image, counts),
    }
}

fn sample_polygons(
    image: &NormalizedImage,
    polygons: &[Vec<Point>],
) -> Result<Vec<[u8; 3]>, VisionError> {
    validate_mask(
        &RegionMask::Polygon {
            polygons: polygons.to_vec(),
        },
        image.width,
        image.height,
    )?;
    let mut samples = Vec::new();
    for y in 0..image.height {
        for x in 0..image.width {
            if !polygons
                .iter()
                .any(|polygon| point_in_polygon(x as f32 + 0.5, y as f32 + 0.5, polygon))
            {
                continue;
            }
            let pixel = image.pixels[(y * image.width + x) as usize];
            if pixel.opaque {
                samples.push(rgb_from_pixel(pixel));
            }
        }
    }
    Ok(samples)
}

fn sample_rle(image: &NormalizedImage, counts: &[u32]) -> Result<Vec<[u8; 3]>, VisionError> {
    let mut samples = Vec::new();
    let mut index = 0usize;
    let mut selected = false;
    for count in counts {
        for _ in 0..*count {
            if selected {
                let pixel = image.pixels.get(index).ok_or_else(|| {
                    VisionError::InvalidMask("rle index out of bounds".to_string())
                })?;
                if pixel.opaque {
                    samples.push(rgb_from_pixel(*pixel));
                }
            }
            index += 1;
        }
        selected = !selected;
    }
    Ok(samples)
}

fn rgb_from_pixel(pixel: crate::quality::LinearPixel) -> [u8; 3] {
    [
        (pixel.r.clamp(0.0, 1.0) * 255.0).round() as u8,
        (pixel.g.clamp(0.0, 1.0) * 255.0).round() as u8,
        (pixel.b.clamp(0.0, 1.0) * 255.0).round() as u8,
    ]
}

fn minimum_samples(kind: RegionKind) -> usize {
    if kind == RegionKind::Skin { 50 } else { 20 }
}

fn polygon_centroid(polygon: &[Point]) -> Result<Point, VisionError> {
    let mut x = 0.0;
    let mut y = 0.0;
    for point in polygon {
        if !point.x.is_finite() || !point.y.is_finite() {
            return Err(VisionError::InvalidMask("point must be finite".to_string()));
        }
        x += point.x;
        y += point.y;
    }
    Ok(Point {
        x: x / polygon.len() as f32,
        y: y / polygon.len() as f32,
    })
}

fn shrink_point(point: Point, centroid: Point, pixels: f32) -> Point {
    let dx = point.x - centroid.x;
    let dy = point.y - centroid.y;
    let distance = (dx * dx + dy * dy).sqrt();
    if distance <= pixels || distance == 0.0 {
        return centroid;
    }
    let scale = (distance - pixels) / distance;
    Point {
        x: centroid.x + dx * scale,
        y: centroid.y + dy * scale,
    }
}

pub fn point_in_polygon(x: f32, y: f32, polygon: &[Point]) -> bool {
    if polygon.len() < 3 {
        return false;
    }
    let mut inside = false;
    let mut previous = polygon.len() - 1;
    for current in 0..polygon.len() {
        let current_point = polygon[current];
        let previous_point = polygon[previous];
        let crosses_scanline = (current_point.y > y) != (previous_point.y > y);
        if !crosses_scanline {
            previous = current;
            continue;
        }
        let intersection_x = (previous_point.x - current_point.x) * (y - current_point.y)
            / (previous_point.y - current_point.y)
            + current_point.x;
        if x < intersection_x {
            inside = !inside;
        }
        previous = current;
    }
    inside
}

#[must_use]
pub fn region_mean_lab(samples: &[[u8; 3]]) -> Option<crate::color::Lab> {
    if samples.is_empty() {
        return None;
    }
    let labs: Vec<_> = samples.iter().map(|rgb| srgb_to_lab(*rgb)).collect();
    let (l, a, b) = labs.iter().fold((0.0, 0.0, 0.0), |acc, lab| {
        (acc.0 + lab.l, acc.1 + lab.a, acc.2 + lab.b)
    });
    Some(crate::color::Lab {
        l: l / labs.len() as f32,
        a: a / labs.len() as f32,
        b: b / labs.len() as f32,
    })
}