use clap::{Parser, ValueEnum};
use std::path::{Path, PathBuf};

#[derive(Parser, Debug, Clone, ValueEnum, PartialEq)]
pub enum Mode {
    TallyKeys,
    CutEdges,
    PolsbyPopper,
    ChangedAssignments,
    RegionSplits,
    RegionPieces,
    UniquePlans,
    ExtractUniquePlans,
}

#[derive(Parser, Debug)]
#[command(
    name = "BEN Process Tool",
    about = "A tool for processing BEN files and saving outputs to Parquet or text.",
    version = "0.1.0"
)]
pub struct Args {
    #[arg(short, long, default_value = "cut-edges")]
    pub mode: Mode,
    #[arg(short, long)]
    pub graph_file: Option<String>,
    #[arg(short, long)]
    pub ben_file: String,
    #[arg(short, long, default_value_t = false)]
    pub normalize: bool,
    #[arg(long)]
    pub max_accepted: Option<usize>,
    /// Randomize merge-split label reassignments (changed-assignments mode only).
    /// Only set this for MCMC merge-split ensembles. Default: off.
    #[arg(long, default_value_t = false)]
    pub randomize_reassignments: bool,
    #[arg(short, long, num_args(1..))]
    pub keys: Vec<String>,
    #[arg(long)]
    pub edge_weight_key: Option<String>,
    #[arg(long)]
    pub area_key: Option<String>,
    #[arg(long)]
    pub perim_key: Option<String>,
    #[arg(long)]
    pub boundary_perim_key: Option<String>,
    #[arg(long)]
    pub shared_perim_key: Option<String>,
    #[arg(long, default_value_t = false)]
    pub no_progress: bool,
    #[arg(long)]
    pub output_dir: Option<String>,
    /// Use Brotli compression for Parquet output (default: Snappy).
    /// Brotli is CPU-heavy and rarely worth it unless you're storage-bound.
    #[arg(long, default_value_t = false)]
    pub high_compression: bool,
}

/// Build an output path next to (or alongside) `in_ben_file` by stripping the
/// BEN extension and appending `suffix`.
///
/// The previous implementation used `String::replace(".jsonl.ben", suffix)`,
/// which silently returned the input unchanged when the input did not end in
/// `.jsonl.ben` — meaning a downstream `File::create` could overwrite the
/// input BEN. This version guarantees the output path differs from the input
/// path: it strips a trailing `.jsonl.ben` (or just `.ben`) before appending,
/// and asserts distinctness as a safety net.
pub fn build_output_path(in_ben_file: &str, suffix: &str, output_dir: Option<&str>) -> String {
    let in_path = Path::new(in_ben_file);
    let stem = ben_stem(in_ben_file);
    let new_name = format!("{}{}", stem, suffix);

    let out = match output_dir {
        Some(dir) => PathBuf::from(dir)
            .join(new_name)
            .to_string_lossy()
            .into_owned(),
        None => match in_path.parent() {
            Some(parent) if !parent.as_os_str().is_empty() => {
                parent.join(new_name).to_string_lossy().into_owned()
            }
            _ => new_name,
        },
    };

    assert!(
        out != in_ben_file,
        "refusing to overwrite input BEN file {:?}: derived output path is identical",
        in_ben_file
    );
    out
}

fn ben_stem(in_ben_file: &str) -> String {
    let in_path = Path::new(in_ben_file);
    let base_name = in_path
        .file_name()
        .expect("Failed to extract basename")
        .to_string_lossy()
        .into_owned();
    base_name
        .strip_suffix(".jsonl.ben")
        .or_else(|| base_name.strip_suffix(".ben"))
        .unwrap_or(&base_name)
        .to_string()
}

pub fn build_tally_output_dir(in_ben_file: &str, output_dir: Option<&str>) -> PathBuf {
    let ben_path = Path::new(in_ben_file);
    let ben_stem = ben_stem(in_ben_file);
    let dir_name = format!("{}_tallies", ben_stem);

    match output_dir {
        Some(dir) => PathBuf::from(dir).join(dir_name),
        None => ben_path
            .parent()
            .unwrap_or_else(|| Path::new("."))
            .join(dir_name),
    }
}

pub fn build_tally_output_path(in_ben_file: &str, key: &str, output_dir: Option<&str>) -> PathBuf {
    let ben_stem = ben_stem(in_ben_file);
    build_tally_output_dir(in_ben_file, output_dir)
        .join(format!("{}_tally_{}.parquet", key, ben_stem))
}

#[cfg(test)]
mod tests {
    use super::{build_output_path, build_tally_output_dir, build_tally_output_path};
    use std::path::PathBuf;

    #[test]
    fn build_output_path_replaces_suffix_in_place_without_output_dir() {
        assert_eq!(
            build_output_path("/tmp/runs/plans.jsonl.ben", "_cut_edges.parquet", None),
            "/tmp/runs/plans_cut_edges.parquet"
        );
    }

    #[test]
    fn build_output_path_uses_basename_when_output_dir_is_set() {
        assert_eq!(
            build_output_path(
                "/tmp/runs/plans.jsonl.ben",
                "_unique_plans.txt",
                Some("/tmp/out"),
            ),
            "/tmp/out/plans_unique_plans.txt"
        );
    }

    #[test]
    fn build_output_path_strips_plain_ben_extension() {
        assert_eq!(
            build_output_path("/tmp/runs/plans.ben", "_cut_edges.parquet", None),
            "/tmp/runs/plans_cut_edges.parquet"
        );
    }

    #[test]
    fn build_output_path_appends_suffix_when_input_has_no_known_extension() {
        // Old behavior silently returned the input unchanged here, which would
        // cause File::create to overwrite the input BEN. New behavior appends.
        assert_eq!(
            build_output_path("/tmp/runs/plans", "_cut_edges.parquet", None),
            "/tmp/runs/plans_cut_edges.parquet"
        );
    }

    #[test]
    fn build_output_path_handles_bare_filename_without_parent_dir() {
        assert_eq!(
            build_output_path("plans.jsonl.ben", "_cut_edges.parquet", None),
            "plans_cut_edges.parquet"
        );
    }

    #[test]
    #[should_panic(expected = "refusing to overwrite input BEN file")]
    fn build_output_path_panics_if_output_would_equal_input() {
        // An empty suffix combined with a no-extension input would yield the
        // same path; the assertion must catch this before File::create runs.
        let _ = build_output_path("/tmp/runs/plans", "", None);
    }

    #[test]
    fn build_tally_output_dir_uses_ben_stem() {
        assert_eq!(
            build_tally_output_dir("/tmp/runs/plans.jsonl.ben", Some("/tmp/out")),
            PathBuf::from("/tmp/out/plans_tallies")
        );
    }

    #[test]
    fn build_tally_output_path_uses_key_and_ben_stem() {
        assert_eq!(
            build_tally_output_path("/tmp/runs/plans.jsonl.ben", "pop", Some("/tmp/out")),
            PathBuf::from("/tmp/out/plans_tallies/pop_tally_plans.parquet")
        );
    }
}
