tyt-material 0.1.6

use crate::{Dependencies, Error, Result};
use clap::Parser;
use std::path::{Path, PathBuf};

/// Creates an MSE png from material texture maps. The output png packs:
///   R = metalness
///   G = smoothness (1 - roughness), or roughness when `--output-rough` is set
///   B = emissive (emissive alpha)
/// Optionally copies the albedo texture alongside.
///
/// By default, metalness and roughness are read from separate textures
/// (`*metalness.png` and `*roughness.png`). Pass `--combine-metal-rough` to
/// read both from a single legacy texture (R = metal, A = roughness).
#[derive(Clone, Debug, Parser)]
pub struct CreateMse {
    /// The output base path. Output files will be `{out_base}-mse.png` and
    /// `{out_base}-albedo.png`.
    #[arg(value_name = "out-base")]
    out_base: String,

    /// Search prefix for texture files. When set, searches for
    /// `{prefix}-metalness.png`, `{prefix}-roughness.png`,
    /// `{prefix}-emission.png`, `{prefix}-albedo.png`.
    #[arg(value_name = "prefix", long)]
    prefix: Option<String>,

    /// Explicit path to the metal texture. Cannot be used with
    /// `--combine-metal-rough`.
    #[arg(value_name = "metal", long, conflicts_with = "combine_metal_rough")]
    metal: Option<PathBuf>,

    /// Explicit path to the rough texture. Cannot be used with
    /// `--combine-metal-rough`.
    #[arg(value_name = "rough", long, conflicts_with = "combine_metal_rough")]
    rough: Option<PathBuf>,

    /// Read metalness and roughness from a single combined texture
    /// (legacy mode: R = metal, A = roughness).
    #[arg(value_name = "combine-metal-rough", long)]
    combine_metal_rough: bool,

    /// Explicit path to the combined metal_rough texture. Only valid with
    /// `--combine-metal-rough`.
    #[arg(value_name = "metal-rough", long, requires = "combine_metal_rough")]
    metal_rough: Option<PathBuf>,

    /// Explicit path to the emissive texture.
    #[arg(value_name = "emissive", long)]
    emissive: Option<PathBuf>,

    /// Explicit path to the albedo texture.
    #[arg(value_name = "albedo", long)]
    albedo: Option<PathBuf>,

    /// Skip the metalness channel (R will be black).
    #[arg(value_name = "ignore-metal", long)]
    ignore_metal: bool,

    /// Skip the roughness channel (G will be black).
    #[arg(value_name = "ignore-rough", long)]
    ignore_rough: bool,

    /// Skip the emissive channel (B will be black).
    #[arg(value_name = "ignore-emissive", long)]
    ignore_emissive: bool,

    /// Skip the albedo pass-through copy.
    #[arg(value_name = "ignore-albedo", long)]
    ignore_albedo: bool,

    /// Output roughness directly into the G channel instead of converting it
    /// to smoothness (1 - roughness).
    #[arg(value_name = "output-rough", long)]
    output_rough: bool,
}

impl CreateMse {
    pub fn execute(self, dependencies: impl Dependencies) -> Result<()> {
        let CreateMse {
            out_base,
            prefix,
            metal,
            rough,
            combine_metal_rough,
            metal_rough,
            emissive,
            albedo,
            ignore_metal,
            ignore_rough,
            ignore_emissive,
            ignore_albedo,
            output_rough,
        } = self;

        // ----------------------------------------------------------------
        // Resolve texture paths
        // ----------------------------------------------------------------
        let metal_rough_path = if !combine_metal_rough || (ignore_metal && ignore_rough) {
            None
        } else {
            Some(match metal_rough {
                Some(p) => coerce_png(p),
                None => match &prefix {
                    Some(pfx) => dependencies.glob_single_match(&format!("{pfx}-metalness.png"))?,
                    None => dependencies.glob_single_match("*metalness.png")?,
                },
            })
        };

        let metal_path = if combine_metal_rough || ignore_metal {
            None
        } else {
            Some(match metal {
                Some(p) => coerce_png(p),
                None => match &prefix {
                    Some(pfx) => dependencies.glob_single_match(&format!("{pfx}-metalness.png"))?,
                    None => dependencies.glob_single_match("*metalness.png")?,
                },
            })
        };

        let rough_path = if combine_metal_rough || ignore_rough {
            None
        } else {
            Some(match rough {
                Some(p) => coerce_png(p),
                None => match &prefix {
                    Some(pfx) => dependencies.glob_single_match(&format!("{pfx}-roughness.png"))?,
                    None => dependencies.glob_single_match("*roughness.png")?,
                },
            })
        };

        let emissive_path = if ignore_emissive {
            None
        } else {
            Some(match emissive {
                Some(p) => coerce_png(p),
                None => match &prefix {
                    Some(pfx) => dependencies.glob_single_match(&format!("{pfx}-emission.png"))?,
                    None => dependencies.glob_single_match("*emission.png")?,
                },
            })
        };

        let albedo_path = if ignore_albedo {
            None
        } else {
            Some(match albedo {
                Some(p) => coerce_png(p),
                None => match &prefix {
                    Some(pfx) => dependencies.glob_single_match(&format!("{pfx}-albedo.png"))?,
                    None => dependencies.glob_single_match("*albedo.png")?,
                },
            })
        };

        // ----------------------------------------------------------------
        // Determine base image for sizing
        // ----------------------------------------------------------------
        let base_img = metal_rough_path
            .as_ref()
            .or(metal_path.as_ref())
            .or(rough_path.as_ref())
            .or(emissive_path.as_ref())
            .or(albedo_path.as_ref())
            .ok_or_else(|| Error::Glob("all channels are ignored; nothing to do".into()))?;

        let size_output = dependencies.exec_magick([
            "identify".into(),
            "-ping".into(),
            "-format".into(),
            "%wx%h".into(),
            base_img.to_string_lossy().into_owned(),
        ])?;
        let size = String::from_utf8_lossy(&size_output).trim().to_string();
        if size.is_empty() {
            return Err(Error::Glob(format!(
                "couldn't determine image size for: {}",
                base_img.display()
            )));
        }

        // ----------------------------------------------------------------
        // Create temp dir for intermediate channel PNGs
        // ----------------------------------------------------------------
        let tmpdir = dependencies.create_temp_dir()?;
        let result = self::create_mse_inner(
            &dependencies,
            &metal_rough_path,
            &metal_path,
            &rough_path,
            &emissive_path,
            &albedo_path,
            output_rough,
            &out_base,
            &size,
            &tmpdir,
        );
        dependencies.remove_dir_all(&tmpdir)?;
        result
    }
}

