use std::borrow::ToOwned;

use crate::workdir::Workdir;

macro_rules! part_eq {
    ($wrk:expr_2021, $path:expr_2021, $expected:expr_2021) => {
        assert_eq!(
            $wrk.from_str::<String>(&$wrk.path($path)),
            $expected.to_owned()
        );
    };
}

fn data(headers: bool) -> Vec<Vec<String>> {
    let mut rows = vec![
        svec!["NY", "Manhattan"],
        svec!["CA", "San Francisco"],
        svec!["TX", "Dallas"],
        svec!["NY", "Buffalo"],
        svec!["TX", "Fort Worth"],
    ];
    if headers {
        rows.insert(0, svec!["state", "city"]);
    }
    rows
}

#[test]
fn partition() {
    let wrk = Workdir::new("partition");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("state").arg(wrk.path(".")).arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
state,city
CA,San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
state,city
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
state,city
TX,Dallas
TX,Fort Worth
"
    );
}

#[test]
fn partition_drop() {
    let wrk = Workdir::new("partition_drop");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("--drop")
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
city
San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
city
Manhattan
Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
city
Dallas
Fort Worth
"
    );
}

#[test]
fn partition_without_headers() {
    let wrk = Workdir::new("partition_without_headers");
    wrk.create("in.csv", data(false));

    let mut cmd = wrk.command("partition");
    cmd.arg("--no-headers")
        .arg("1")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
CA,San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
TX,Dallas
TX,Fort Worth
"
    );
}

#[test]
fn partition_drop_without_headers() {
    let wrk = Workdir::new("partition_drop_without_headers");
    wrk.create("in.csv", data(false));

    let mut cmd = wrk.command("partition");
    cmd.arg("--drop")
        .arg("--no-headers")
        .arg("1")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "CA.csv",
        "\
San Francisco
"
    );
    part_eq!(
        wrk,
        "NY.csv",
        "\
Manhattan
Buffalo
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
Dallas
Fort Worth
"
    );
}

#[test]
fn partition_into_new_directory() {
    let wrk = Workdir::new("partition_into_new_directory");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.arg("state").arg(wrk.path("out")).arg("in.csv");
    wrk.run(&mut cmd);

    assert!(wrk.path("out/NY.csv").exists());
}

#[test]
fn partition_custom_filename() {
    let wrk = Workdir::new("partition_custom_filename");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "state-{}-partition.csv"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    assert!(wrk.path("state-NY-partition.csv").exists());
}

#[test]
fn partition_custom_filename_with_directory() {
    let wrk = Workdir::new("partition_custom_filename_with_directory");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "{}/cities.csv"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    // This variation also helps with parallel partition jobs.
    assert!(wrk.path("NY/cities.csv").exists());
}

