use clap::{Parser, ValueEnum};
use fwf::{
    Alignment, FwReader, FwReaderField, FwReaderSchema, FwWriter, FwWriterField, FwWriterSchema,
    RecordSeparator, TruncationPolicy,
};
use std::fs::File;
use std::path::{Path, PathBuf};
use std::time::Instant;
use tempfile::NamedTempFile;

const RECORD_WIDTH: usize = 130;

#[derive(Parser)]
#[command(about = "Profile fwf read/write throughput")]
struct Cli {
    /// Which phase to run
    #[arg(long, short, default_value = "both")]
    phase: Phase,

    /// Number of records to generate
    #[arg(long, short, default_value_t = 1_000_000)]
    records: usize,

    /// Path to use for the data file. If omitted, a temporary file is used.
    #[arg(long, short)]
    file: Option<PathBuf>,
}

#[derive(Clone, PartialEq, ValueEnum)]
enum Phase {
    Write,
    Read,
    Both,
    #[cfg(feature = "tokio")]
    AsyncWrite,
    #[cfg(feature = "tokio")]
    AsyncRead,
    #[cfg(feature = "tokio")]
    AsyncBoth,
    #[cfg(feature = "tokio")]
    All,
}

fn writer_schema() -> FwWriterSchema {
    FwWriterSchema::builder()
        .field(FwWriterField::builder(12).build())
        .field(FwWriterField::builder(30).build())
        .field(
            FwWriterField::builder(12)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwWriterField::builder(10).build())
        .field(FwWriterField::builder(4).build())
        .field(
            FwWriterField::builder(8)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwWriterField::builder(6).build())
        .field(FwWriterField::builder(16).build())
        .field(
            FwWriterField::builder(12)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwWriterField::builder(20).build())
        .separator(RecordSeparator::Lf)
        .truncation(TruncationPolicy::Right)
        .build()
}

fn reader_schema() -> FwReaderSchema {
    FwReaderSchema::builder(RECORD_WIDTH)
        .field(FwReaderField::builder(0..12).build())
        .field(FwReaderField::builder(12..42).build())
        .field(
            FwReaderField::builder(42..54)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwReaderField::builder(54..64).build())
        .field(FwReaderField::builder(64..68).build())
        .field(
            FwReaderField::builder(68..76)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwReaderField::builder(76..82).build())
        .field(FwReaderField::builder(82..98).build())
        .field(
            FwReaderField::builder(98..110)
                .alignment(Alignment::Right)
                .padding(b'0')
                .build(),
        )
        .field(FwReaderField::builder(110..130).build())
        .separator(RecordSeparator::Lf)
        .build()
        .unwrap()
}

fn build_template_record() -> Vec<&'static [u8]> {
    vec![
        b"ACC-00123456",
        b"John Q. Public",
        b"123456.78",
        b"2024-01-15",
        b"SALE",
        b"42",
        b"POSTED",
        b"TXN-20240115-01",
        b"98765.43",
        b"Monthly payment",
    ]
}

fn run_write(path: &Path, record_count: usize) -> (std::time::Duration, u64) {
    let template = build_template_record();
    let file = File::create(path).unwrap();

    let start = Instant::now();
    let mut writer = FwWriter::options(writer_schema()).from_file(file);
    for _ in 0..record_count {
        writer.write_byte_record(&template).unwrap();
    }
    writer.flush().unwrap();
    let elapsed = start.elapsed();

    let file_size = std::fs::metadata(path).unwrap().len();
    (elapsed, file_size)
}

fn run_read(path: &Path) -> (std::time::Duration, u64, u64) {
    let file_size = std::fs::metadata(path).unwrap().len();
    let file = File::open(path).unwrap();

    let start = Instant::now();
    let mut reader = FwReader::options(reader_schema()).from_file(file);
    let mut record_count: u64 = 0;
    while reader.read_byte_record().unwrap().is_some() {
        record_count += 1;
    }
    let elapsed = start.elapsed();

    (elapsed, record_count, file_size)
}

#[cfg(feature = "tokio")]
async fn run_async_write(path: &Path, record_count: usize) -> (std::time::Duration, u64) {
    let template = build_template_record();
    let file = tokio::fs::File::create(path).await.unwrap();

    let start = Instant::now();
    let mut writer = FwWriter::options(writer_schema()).from_tokio_file(file);
    for _ in 0..record_count {
        writer.write_byte_record(&template).await.unwrap();
    }
    writer.flush().await.unwrap();
    let elapsed = start.elapsed();

    let file_size = std::fs::metadata(path).unwrap().len();
    (elapsed, file_size)
}

