// Rust guideline compliant 2026-06-22

#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

use std::ffi::OsStr;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::SystemTime;

use anyhow::{Context, Result};
use clap::{Args, Parser, Subcommand, ValueEnum};
use moodbar_core::{
    analysis_to_raw_rgb_bytes, analyze_path, analyze_pcm_mono, render_png, render_svg,
    DetectionMode, GenerateOptions, NormalizeMode, PngOptions, SvgOptions, SvgShape,
};
use rayon::prelude::*;
use serde::Serialize;
use walkdir::WalkDir;

#[derive(Parser, Debug)]
#[command(name = "moodbar")]
#[command(about = "Generate moodbar data from audio files", long_about = None)]
struct Cli {
    #[arg(long, global = true, help = "Emit machine-readable JSON output")]
    json: bool,

    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand, Debug)]
enum Command {
    /// Generate a single moodbar file.
    Generate {
        #[arg(short = 'i', long)]
        input: Option<PathBuf>,
        #[arg(short = 'o', long)]
        output: PathBuf,
        #[command(flatten)]
        dsp: DspArgs,
        #[command(flatten)]
        render: RenderArgs,
        #[arg(long, help = "Read PCM samples from stdin instead of input file path")]
        stdin: bool,
        #[arg(long, value_enum, default_value_t = StdinFormat::F32Le)]
        stdin_format: StdinFormat,
        #[arg(long, default_value_t = 44_100)]
        sample_rate: u32,
        #[arg(long, default_value_t = 1)]
        channels: usize,
        #[arg(long)]
        force: bool,
    },

    /// Recursively generate moodbar files from a directory.
    Batch {
        #[arg(short = 'i', long)]
        input_dir: PathBuf,
        #[arg(short = 'o', long)]
        output_dir: PathBuf,
        #[command(flatten)]
        dsp: DspArgs,
        #[command(flatten)]
        render: RenderArgs,
        #[arg(long, default_value = "mood")]
        output_ext: String,
        #[arg(long, default_value_t = 0, help = "Parallel worker count (0 = auto)")]
        jobs: usize,
        #[arg(long, help = "Show per-file batch progress on stderr")]
        progress: bool,
        #[arg(long)]
        force: bool,
    },

    /// Print basic information about an existing moodbar file.
    Inspect {
        #[arg(short = 'i', long)]
        input: PathBuf,
    },
}

#[derive(Args, Debug, Clone)]
struct DspArgs {
    #[arg(long, default_value_t = 2048)]
    fft_size: usize,
    #[arg(long, default_value_t = 500.0)]
    low_cut_hz: f32,
    #[arg(long, default_value_t = 2000.0)]
    mid_cut_hz: f32,
    #[arg(long, value_enum, default_value_t = NormalizeModeArg::PerChannelPeak)]
    normalize_mode: NormalizeModeArg,
    #[arg(long, default_value_t = 1e-12)]
    deterministic_floor: f64,
    #[arg(long, value_enum, default_value_t = DetectionModeArg::SpectralEnergy)]
    detection_mode: DetectionModeArg,
    #[arg(long, default_value_t = 1)]
    frames_per_color: usize,
    #[arg(long, value_delimiter = ',')]
    band_edges_hz: Vec<f32>,
    /// Effective playback rate for frequency-to-color mapping (e.g. 1.09 = +9% pitch).
    #[arg(long)]
    playback_rate: Option<f32>,
    #[arg(long, value_enum, default_value_t = ThemeArg::Classic)]
    theme: ThemeArg,
    #[arg(long, value_delimiter = ',')]
    colors: Vec<String>,
}

