chromaframe-mcp 0.1.1

use std::{borrow::Cow, path::PathBuf};

use chromaframe_sdk::{GoalVector, Lab};
use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
use serde::{Deserialize, Deserializer, Serialize};

pub const DEFAULT_CONFIDENCE: f32 = 0.65;
pub const DEFAULT_RANKING_LIMIT: usize = 5;
pub const MAX_RANKING_LIMIT: usize = 25;
pub const MAX_CANDIDATE_COUNT: usize = 64;
pub const MAX_CANDIDATE_NAME_BYTES: usize = 80;
pub const MAX_PATH_BYTES: usize = 4096;
pub const MAX_ENCODED_IMAGE_BYTES: u64 = 64 * 1024 * 1024;

#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ReadinessInput {
    #[serde(default)]
    pub include_warnings: bool,
}

impl JsonSchema for ReadinessInput {
    fn schema_name() -> Cow<'static, str> {
        "ReadinessInput".into()
    }

    fn schema_id() -> Cow<'static, str> {
        concat!(module_path!(), "::ReadinessInput").into()
    }

    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
        json_schema!({
            "type": "object",
            "additionalProperties": false,
            "properties": {
                "include_warnings": { "type": "boolean", "default": false }
            }
        })
    }
}

#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ManualRankInput {
    pub skin_lab: LabInput,

    #[serde(default)]
    pub brow_lab: Option<LabInput>,

    #[serde(default)]
    pub iris_lab: Option<LabInput>,

    #[serde(default)]
    pub sclera_lab: Option<LabInput>,

    #[serde(default)]
    pub lip_lab: Option<LabInput>,

    #[serde(default)]
    pub hair_lab: Option<LabInput>,

    #[serde(default)]
    pub beard_lab: Option<LabInput>,

    #[serde(default)]
    pub goal_vector: GoalVectorInput,

    #[serde(default = "default_confidence")]
    pub confidence: f32,

    #[serde(default = "default_ranking_limit")]
    pub limit: usize,

    pub candidates: Vec<CandidateInput>,
}

impl JsonSchema for ManualRankInput {
    fn schema_name() -> Cow<'static, str> {
        "ManualRankInput".into()
    }

    fn schema_id() -> Cow<'static, str> {
        concat!(module_path!(), "::ManualRankInput").into()
    }

    fn json_schema(generator: &mut SchemaGenerator) -> Schema {
        let lab_schema = generator.subschema_for::<LabInput>();
        let candidate_schema = generator.subschema_for::<CandidateInput>();
        let goal_schema = generator.subschema_for::<GoalVectorInput>();
        json_schema!({
            "type": "object",
            "additionalProperties": false,
            "required": ["skin_lab", "candidates"],
            "properties": {
                "skin_lab": lab_schema,
                "brow_lab": lab_schema,
                "iris_lab": lab_schema,
                "sclera_lab": lab_schema,
                "lip_lab": lab_schema,
                "hair_lab": lab_schema,
                "beard_lab": lab_schema,
                "goal_vector": goal_schema,
                "confidence": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": DEFAULT_CONFIDENCE },
                "limit": { "type": "integer", "minimum": 1, "maximum": MAX_RANKING_LIMIT, "default": DEFAULT_RANKING_LIMIT },
                "candidates": { "type": "array", "minItems": 1, "maxItems": MAX_CANDIDATE_COUNT, "items": candidate_schema }
            }
        })
    }
}

#[derive(Debug, Clone)]
pub struct AnalyzeImageInput {
    pub image_path: PathBuf,
    pub goal_vector: GoalVectorInput,
    pub limit: usize,
    pub candidates: Vec<CandidateInput>,
}

impl<'de> Deserialize<'de> for AnalyzeImageInput {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(deny_unknown_fields)]
        struct RawAnalyzeImageInput {
            image_path: String,
            #[serde(default)]
            goal_vector: GoalVectorInput,
            #[serde(default = "default_ranking_limit")]
            limit: usize,
            candidates: Vec<CandidateInput>,
        }

        let raw = RawAnalyzeImageInput::deserialize(deserializer)?;
        let image_path = parse_non_blank_path(raw.image_path, "image_path")?;

        Ok(Self {
            image_path,
            goal_vector: raw.goal_vector,
            limit: raw.limit,
            candidates: raw.candidates,
        })
    }
}