#[cfg(feature = "tokio")]
async fn run_async_read(path: &Path) -> (std::time::Duration, u64, u64) {
    let file_size = std::fs::metadata(path).unwrap().len();
    let file = tokio::fs::File::open(path).await.unwrap();

    let start = Instant::now();
    let mut reader = FwReader::options(reader_schema()).from_tokio_file(file);
    let mut record_count: u64 = 0;
    while reader.read_byte_record().await.unwrap().is_some() {
        record_count += 1;
    }
    let elapsed = start.elapsed();

    (elapsed, record_count, file_size)
}

#[allow(clippy::cast_precision_loss)]
fn print_row(phase: &str, records: u64, size_bytes: u64, elapsed: std::time::Duration) {
    let size_mb = size_bytes as f64 / (1024.0 * 1024.0);
    let secs = elapsed.as_secs_f64();
    let records_per_sec = records as f64 / secs;
    let mb_per_sec = size_mb / secs;
    println!(
        "{phase:<12} {records:>12} {size_mb:>10.1} {secs:>10.3} {records_per_sec:>12.0} {mb_per_sec:>10.1}",
    );
}

/// Returns true for the `All` variant (only exists with tokio).
fn cfg_all(phase: &Phase) -> bool {
    #[cfg(feature = "tokio")]
    if *phase == Phase::All {
        return true;
    }
    false
}

/// Returns true if any read phase is selected that needs a pre-existing file.
fn needs_seed_file(phase: &Phase) -> bool {
    let sync_read = matches!(phase, Phase::Read | Phase::Both);
    #[cfg(feature = "tokio")]
    let async_read = matches!(phase, Phase::AsyncRead | Phase::AsyncBoth | Phase::All);
    #[cfg(not(feature = "tokio"))]
    let async_read = false;
    sync_read || async_read
}

fn main() {
    let cli = Cli::parse();

    println!("Records: {}", cli.records);
    println!();
    println!(
        "{:<12} {:>12} {:>10} {:>10} {:>12} {:>10}",
        "Phase", "Records", "Size (MB)", "Time (s)", "Records/s", "MB/s"
    );
    println!("{}", "-".repeat(70));

    // Use the provided path or fall back to a temporary file.
    let (_tmp, file_path) = if let Some(path) = cli.file.clone() {
        (None, path)
    } else {
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        (Some(tmp), path)
    };

    let run_sync_write = matches!(cli.phase, Phase::Write | Phase::Both) || cfg_all(&cli.phase);
    let run_sync_read = matches!(cli.phase, Phase::Read | Phase::Both) || cfg_all(&cli.phase);

    // Ensure a file exists before any read-only phase.
    if needs_seed_file(&cli.phase) && !run_sync_write && !file_path.exists() {
        run_write(&file_path, cli.records);
    }

    if run_sync_write {
        let (elapsed, file_size) = run_write(&file_path, cli.records);
        print_row("Write", cli.records as u64, file_size, elapsed);
    }

    if run_sync_read {
        let (elapsed, records_read, file_size) = run_read(&file_path);
        print_row("Read", records_read, file_size, elapsed);
    }

    #[cfg(feature = "tokio")]
    run_async_phases(&cli, &file_path, run_sync_write);
}

#[cfg(feature = "tokio")]
fn run_async_phases(cli: &Cli, file_path: &Path, sync_wrote: bool) {
    let should_write = matches!(cli.phase, Phase::AsyncWrite | Phase::AsyncBoth | Phase::All);
    let should_read = matches!(cli.phase, Phase::AsyncRead | Phase::AsyncBoth | Phase::All);

    if !should_write && !should_read {
        return;
    }

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    rt.block_on(async {
        if should_write {
            let (elapsed, file_size) = run_async_write(file_path, cli.records).await;
            print_row("AsyncWrite", cli.records as u64, file_size, elapsed);
        }

        if should_read {
            if !sync_wrote && !should_write && !file_path.exists() {
                run_write(file_path, cli.records);
            }
            let (elapsed, records_read, file_size) = run_async_read(file_path).await;
            print_row("AsyncRead", records_read, file_size, elapsed);
        }
    });
}
