mod settings;
use std::io::{Result, ErrorKind};
use std::io::SeekFrom;
use std::time::Instant;
use rand::Rng;
use aws_sdk_s3::Client;
use human_bytes::human_bytes;
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use hyperfile::file::hyper::Hyper;
use hyperfile::file::tokio_wrapper::HyperFileTokio;
use hyperfile::file::flags::FileFlags;
use hyperfile::file::mode::FileMode;
use hyperfile::config::{
    HyperFileConfigBuilder, HyperFileConfig,
    HyperFileRuntimeConfig, HyperFileMetaConfig,
};
use hyperfile::staging::config::StagingConfig;
#[cfg(feature = "wal")]
use hyperfile::wal::config::HyperFileWalConfig;
use settings::*;

async fn random_write(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>, write_zero: bool) -> Result<()> {
    let rand_data_offset = rand::random_range(DEFAULT_RANDOM_WRITE_OFFSET);
    let mut rand_data_bytes = rand::random_range(DEFAULT_RANDOM_WRITE_BYTES);
    if rand_data_offset + rand_data_bytes >= DEFAULT_MAX_FILE_SIZE {
        rand_data_bytes = DEFAULT_MAX_FILE_SIZE - rand_data_offset;
    }
    let total_len = rand_data_offset + rand_data_bytes;
    // only resize buffer for data extend case
    if total_len > data.len() {
        data.resize(total_len, 0);
    }
    if write_zero {
        print!("random write zero offset {} bytes {} ..", rand_data_offset, rand_data_bytes);
    } else {
        print!("random write offset {} bytes {} ..", rand_data_offset, rand_data_bytes);
    }

    let flags = FileFlags::rdwr();
    let mode = FileMode::default_file();
    let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
    // write content
    let filled_buf = &mut data[rand_data_offset..total_len];
    if write_zero {
        filled_buf.fill(0);
    } else {
        rand::fill(filled_buf);
    }
    file.seek(SeekFrom::Start(rand_data_offset as u64)).await?;
    let write_bytes = if write_zero {
        file.write_zero(rand_data_bytes).await?
    } else {
        file.write(filled_buf).await?
    };
    assert!(write_bytes == rand_data_bytes);
    file.flush().await?;
    let last_cno = file.last_cno().await?;
    file.shutdown().await?;
    println!(". Done with last cno {}", last_cno);
    Ok(())
}

async fn random_truncate(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>) -> Result<()> {
    let curr_len = data.len();
    // ratio to truncate only one byte
    let new_file_len = if curr_len == 0 {
        rand::random_range(0..DEFAULT_MAX_FILE_SIZE)
    } else if rand::rng().random_ratio(2, 10) {
        curr_len - 1
    } else if rand::rng().random_ratio(2, 10) {
        curr_len + 1
    } else {
        rand::random_range(0..DEFAULT_MAX_FILE_SIZE)
    };
    print!("random truncate file to new size {} ..", new_file_len);
    // open file for read
    let flags = FileFlags::rdwr();
    let mode = FileMode::default_file();
    let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
    // get stat
    let stat = file.metadata().await?;
    assert!(stat.st_size == data.len() as i64);
    if new_file_len > stat.st_size as usize {
        print!(".[EXPEND]..");
    } else if new_file_len < stat.st_size as usize {
        print!(".[SHRINK]..");
    } else {
        print!(".[NO CHANGE]..");
    }

    // extend file by truncate
    file.set_len(new_file_len as u64).await?;
    file.flush().await?;
    let last_cno = file.last_cno().await?;
    data.resize(new_file_len, 0);
    file.shutdown().await?;
    println!(". Done with last cno {}", last_cno);
    Ok(())
}

async fn read_check(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>) -> Result<()> {
    let total_bytes = data.len();
    // open file for read
    let flags = FileFlags::rdonly();
    let mode = FileMode::default_file();
    let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
    // get stat
    let stat = file.metadata().await?;
    assert!(stat.st_size == total_bytes as i64);

    let mut buf = Vec::new();
    buf.resize(total_bytes, 0);
    file.seek(SeekFrom::Start(0)).await?;
    let start = Instant::now();
    let read_bytes = file.read(&mut buf).await?;
    let nanos = start.elapsed().as_nanos();
    let md5_data = md5::compute(&data);
    let md5_file = md5::compute(&buf);
    println!("real data in memory md5: {:?} - read data from file md5: {:?}", md5_data, md5_file);
    assert!(read_bytes == total_bytes);
    let throughput = (read_bytes as u128 * 1000_000_000) / nanos;
    println!("bytes read {} throughput {}/s", human_bytes(read_bytes as f64), human_bytes(throughput as f64));
    if md5_data != md5_file {
        panic!("md5 check failed");
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();
	let uri = std::env::var("HYPERFILE_STAGING_ROOT_URI").unwrap();
    #[cfg(feature = "wal")]
	let wal_uri = std::env::var("HYPERFILE_WAL_ROOT_URI").unwrap();

    // build s3 client
	let config = aws_config::load_from_env().await;
    let client = aws_sdk_s3::Client::new(&config);

    // clear any existing on staging
    let _ = Hyper::fs_unlink(&client, &uri).await
                        .or_else(|e| {
                            if e.kind() == ErrorKind::NotFound {
                                return Ok(());
                            }
                            Err(e)
                        })?;

    let mut data: Vec<u8> = Vec::new();

    let meta_config = HyperFileMetaConfig::default();
    let staging_config = StagingConfig::new_s3_uri(&uri, None);
    let runtime_config = HyperFileRuntimeConfig::default_large();
    #[cfg(not(feature = "wal"))]
    let file_config = HyperFileConfigBuilder::new()
                        .with_meta_config(&meta_config)
                        .with_staging_config(&staging_config)
                        .with_runtime_config(&runtime_config)
                        .build();
    #[cfg(feature = "wal")]
    let wal_config = HyperFileWalConfig::new(&wal_uri);
    #[cfg(feature = "wal")]
    let file_config = HyperFileConfigBuilder::new()
                        .with_meta_config(&meta_config)
                        .with_staging_config(&staging_config)
                        .with_runtime_config(&runtime_config)
                        .with_wal_config(&wal_config)
                        .build();

    for _ in 0..NUM_ITER {
        if rand::rng().random_ratio(5, 10) {
            random_write(&client, &file_config, &mut data, false).await?;
        } else if rand::rng().random_ratio(2, 10) {
            random_write(&client, &file_config, &mut data, true).await?;
        } else {
            random_truncate(&client, &file_config, &mut data).await?;
        }
    }

    println!("Final read check...");
    read_check(&client, &file_config, &mut data).await?;

    Ok(())
}