#[derive(Args, Debug, Clone)]
struct RenderArgs {
    #[arg(long, value_enum, default_value_t = OutputFormat::RawRgbV1)]
    format: OutputFormat,
    #[arg(long, value_enum, default_value_t = SvgShapeArg::Strip)]
    svg_shape: SvgShapeArg,
    #[arg(long, default_value_t = 1200)]
    svg_width: u32,
    #[arg(long, default_value_t = 96)]
    svg_height: u32,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum OutputFormat {
    RawRgbV1,
    Svg,
    Png,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum NormalizeModeArg {
    PerChannelPeak,
    GlobalPeak,
}

impl NormalizeModeArg {
    fn into_core(self) -> NormalizeMode {
        match self {
            NormalizeModeArg::PerChannelPeak => NormalizeMode::PerChannelPeak,
            NormalizeModeArg::GlobalPeak => NormalizeMode::GlobalPeak,
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum ThemeArg {
    Classic,
    Cool,
    Light,
}

impl ThemeArg {
    fn into_core(self) -> moodbar_core::Theme {
        match self {
            ThemeArg::Classic => moodbar_core::Theme::Classic,
            ThemeArg::Cool => moodbar_core::Theme::Cool,
            ThemeArg::Light => moodbar_core::Theme::Light,
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum DetectionModeArg {
    SpectralEnergy,
    SpectralFlux,
}

impl DetectionModeArg {
    fn into_core(self) -> DetectionMode {
        match self {
            DetectionModeArg::SpectralEnergy => DetectionMode::SpectralEnergy,
            DetectionModeArg::SpectralFlux => DetectionMode::SpectralFlux,
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum SvgShapeArg {
    Strip,
    Waveform,
    SplitStacked,
    SplitWaveform,
    SplitLanes,
    SplitCentrifugal,
    SplitOverlapping,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum StdinFormat {
    F32Le,
    S16Le,
}

impl SvgShapeArg {
    fn into_core(self) -> SvgShape {
        match self {
            SvgShapeArg::Strip => SvgShape::Strip,
            SvgShapeArg::Waveform => SvgShape::Waveform,
            SvgShapeArg::SplitStacked => SvgShape::SplitStacked,
            SvgShapeArg::SplitWaveform => SvgShape::SplitWaveform,
            SvgShapeArg::SplitLanes => SvgShape::SplitLanes,
            SvgShapeArg::SplitCentrifugal => SvgShape::SplitCentrifugal,
            SvgShapeArg::SplitOverlapping => SvgShape::SplitOverlapping,
        }
    }
}

#[derive(Serialize)]
struct GenerateResult {
    input: String,
    output: String,
    bytes_written: usize,
    frames: usize,
    channels: usize,
    format: String,
}

#[derive(Serialize)]
struct BatchResult {
    processed: usize,
    succeeded: usize,
    skipped: usize,
    failed: usize,
}

#[derive(Serialize)]
struct InspectResult {
    input: String,
    bytes: usize,
    frames: usize,
}

fn main() {
    let cli = Cli::parse();
    let exit_code = match run(cli) {
        Ok(code) => code,
        Err(err) => {
            eprintln!("error: {err:#}");
            2
        }
    };
    std::process::exit(exit_code);
}

fn run(cli: Cli) -> Result<i32> {
    match cli.command {
        Command::Generate {
            input,
            output,
            dsp,
            render,
            stdin,
            stdin_format,
            sample_rate,
            channels,
            force,
        } => {
            if !stdin && input.is_none() {
                anyhow::bail!("input path is required unless --stdin is provided");
            }
            if stdin && input.is_some() {
                anyhow::bail!("use either --stdin or --input, not both");
            }
            if channels == 0 {
                anyhow::bail!("--channels must be >= 1");
            }
            ensure_can_write(&output, force)?;

            let options = build_options(&dsp)?;
            let analysis = if stdin {
                let mono = read_stdin_pcm_mono(stdin_format, channels)?;
                analyze_pcm_mono(sample_rate, &mono, &options)
            } else {
                let inpath = input.as_ref().expect("validated above");
                analyze_path(inpath, &options).with_context(|| {
                    format!("failed to generate moodbar for {}", inpath.display())
                })?
            };

            let bytes_written = match render.format {
                OutputFormat::RawRgbV1 => {
                    let bytes = analysis_to_raw_rgb_bytes(&analysis);
                    write_bytes(&output, &bytes)?;
                    bytes.len()
                }
                OutputFormat::Svg => {
                    let svg = render_svg(
                        &analysis,
                        &SvgOptions {
                            width: render.svg_width,
                            height: render.svg_height,
                            shape: render.svg_shape.into_core(),
                            ..SvgOptions::default()
                        },
                    );
                    write_text(&output, &svg)?;
                    svg.len()
                }
                OutputFormat::Png => {
                    let png = render_png(
                        &analysis,
                        &PngOptions {
                            width: render.svg_width,
                            height: render.svg_height,
                            shape: render.svg_shape.into_core(),
                        },
                    )?;
                    write_bytes(&output, &png)?;
                    png.len()
                }
            };

            let result = GenerateResult {
                input: input
                    .as_ref()
                    .map(|p| p.display().to_string())
                    .unwrap_or_else(|| "<stdin>".to_string()),
                output: output.display().to_string(),
                bytes_written,
                frames: analysis.frames.len() / analysis.channel_count.max(1),
                channels: analysis.channel_count,
                format: match render.format {
                    OutputFormat::RawRgbV1 => "raw_rgb_v1",
                    OutputFormat::Svg => "svg",
                    OutputFormat::Png => "png",
                }
                .to_string(),
            };
            print_result(cli.json, &result)?;
            Ok(0)
        }
        Command::Batch {
            input_dir,
            output_dir,
            dsp,
            render,
            output_ext,
            jobs,
            progress,
            force,
        } => {
            let options = build_options(&dsp)?;
            let candidates = WalkDir::new(&input_dir)
                .follow_links(false)
                .into_iter()
                .filter_map(std::result::Result::ok)
                .filter(|e| e.file_type().is_file())
                .filter(|e| looks_like_audio(e.path()))
                .map(|e| e.path().to_path_buf())
                .collect::<Vec<_>>();

            let processed = candidates.len();
            let input_dir = Arc::new(input_dir);
            let output_dir = Arc::new(output_dir);
            let output_ext = Arc::new(output_ext);
            let options = Arc::new(options);
            let render = Arc::new(render);

            let worker_count = if jobs == 0 {
                std::thread::available_parallelism()
                    .map(|n| n.get())
                    .unwrap_or(1)
            } else {
                jobs.max(1)
            };

            let pool = rayon::ThreadPoolBuilder::new()
                .num_threads(worker_count)
                .build()
                .context("failed to build rayon thread pool")?;

            let progress_counter = Arc::new(AtomicUsize::new(0));
            let progress_lock = Arc::new(std::sync::Mutex::new(()));

            let outcomes = pool.install(|| {
                candidates
                    .par_iter()
                    .map(|src| {
                        let rel = src.strip_prefix(&*input_dir).unwrap_or(src);
                        let mut dst = output_dir.join(rel);
                        dst.set_extension(output_ext.as_ref());
                        let outcome = process_batch_item(src, &dst, &options, &render, force)
                            .with_context(|| format!("{} -> {}", src.display(), dst.display()));

                        if progress {
                            let done = progress_counter.fetch_add(1, Ordering::SeqCst) + 1;
                            let state = match &outcome {
                                Ok(BatchItemStatus::Generated) => "generated",
                                Ok(BatchItemStatus::Skipped) => "skipped",
                                Err(_) => "failed",
                            };
                            if let Ok(_guard) = progress_lock.lock() {
                                eprintln!("[{done}/{processed}] {state} {}", src.display());
                            }
                        }
                        outcome
                    })
                    .collect::<Vec<_>>()
            });

            let mut succeeded = 0usize;
            let mut skipped = 0usize;
            let mut failed = 0usize;
            for (src, outcome) in candidates.iter().zip(outcomes) {
                match outcome {
                    Ok(BatchItemStatus::Generated) => succeeded += 1,
                    Ok(BatchItemStatus::Skipped) => skipped += 1,
                    Err(err) => {
                        failed += 1;
                        if cli.json {
                            eprintln!(
                                "{{\"input\":\"{}\",\"error\":\"{}\"}}",
                                src.display(),
                                escape_json_string(&format!("{err:#}"))
                            );
                        } else {
                            eprintln!("failed: {} ({err:#})", src.display());
                        }
                    }
                }
            }

            let result = BatchResult {
                processed,
                succeeded,
                skipped,
                failed,
            };
            print_result(cli.json, &result)?;
            Ok(if failed == 0 { 0 } else { 1 })
        }
        Command::Inspect { input } => {
            let bytes = fs::read(&input)?;
            let result = InspectResult {
                input: input.display().to_string(),
                bytes: bytes.len(),
                frames: bytes.len() / 3,
            };
            print_result(cli.json, &result)?;
            Ok(0)
        }
    }
}
fn build_options(dsp: &DspArgs) -> Result<GenerateOptions> {
    use anyhow::Context as _;
    let custom_colors = if dsp.colors.is_empty() {
        None
    } else {
        let mut parsed = Vec::new();
        for s in &dsp.colors {
            let color = parse_cli_hex_color(s)
                .with_context(|| format!("invalid color hex value: '{}'", s))?;
            parsed.push(color);
        }
        Some(parsed)
    };

    Ok(GenerateOptions {
        fft_size: dsp.fft_size,
        low_cut_hz: dsp.low_cut_hz,
        mid_cut_hz: dsp.mid_cut_hz,
        normalize_mode: dsp.normalize_mode.into_core(),
        deterministic_floor: dsp.deterministic_floor,
        detection_mode: dsp.detection_mode.into_core(),
        frames_per_color: dsp.frames_per_color,
        band_edges_hz: dsp.band_edges_hz.clone(),
        playback_rate: dsp.playback_rate,
        theme: dsp.theme.into_core(),
        custom_colors,
    })
}

fn parse_cli_hex_color(s: &str) -> Result<[u8; 3]> {
    let cleaned = s.trim().trim_start_matches('#');
    if cleaned.len() != 6 {
        anyhow::bail!("length must be exactly 6 characters (e.g. #FF0000 or FF0000)");
    }
    let r = u8::from_str_radix(&cleaned[0..2], 16).context("invalid red channel hex")?;
    let g = u8::from_str_radix(&cleaned[2..4], 16).context("invalid green channel hex")?;
    let b = u8::from_str_radix(&cleaned[4..6], 16).context("invalid blue channel hex")?;
    Ok([r, g, b])
}

#[derive(Copy, Clone)]
enum BatchItemStatus {
    Generated,
    Skipped,
}

fn process_batch_item(
    src: &Path,
    dst: &Path,
    options: &GenerateOptions,
    render: &RenderArgs,
    force: bool,
) -> Result<BatchItemStatus> {
    if should_skip_unchanged(src, dst, force)? {
        return Ok(BatchItemStatus::Skipped);
    }
    let analysis = analyze_path(src, options)
        .with_context(|| format!("decode/generate failed for {}", src.display()))?;

    match render.format {
        OutputFormat::RawRgbV1 => {
            let bytes = analysis_to_raw_rgb_bytes(&analysis);
            write_bytes(dst, &bytes)?;
        }
        OutputFormat::Svg => {
            let svg = render_svg(
                &analysis,
                &SvgOptions {
                    width: render.svg_width,
                    height: render.svg_height,
                    shape: render.svg_shape.into_core(),
                    ..SvgOptions::default()
                },
            );
            write_text(dst, &svg)?;
        }
        OutputFormat::Png => {
            let png = render_png(
                &analysis,
                &PngOptions {
                    width: render.svg_width,
                    height: render.svg_height,
                    shape: render.svg_shape.into_core(),
                },
            )?;
            write_bytes(dst, &png)?;
        }
    }
    Ok(BatchItemStatus::Generated)
}

fn should_skip_unchanged(src: &Path, dst: &Path, force: bool) -> Result<bool> {
    if force || !dst.exists() {
        return Ok(false);
    }

    let src_meta = fs::metadata(src)?;
    let dst_meta = fs::metadata(dst)?;
    let src_mtime = src_meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
    let dst_mtime = dst_meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
    Ok(dst_mtime >= src_mtime)
}

fn read_stdin_pcm_mono(format: StdinFormat, channels: usize) -> Result<Vec<f32>> {
    let mut bytes = Vec::new();
    std::io::stdin().read_to_end(&mut bytes)?;
    if bytes.is_empty() {
        return Ok(Vec::new());
    }

    match format {
        StdinFormat::F32Le => {
            if bytes.len() % 4 != 0 {
                anyhow::bail!("stdin length must be a multiple of 4 bytes for f32le");
            }
            let samples = bytes
                .chunks_exact(4)
                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                .collect::<Vec<_>>();
            downmix_to_mono(&samples, channels)
        }
        StdinFormat::S16Le => {
            if bytes.len() % 2 != 0 {
                anyhow::bail!("stdin length must be a multiple of 2 bytes for s16le");
            }
            let samples = bytes
                .chunks_exact(2)
                .map(|b| i16::from_le_bytes([b[0], b[1]]) as f32 / i16::MAX as f32)
                .collect::<Vec<_>>();
            downmix_to_mono(&samples, channels)
        }
    }
}

fn downmix_to_mono(samples: &[f32], channels: usize) -> Result<Vec<f32>> {
    if channels == 1 {
        return Ok(samples.to_vec());
    }
    if samples.len() % channels != 0 {
        anyhow::bail!("sample count is not divisible by channel count");
    }

    let mut mono = Vec::with_capacity(samples.len() / channels);
    for frame in samples.chunks_exact(channels) {
        let sum = frame.iter().copied().sum::<f32>();
        mono.push(sum / channels as f32);
    }
    Ok(mono)
}

fn ensure_can_write(path: &Path, force: bool) -> Result<()> {
    if path.exists() && !force {
        anyhow::bail!("{} exists; pass --force to overwrite", path.display());
    }
    Ok(())
}

fn write_bytes(path: &Path, bytes: &[u8]) -> Result<()> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(path, bytes)?;
    Ok(())
}

fn write_text(path: &Path, text: &str) -> Result<()> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(path, text.as_bytes())?;
    Ok(())
}

fn print_result<T: Serialize>(json: bool, value: &T) -> Result<()> {
    if json {
        println!("{}", serde_json::to_string_pretty(value)?);
    } else {
        match serde_json::to_value(value)? {
            serde_json::Value::Object(map) => {
                for (k, v) in map {
                    println!("{k}: {v}");
                }
            }
            other => println!("{other}"),
        }
    }
    Ok(())
}

fn looks_like_audio(path: &Path) -> bool {
    let ext = path.extension().and_then(OsStr::to_str).unwrap_or_default();
    matches!(
        ext.to_ascii_lowercase().as_str(),
        "mp3" | "ogg" | "flac" | "wav" | "m4a" | "aac" | "opus"
    )
}

fn escape_json_string(s: &str) -> String {
    serde_json::to_string(s)
        .unwrap_or_else(|_| "\"<encode-error>\"".to_string())
        .trim_matches('"')
        .to_string()
}
