use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum PbrTextureRole {
    BaseColor,
    Normal,
    Roughness,
    Metallic,
    Occlusion,
    PackedOrm,
    Height,
    Unknown,
}

impl PbrTextureRole {
    pub const ALL: [Self; 8] = [
        Self::BaseColor,
        Self::Normal,
        Self::Roughness,
        Self::Metallic,
        Self::Occlusion,
        Self::PackedOrm,
        Self::Height,
        Self::Unknown,
    ];

    pub const fn label(self) -> &'static str {
        match self {
            Self::BaseColor => "Base color",
            Self::Normal => "Normal",
            Self::Roughness => "Roughness",
            Self::Metallic => "Metallic",
            Self::Occlusion => "Occlusion",
            Self::PackedOrm => "Packed ORM/ARM",
            Self::Height => "Height",
            Self::Unknown => "Unresolved",
        }
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum PbrMatchConfidence {
    High,
    Medium,
    Low,
    Unresolved,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct DetectedPbrTexture {
    pub path: PathBuf,
    pub group_key: String,
    pub role: PbrTextureRole,
    pub confidence: PbrMatchConfidence,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct DetectedPbrGroup {
    pub key: String,
    pub textures: Vec<DetectedPbrTexture>,
}

pub fn detect_pbr_texture_groups(
    paths: impl IntoIterator<Item = PathBuf>,
) -> Vec<DetectedPbrGroup> {
    let mut groups: BTreeMap<String, Vec<DetectedPbrTexture>> = BTreeMap::new();
    for path in paths {
        if !is_texture_path(&path) {
            continue;
        }
        let (group_key, role, confidence) = classify_texture_name(&path);
        groups
            .entry(group_key.clone())
            .or_default()
            .push(DetectedPbrTexture {
                path,
                group_key,
                role,
                confidence,
            });
    }
    groups
        .into_iter()
        .map(|(key, mut textures)| {
            textures.sort_by(|left, right| left.path.cmp(&right.path));
            DetectedPbrGroup { key, textures }
        })
        .collect()
}

fn classify_texture_name(path: &Path) -> (String, PbrTextureRole, PbrMatchConfidence) {
    let stem = path
        .file_stem()
        .and_then(|stem| stem.to_str())
        .unwrap_or("texture")
        .to_ascii_lowercase();
    let mut tokens: Vec<&str> = stem
        .split(|character: char| !character.is_ascii_alphanumeric())
        .filter(|token| !token.is_empty())
        .collect();
    while tokens
        .last()
        .is_some_and(|token| is_resolution_token(token))
    {
        tokens.pop();
    }
    let joined = tokens.join("_");
    let aliases: &[(&[&str], PbrTextureRole, PbrMatchConfidence)] = &[
        (
            &["occlusion", "roughness", "metallic"],
            PbrTextureRole::PackedOrm,
            PbrMatchConfidence::High,
        ),
        (
            &["base", "color"],
            PbrTextureRole::BaseColor,
            PbrMatchConfidence::High,
        ),
        (
            &["ambient", "occlusion"],
            PbrTextureRole::Occlusion,
            PbrMatchConfidence::High,
        ),
        (
            &["metallic", "roughness"],
            PbrTextureRole::PackedOrm,
            PbrMatchConfidence::Medium,
        ),
        (
            &["basecolor"],
            PbrTextureRole::BaseColor,
            PbrMatchConfidence::High,
        ),
        (
            &["albedo"],
            PbrTextureRole::BaseColor,
            PbrMatchConfidence::High,
        ),
        (
            &["diffuse"],
            PbrTextureRole::BaseColor,
            PbrMatchConfidence::High,
        ),
        (
            &["diff"],
            PbrTextureRole::BaseColor,
            PbrMatchConfidence::Medium,
        ),
        (
            &["normal"],
            PbrTextureRole::Normal,
            PbrMatchConfidence::High,
        ),
        (
            &["nor", "gl"],
            PbrTextureRole::Normal,
            PbrMatchConfidence::High,
        ),
        (
            &["nor", "dx"],
            PbrTextureRole::Normal,
            PbrMatchConfidence::High,
        ),
        (&["nrm"], PbrTextureRole::Normal, PbrMatchConfidence::Medium),
        (
            &["roughness"],
            PbrTextureRole::Roughness,
            PbrMatchConfidence::High,
        ),
        (
            &["rough"],
            PbrTextureRole::Roughness,
            PbrMatchConfidence::Medium,
        ),
        (
            &["metallic"],
            PbrTextureRole::Metallic,
            PbrMatchConfidence::High,
        ),
        (
            &["metalness"],
            PbrTextureRole::Metallic,
            PbrMatchConfidence::High,
        ),
        (
            &["metal"],
            PbrTextureRole::Metallic,
            PbrMatchConfidence::Medium,
        ),
        (
            &["occlusion"],
            PbrTextureRole::Occlusion,
            PbrMatchConfidence::High,
        ),
        (&["ao"], PbrTextureRole::Occlusion, PbrMatchConfidence::High),
        (
            &["orm"],
            PbrTextureRole::PackedOrm,
            PbrMatchConfidence::High,
        ),
        (
            &["arm"],
            PbrTextureRole::PackedOrm,
            PbrMatchConfidence::High,
        ),
        (
            &["height"],
            PbrTextureRole::Height,
            PbrMatchConfidence::High,
        ),
        (
            &["displacement"],
            PbrTextureRole::Height,
            PbrMatchConfidence::High,
        ),
        (
            &["disp"],
            PbrTextureRole::Height,
            PbrMatchConfidence::Medium,
        ),
        (&["n"], PbrTextureRole::Normal, PbrMatchConfidence::Low),
        (&["r"], PbrTextureRole::Roughness, PbrMatchConfidence::Low),
        (&["m"], PbrTextureRole::Metallic, PbrMatchConfidence::Low),
    ];
    for (suffix, role, confidence) in aliases {
        if tokens.ends_with(suffix) {
            let base_len = tokens.len().saturating_sub(suffix.len());
            let group = tokens[..base_len].join("_");
            return (
                if group.is_empty() {
                    joined.clone()
                } else {
                    group
                },
                *role,
                *confidence,
            );
        }
    }
    (
        joined,
        PbrTextureRole::Unknown,
        PbrMatchConfidence::Unresolved,
    )
}

fn is_resolution_token(token: &str) -> bool {
    token
        .strip_suffix('k')
        .is_some_and(|number| matches!(number, "1" | "2" | "4" | "8" | "16"))
}

fn is_texture_path(path: &Path) -> bool {
    matches!(
        path.extension()
            .and_then(|extension| extension.to_str())
            .unwrap_or_default()
            .to_ascii_lowercase()
            .as_str(),
        "png" | "jpg" | "jpeg" | "webp" | "ktx2" | "dds" | "tga" | "exr" | "hdr"
    )
}

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

    #[test]
    fn poly_haven_suffixes_form_one_high_confidence_group() {
        let groups = detect_pbr_texture_groups([
            "assets/metal_office_desk_diff_2k.jpg".into(),
            "assets/metal_office_desk_nor_gl_2k.jpg".into(),
            "assets/metal_office_desk_arm_2k.jpg".into(),
        ]);
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].key, "metal_office_desk");
        assert!(groups[0]
            .textures
            .iter()
            .any(|texture| texture.role == PbrTextureRole::BaseColor));
        assert!(groups[0]
            .textures
            .iter()
            .any(|texture| texture.role == PbrTextureRole::Normal));
        assert!(groups[0]
            .textures
            .iter()
            .any(|texture| texture.role == PbrTextureRole::PackedOrm));
    }

    #[test]
    fn unknown_images_remain_visible_for_manual_resolution() {
        let groups = detect_pbr_texture_groups(["assets/wood_preview.png".into()]);
        assert_eq!(groups[0].textures[0].role, PbrTextureRole::Unknown);
        assert_eq!(
            groups[0].textures[0].confidence,
            PbrMatchConfidence::Unresolved
        );
    }
}
