use std::borrow::ToOwned;

use crate::workdir::Workdir;

macro_rules! part_eq {
    ($wrk:expr, $path:expr, $expected:expr) => {
        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", "Manhatten"],
        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,Manhatten
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
Manhatten
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,Manhatten
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",
        "\
Manhatten
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
"
    );
}

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
"
    );
}