fn create_mse_inner(
    dependencies: &impl Dependencies,
    metal_rough_path: &Option<PathBuf>,
    metal_path: &Option<PathBuf>,
    rough_path: &Option<PathBuf>,
    emissive_path: &Option<PathBuf>,
    albedo_path: &Option<PathBuf>,
    output_rough: bool,
    out_base: &str,
    size: &str,
    tmpdir: &Path,
) -> Result<()> {
    let r_img = tmpdir.join("r.png");
    let g_img = tmpdir.join("g.png");
    let b_img = tmpdir.join("b.png");

    let r_str = r_img.to_string_lossy().into_owned();
    let g_str = g_img.to_string_lossy().into_owned();
    let b_str = b_img.to_string_lossy().into_owned();

    // R channel: metalness
    match (metal_rough_path, metal_path) {
        (Some(mr), _) => {
            let mr_str = mr.to_string_lossy().into_owned();
            dependencies.exec_magick([
                &mr_str,
                "-colorspace",
                "sRGB",
                "-alpha",
                "on",
                "-channel",
                "R",
                "-separate",
                "+channel",
                "-resize",
                &format!("{size}!"),
                &r_str,
            ])?;
        }
        (None, Some(m)) => {
            let metal_str = m.to_string_lossy().into_owned();
            dependencies.exec_magick([
                &metal_str,
                "-colorspace",
                "sRGB",
                "-channel",
                "R",
                "-separate",
                "+channel",
                "-resize",
                &format!("{size}!"),
                &r_str,
            ])?;
        }
        (None, None) => {
            dependencies.exec_magick(["-size", size, "xc:black", &r_str])?;
        }
    }

    // G channel: smoothness = 1 - roughness, or roughness if output_rough
    match (metal_rough_path, rough_path) {
        (Some(mr), _) => {
            let mr_str = mr.to_string_lossy().into_owned();
            if output_rough {
                dependencies.exec_magick([
                    &mr_str,
                    "-colorspace",
                    "sRGB",
                    "-alpha",
                    "on",
                    "-alpha",
                    "extract",
                    "-resize",
                    &format!("{size}!"),
                    &g_str,
                ])?;
            } else {
                dependencies.exec_magick([
                    &mr_str,
                    "-colorspace",
                    "sRGB",
                    "-alpha",
                    "on",
                    "-alpha",
                    "extract",
                    "-fx",
                    "u==0 ? 0 : 1-u",
                    "-resize",
                    &format!("{size}!"),
                    &g_str,
                ])?;
            }
        }
        (None, Some(r)) => {
            let rough_str = r.to_string_lossy().into_owned();
            if output_rough {
                dependencies.exec_magick([
                    &rough_str,
                    "-colorspace",
                    "sRGB",
                    "-channel",
                    "R",
                    "-separate",
                    "+channel",
                    "-resize",
                    &format!("{size}!"),
                    &g_str,
                ])?;
            } else {
                dependencies.exec_magick([
                    &rough_str,
                    "-colorspace",
                    "sRGB",
                    "-channel",
                    "R",
                    "-separate",
                    "+channel",
                    "-fx",
                    "1-u",
                    "-resize",
                    &format!("{size}!"),
                    &g_str,
                ])?;
            }
        }
        (None, None) => {
            dependencies.exec_magick(["-size", size, "xc:black", &g_str])?;
        }
    }

    // B channel: emissive (alpha channel of emissive)
    match emissive_path {
        None => {
            dependencies.exec_magick(["-size", size, "xc:black", &b_str])?;
        }
        Some(em) => {
            let em_str = em.to_string_lossy().into_owned();
            dependencies.exec_magick([
                &em_str,
                "-colorspace",
                "sRGB",
                "-alpha",
                "on",
                "-alpha",
                "extract",
                "-resize",
                &format!("{size}!"),
                &b_str,
            ])?;
        }
    }

    // Combine R/G/B into MSE
    let mse_out = format!("{out_base}-mse.png");
    dependencies.exec_magick([
        &r_str,
        &g_str,
        &b_str,
        "-combine",
        "-colorspace",
        "sRGB",
        &mse_out,
    ])?;

    dependencies.write_stdout(format!("Wrote: {mse_out}\n").as_bytes())?;

    // Copy albedo if not ignored
    if let Some(albedo) = albedo_path {
        let albedo_out = format!("{out_base}-albedo.png");
        dependencies.copy_file(albedo, &albedo_out)?;
        dependencies.write_stdout(format!("Wrote: {albedo_out}\n").as_bytes())?;
    }

    Ok(())
}

/// If the path has no extension, assume `.png`.
fn coerce_png(p: PathBuf) -> PathBuf {
    if p.extension().is_none() {
        p.with_extension("png")
    } else {
        p
    }
}