#[test]
fn partition_invalid_filename() {
    let wrk = Workdir::new("partition_invalid_filename");
    wrk.create("in.csv", data(true));

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "foo.csv"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.assert_err(&mut cmd);

    let mut cmd = wrk.command("partition");
    cmd.args(["--filename", "{}{}.csv"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.assert_err(&mut cmd);
}

fn tricky_data() -> Vec<Vec<String>> {
    vec![
        svec!["key", "explanation"],
        svec!["", "empty key"],
        svec!["empty", "the string empty"],
        svec!["unsafe _1$!,\"", "unsafe in shell"],
        svec!["collision", "ordinary value"],
        svec!["collision", "in same file"],
        svec!["coll ision", "collides"],
        svec!["collision!", "collides again"],
        svec!["collision_2", "collides with disambiguated"],
    ]
}

#[test]
fn partition_with_tricky_key_values() {
    let wrk = Workdir::new("partition_with_tricky_key_values");
    wrk.create("in.csv", tricky_data());

    let mut cmd = wrk.command("partition");
    cmd.arg("key").arg(wrk.path(".")).arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "empty.csv",
        "\
key,explanation
,empty key
"
    );
    part_eq!(
        wrk,
        "empty_1.csv",
        "\
key,explanation
empty,the string empty
"
    );
    part_eq!(
        wrk,
        "unsafe_1.csv",
        r#"key,explanation
"unsafe _1$!,""",unsafe in shell
"#
    );
    part_eq!(
        wrk,
        "collision.csv",
        "\
key,explanation
collision,ordinary value
collision,in same file
"
    );
    part_eq!(
        wrk,
        "collision_2.csv",
        "\
key,explanation
coll ision,collides
"
    );
    part_eq!(
        wrk,
        "collision_3.csv",
        "\
key,explanation
collision!,collides again
"
    );
    // Tricky! We didn't see this an input, but we did generate it as an
    // output already.
    part_eq!(
        wrk,
        "collision_2_4.csv",
        "\
key,explanation
collision_2,collides with disambiguated
"
    );
}

#[test]
fn partition_with_limit() {
    let wrk = Workdir::new("partition_with_limit");
    // Create data with more unique values than the limit
    let data = vec![
        svec!["state", "city"],
        svec!["NY", "Manhattan"],
        svec!["CA", "San Francisco"],
        svec!["TX", "Dallas"],
        svec!["FL", "Miami"],
        svec!["WA", "Seattle"],
        svec!["IL", "Chicago"],
        svec!["PA", "Philadelphia"],
        svec!["OH", "Columbus"],
        svec!["GA", "Atlanta"],
        svec!["NC", "Charlotte"],
        svec!["MI", "Detroit"],
        svec!["NJ", "Newark"],
        svec!["VA", "Virginia Beach"],
        svec!["NY", "Buffalo"],
        svec!["CA", "Los Angeles"],
        svec!["TX", "Houston"],
        svec!["FL", "Orlando"],
        svec!["WA", "Spokane"],
        svec!["IL", "Springfield"],
        svec!["PA", "Pittsburgh"],
        svec!["OH", "Cleveland"],
        svec!["GA", "Augusta"],
        svec!["NC", "Raleigh"],
        svec!["MI", "Grand Rapids"],
        svec!["NJ", "Jersey City"],
        svec!["VA", "Norfolk"],
    ];
    wrk.create("in.csv", data);

    let mut cmd = wrk.command("partition");
    cmd.args(["--limit", "7"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    // Should still create all 10 unique state files despite the limit
    assert!(wrk.path("NY.csv").exists());
    assert!(wrk.path("CA.csv").exists());
    assert!(wrk.path("TX.csv").exists());
    assert!(wrk.path("FL.csv").exists());
    assert!(wrk.path("WA.csv").exists());
    assert!(wrk.path("IL.csv").exists());
    assert!(wrk.path("PA.csv").exists());
    assert!(wrk.path("OH.csv").exists());
    assert!(wrk.path("GA.csv").exists());
    assert!(wrk.path("NC.csv").exists());
    assert!(wrk.path("MI.csv").exists());
    assert!(wrk.path("NJ.csv").exists());
    assert!(wrk.path("VA.csv").exists());

    // Verify that all rows for each state are properly grouped
    part_eq!(
        wrk,
        "NY.csv",
        "\
state,city
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "CA.csv",
        "\
state,city
CA,San Francisco
CA,Los Angeles
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
state,city
TX,Dallas
TX,Houston
"
    );
    part_eq!(
        wrk,
        "FL.csv",
        "\
state,city
FL,Miami
FL,Orlando
"
    );
    part_eq!(
        wrk,
        "WA.csv",
        "\
state,city
WA,Seattle
WA,Spokane
"
    );
    part_eq!(
        wrk,
        "IL.csv",
        "\
state,city
IL,Chicago
IL,Springfield
"
    );
    part_eq!(
        wrk,
        "PA.csv",
        "\
state,city
PA,Philadelphia
PA,Pittsburgh
"
    );
    part_eq!(
        wrk,
        "OH.csv",
        "\
state,city
OH,Columbus
OH,Cleveland
"
    );
    part_eq!(
        wrk,
        "GA.csv",
        "\
state,city
GA,Atlanta
GA,Augusta
"
    );
    part_eq!(
        wrk,
        "NC.csv",
        "\
state,city
NC,Charlotte
NC,Raleigh
"
    );
    part_eq!(
        wrk,
        "MI.csv",
        "\
state,city
MI,Detroit
MI,Grand Rapids
"
    );
    part_eq!(
        wrk,
        "NJ.csv",
        "\
state,city
NJ,Newark
NJ,Jersey City
"
    );
    part_eq!(
        wrk,
        "VA.csv",
        "\
state,city
VA,Virginia Beach
VA,Norfolk
"
    );
}

fn prefix_data() -> Vec<Vec<String>> {
    vec![
        svec!["state", "city"],
        svec!["MA", "Boston"],
        svec!["ME", "Portland"],
        svec!["M", "Too short"],
        svec!["CA", "San Francisco"],
        svec!["CO", "Denver"],
    ]
}

#[test]
fn partition_with_prefix_length() {
    let wrk = Workdir::new("partition_with_prefix_length");
    wrk.create("in.csv", prefix_data());

    let mut cmd = wrk.command("partition");
    cmd.args(["--prefix-length", "1"])
        .arg("state")
        .arg(wrk.path("."))
        .arg("in.csv");
    wrk.run(&mut cmd);

    part_eq!(
        wrk,
        "M.csv",
        "\
state,city
MA,Boston
ME,Portland
M,Too short
"
    );
    part_eq!(
        wrk,
        "C.csv",
        "\
state,city
CA,San Francisco
CO,Denver
"
    );
}

#[test]
fn partition_case_insensitive_collision() {
    let wrk = Workdir::new("partition_case_insensitive_collision");
    wrk.create(
        "in.csv",
        vec![
            svec!["state", "city"],
            svec!["CO", "Denver"],
            svec!["Co", "Boulder"],
            svec!["CO", "Colorado Springs"],
            svec!["Co", "Aspen"],
        ],
    );

    let mut cmd = wrk.command("partition");
    cmd.arg("state").arg(wrk.path(".")).arg("in.csv");
    wrk.run(&mut cmd);

    // On case-insensitive file systems, "CO" and "Co" would collide
    // The fix should create separate files: CO.csv and Co_1.csv
    part_eq!(
        wrk,
        "CO.csv",
        "\
state,city
CO,Denver
CO,Colorado Springs
"
    );
    part_eq!(
        wrk,
        "Co_1.csv",
        "\
state,city
Co,Boulder
Co,Aspen
"
    );
}

#[test]
fn partition_stdin_pipeline() {
    use std::io::Write;

    let wrk = Workdir::new("partition_stdin_pipeline");

    // Create test data with multiple unique values to trigger batching
    let data = vec![
        svec!["state", "city"],
        svec!["NY", "Manhattan"],
        svec!["CA", "San Francisco"],
        svec!["TX", "Dallas"],
        svec!["FL", "Miami"],
        svec!["WA", "Seattle"],
        svec!["IL", "Chicago"],
        svec!["PA", "Philadelphia"],
        svec!["NY", "Buffalo"],
        svec!["CA", "Los Angeles"],
        svec!["TX", "Houston"],
    ];

    // Create a temporary file with the CSV data
    wrk.create("stdin_data.csv", data);

    // Run the partition command with stdin input and --limit to trigger process_in_batches
    let mut cmd = wrk.command("partition");
    cmd.args(["--limit", "3"]).arg("state").arg(wrk.path("."));

    // Set up stdin for the command (no input file argument = implicit stdin)
    let stdin_data = wrk.read_to_string("stdin_data.csv").unwrap();
    cmd.stdin(std::process::Stdio::piped());

    // Run the command
    let mut child = cmd.spawn().unwrap();
    let mut stdin = child.stdin.take().unwrap();
    std::thread::spawn(move || {
        stdin.write_all(stdin_data.as_bytes()).unwrap();
    });

    // Wait for the command to complete
    let status = child.wait().unwrap();
    assert!(
        status.success(),
        "Partition should succeed with stdin input"
    );

    // Verify that partition files were created correctly
    assert!(wrk.path("NY.csv").exists());
    assert!(wrk.path("CA.csv").exists());
    assert!(wrk.path("TX.csv").exists());
    assert!(wrk.path("FL.csv").exists());
    assert!(wrk.path("WA.csv").exists());
    assert!(wrk.path("IL.csv").exists());
    assert!(wrk.path("PA.csv").exists());

    // Verify the content of partition files
    part_eq!(
        wrk,
        "NY.csv",
        "\
state,city
NY,Manhattan
NY,Buffalo
"
    );
    part_eq!(
        wrk,
        "CA.csv",
        "\
state,city
CA,San Francisco
CA,Los Angeles
"
    );
    part_eq!(
        wrk,
        "TX.csv",
        "\
state,city
TX,Dallas
TX,Houston
"
    );
    part_eq!(
        wrk,
        "FL.csv",
        "\
state,city
FL,Miami
"
    );
    part_eq!(
        wrk,
        "WA.csv",
        "\
state,city
WA,Seattle
"
    );
    part_eq!(
        wrk,
        "IL.csv",
        "\
state,city
IL,Chicago
"
    );
    part_eq!(
        wrk,
        "PA.csv",
        "\
state,city
PA,Philadelphia
"
    );
}