impl JsonSchema for AnalyzeImageInput {
    fn schema_name() -> Cow<'static, str> {
        "AnalyzeImageInput".into()
    }

    fn schema_id() -> Cow<'static, str> {
        concat!(module_path!(), "::AnalyzeImageInput").into()
    }

    fn json_schema(generator: &mut SchemaGenerator) -> Schema {
        let candidate_schema = generator.subschema_for::<CandidateInput>();
        let goal_schema = generator.subschema_for::<GoalVectorInput>();
        json_schema!({
            "type": "object",
            "additionalProperties": false,
            "required": ["image_path", "candidates"],
            "properties": {
                "image_path": { "type": "string", "minLength": 1, "maxLength": MAX_PATH_BYTES, "pattern": ".*\\S.*" },
                "goal_vector": goal_schema,
                "limit": { "type": "integer", "minimum": 1, "maximum": MAX_RANKING_LIMIT, "default": DEFAULT_RANKING_LIMIT },
                "candidates": { "type": "array", "minItems": 1, "maxItems": MAX_CANDIDATE_COUNT, "items": candidate_schema }
            }
        })
    }
}

#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct LabInput {
    pub l: f32,
    pub a: f32,
    pub b: f32,
}

impl From<LabInput> for Lab {
    fn from(value: LabInput) -> Self {
        Self {
            l: value.l,
            a: value.a,
            b: value.b,
        }
    }
}

impl From<Lab> for LabInput {
    fn from(value: Lab) -> Self {
        Self {
            l: value.l,
            a: value.a,
            b: value.b,
        }
    }
}

#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct GoalVectorInput {
    #[serde(default = "default_contrast_target")]
    pub contrast_target: f32,
    #[serde(default)]
    pub warmth_target: f32,
    #[serde(default = "default_chroma_target")]
    pub chroma_target: f32,
    #[serde(default = "default_feature_readability_target")]
    pub feature_readability_target: f32,
    #[serde(default = "default_artificiality_tolerance")]
    pub artificiality_tolerance: f32,
}

impl JsonSchema for GoalVectorInput {
    fn schema_name() -> Cow<'static, str> {
        "GoalVectorInput".into()
    }

    fn schema_id() -> Cow<'static, str> {
        concat!(module_path!(), "::GoalVectorInput").into()
    }

    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
        json_schema!({
            "type": "object",
            "additionalProperties": false,
            "properties": {
                "contrast_target": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": default_contrast_target() },
                "warmth_target": { "type": "number", "minimum": -1.0, "maximum": 1.0, "default": 0.0 },
                "chroma_target": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": default_chroma_target() },
                "feature_readability_target": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": default_feature_readability_target() },
                "artificiality_tolerance": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": default_artificiality_tolerance() }
            }
        })
    }
}

impl Default for GoalVectorInput {
    fn default() -> Self {
        let goal = GoalVector::default();
        Self {
            contrast_target: goal.contrast_target,
            warmth_target: goal.warmth_target,
            chroma_target: goal.chroma_target,
            feature_readability_target: goal.feature_readability_target,
            artificiality_tolerance: goal.artificiality_tolerance,
        }
    }
}

impl From<GoalVectorInput> for GoalVector {
    fn from(value: GoalVectorInput) -> Self {
        Self {
            contrast_target: value.contrast_target,
            warmth_target: value.warmth_target,
            chroma_target: value.chroma_target,
            feature_readability_target: value.feature_readability_target,
            artificiality_tolerance: value.artificiality_tolerance,
        }
    }
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct CandidateInput {
    pub name: String,
    pub srgb: [u8; 3],
}

impl<'de> Deserialize<'de> for CandidateInput {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(deny_unknown_fields)]
        struct RawCandidateInput {
            name: String,
            srgb: [u8; 3],
        }

        let raw = RawCandidateInput::deserialize(deserializer)?;
        let name = parse_candidate_name::<D::Error>(raw.name)?;
        Ok(Self {
            name,
            srgb: raw.srgb,
        })
    }
}

impl JsonSchema for CandidateInput {
    fn schema_name() -> Cow<'static, str> {
        "CandidateInput".into()
    }

    fn schema_id() -> Cow<'static, str> {
        concat!(module_path!(), "::CandidateInput").into()
    }

    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
        json_schema!({
            "type": "object",
            "additionalProperties": false,
            "required": ["name", "srgb"],
            "properties": {
                "name": { "type": "string", "minLength": 1, "maxLength": MAX_CANDIDATE_NAME_BYTES, "pattern": ".*\\S.*" },
                "srgb": {
                    "type": "array",
                    "minItems": 3,
                    "maxItems": 3,
                    "items": { "type": "integer", "minimum": 0, "maximum": 255 }
                }
            }
        })
    }
}

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

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct ManualRankOutput {
    pub status: String,
    pub subject: SubjectSummary,
    pub rankings: Vec<CandidateRankingSummary>,
    pub uncertainty: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct AnalyzeImageOutput {
    pub image: ImageSummary,
    pub extraction: ExtractionSummary,
    pub measurement: MeasurementSummary,
    pub rankings: Vec<CandidateRankingSummary>,
    pub evidence_cards: Vec<EvidenceCardSummary>,
    pub warnings: Vec<String>,
    pub limitations: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct ImageSummary {
    pub dimensions: [u32; 2],
    pub measurement_mode: String,
    pub icc_status: String,
    pub metadata_retained: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct ExtractionSummary {
    pub backend: String,
    pub faces_detected: u32,
    pub selected_face_index: Option<u32>,
    pub regions: Vec<RegionSummary>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct RegionSummary {
    pub kind: String,
    pub status: String,
    pub source: String,
    pub confidence: f32,
    pub sample_hint: Option<usize>,
    pub reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct MeasurementSummary {
    pub status: String,
    pub confidence: Option<f32>,
    pub subject: Option<SubjectSummary>,
    pub contrasts: Vec<ContrastSummary>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct SubjectSummary {
    pub skin_lab: LabInput,
    pub skin_ita: f32,
    pub skin_depth_proxy: f32,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct ContrastSummary {
    pub from: String,
    pub to: String,
    pub delta_l: f32,
    pub delta_a: f32,
    pub delta_b: f32,
    pub delta_e00: f32,
    pub michelson_lightness_contrast: f32,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct CandidateRankingSummary {
    pub name: String,
    pub score: f32,
    pub confidence: f32,
    pub harshness: f32,
    pub label: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct EvidenceCardSummary {
    pub kind: String,
    pub summary: String,
}

fn parse_non_blank_path<E>(raw_path: String, field: &'static str) -> Result<PathBuf, E>
where
    E: serde::de::Error,
{
    let trimmed = raw_path.trim();
    if trimmed.is_empty() {
        return Err(E::custom(format!("`{field}` cannot be blank")));
    }
    if trimmed.len() > MAX_PATH_BYTES {
        return Err(E::custom(format!(
            "`{field}` exceeds {MAX_PATH_BYTES} UTF-8 bytes"
        )));
    }

    Ok(PathBuf::from(trimmed))
}

fn parse_candidate_name<E>(raw_name: String) -> Result<String, E>
where
    E: serde::de::Error,
{
    let trimmed = raw_name.trim();
    if trimmed.is_empty() {
        return Err(E::custom("candidate `name` cannot be blank"));
    }
    if trimmed.len() > MAX_CANDIDATE_NAME_BYTES {
        return Err(E::custom(format!(
            "candidate `name` exceeds {MAX_CANDIDATE_NAME_BYTES} UTF-8 bytes"
        )));
    }

    Ok(trimmed.to_string())
}

const fn default_confidence() -> f32 {
    DEFAULT_CONFIDENCE
}

const fn default_ranking_limit() -> usize {
    DEFAULT_RANKING_LIMIT
}

fn default_contrast_target() -> f32 {
    GoalVector::default().contrast_target
}

fn default_chroma_target() -> f32 {
    GoalVector::default().chroma_target
}

fn default_feature_readability_target() -> f32 {
    GoalVector::default().feature_readability_target
}

fn default_artificiality_tolerance() -> f32 {
    GoalVector::default().artificiality_tolerance
}

pub fn validate_limit(limit: usize) -> Result<usize, String> {
    if !(1..=MAX_RANKING_LIMIT).contains(&limit) {
        return Err(format!("limit must be between 1 and {MAX_RANKING_LIMIT}"));
    }

    Ok(limit)
}

pub fn validate_candidate_count(count: usize) -> Result<usize, String> {
    if !(1..=MAX_CANDIDATE_COUNT).contains(&count) {
        return Err(format!(
            "candidate count must be between 1 and {MAX_CANDIDATE_COUNT}"
        ));
    }

    Ok(count)
}