use std::time::Duration;

use bytesize::ByteSize;
use criterion::{
    black_box, criterion_group, criterion_main,
    measurement::{Measurement, ValueFormatter},
    Criterion, Throughput,
};
use jemalloc_ctl::{epoch, stats};

#[path = "../tests/automerge_paper_helpers.rs"]
mod automerge_helpers;

use automerge_helpers::load_automerge_paper_from_json;
use jmbl::stream_format::OpSerializer;

// fn iter_measuring_memory<T, C: Fn() -> T>(creator: C) -> impl Fn(u64) -> u64 {
//     move |count: u64| {
//         let mut sum: u64 = 0;

//         println!("count {}", count);

//         for i in 0..count {
//             epoch::advance().unwrap();
//             let start = stats::allocated::read().unwrap();

//             let val = creator();

//             epoch::advance().unwrap();
//             let end = stats::allocated::read().unwrap();

//             black_box(val);
//             // info!("{}", ByteSize((end - start) as u64));
//             sum += (end - start) as u64;
//         }

//         sum
//     }
// }

#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

fn automerge_paper_memory(c: &mut Criterion<AllocationMeasurement>) {
    let (doc, list_id, ops) = load_automerge_paper_from_json(true);
    // info!("{:#?}", doc.apply_op_diff(&diff).mem_use());
    // panic!("");

    {
        let mut group = c.benchmark_group("automerge_paper_memory_op_stream");
        group.throughput(Throughput::Elements(ops.len() as u64));
        group.sample_size(10);

        // group.bench_function("clone_op_stream", |b| {
        //     b.iter_custom(iter_measuring_memory(|| {
        //         // sleep a little to prevent criterion from running lots of samples and draining system memory
        //         std::thread::sleep(Duration::from_millis(100));
        //         op_stream.clone()
        //     }))
        // });

        group.bench_function("serialize_op_stream", |b| {
            b.iter_custom(|count| {
                let mut sum = 0;
                for _ in 0..count {
                    // sleep a little to prevent criterion from running lots of samples and draining system memory
                    std::thread::sleep(Duration::from_millis(100));
                    let mut vec = Vec::new();
                    let mut op_serializer = OpSerializer::new();
                    for op in &ops {
                        litl::to_writer(&mut vec, &op_serializer.serialize(op)).unwrap();
                    }

                    let compressed = lil::compress_bytes(&vec);

                    let len = compressed.len();

                    black_box(compressed);
                    sum += len as u64
                }
                sum
            })
        });
    }

    // {
    //     let mut group = c.benchmark_group("automerge_paper_memory_doc");
    //     group.throughput(Throughput::Elements(diff.ops[0].ops.len() as u64));
    //     group.sample_size(10);
    //     group.sampling_mode(SamplingMode::Flat);

    //     group.bench_function("doc_with_ops_applied", |b| {
    //         b.iter_custom(iter_measuring_memory(|| doc.apply_op_diff(&diff)))
    //     });

    // }
}

fn automerge_paper_runtime(c: &mut Criterion) {
    let (doc, list_id, ops) = load_automerge_paper_from_json(true);

    let mut group = c.benchmark_group("automerge_paper_runtime");
    group.throughput(Throughput::Elements(ops.len() as u64));

    group.bench_function("apply_op_diff", |b| {
        b.iter(|| {
            let doc2 = black_box(doc.with_ops_applied(&ops));
        })
    });

    let doc2 = doc.with_ops_applied(&ops);

    // info!("Mem use of doc {:#?}", doc2.mem_use());

    group.bench_function("get_final_output", |b| {
        b.iter(|| {
            let output = black_box(
                doc2.get_root()
                    .if_list()
                    .unwrap()
                    .val()
                    .iter()
                    .map(|elem| elem.if_plain().unwrap().as_str().unwrap())
                    .collect::<String>(),
            );
        })
    });

    group.bench_function("serialize_op_stream", |b| {
        b.iter(|| {
            let mut vec = Vec::new();
            let mut op_serializer = OpSerializer::new();
            for op in &ops {
                litl::to_writer(&mut vec, &op_serializer.serialize(op)).unwrap();
            }
            black_box(vec);
        })
    });

    // info!(
    //     "{}",
    //     doc2.get_root()
    //         .expect_list()
    //         .val()
    //         .iter()
    //         .map(|elem| elem.expect_plain().expect_str())
    //         .collect::<String>()
    // );
}

struct AllocationMeasurement;

impl Measurement for AllocationMeasurement {
    type Intermediate = usize;

    type Value = u64;

    fn start(&self) -> Self::Intermediate {
        epoch::advance().unwrap();
        let start = stats::allocated::read().unwrap();
        println!("Start {}", ByteSize(start as u64));
        start
    }

    fn end(&self, start: Self::Intermediate) -> Self::Value {
        epoch::advance().unwrap();
        let end = stats::allocated::read().unwrap();
        println!("End {}", ByteSize(end as u64));

        (end - start) as u64
    }

    fn add(&self, v1: &Self::Value, v2: &Self::Value) -> Self::Value {
        v1 + v2
    }

    fn zero(&self) -> Self::Value {
        0
    }

    fn to_f64(&self, value: &Self::Value) -> f64 {
        *value as f64
    }

    fn formatter(&self) -> &dyn criterion::measurement::ValueFormatter {
        &AllocationFormatter
    }
}

struct AllocationFormatter;

impl ValueFormatter for AllocationFormatter {
    fn format_value(&self, value: f64) -> String {
        ByteSize(value as u64).to_string()
    }
    fn format_throughput(&self, throughput: &Throughput, value: f64) -> String {
        match throughput {
            Throughput::Elements(elems) => {
                format!("{}/elem", ByteSize((value / *elems as f64) as u64),)
            }
            Throughput::Bytes(bytes) => {
                todo!()
            }
        }
    }
    fn scale_values(&self, typical_value: f64, values: &mut [f64]) -> &'static str {
        let exp_1000 = (typical_value.log10() / 3.0).floor() as usize;
        if exp_1000 == 0 {
            return "B";
        };
        if exp_1000 == 1 {
            for value in values {
                *value /= 1_000.0
            }
            return "KB";
        }
        if exp_1000 == 2 {
            for value in values {
                *value /= 1_000_000.0
            }
            return "MB";
        }
        if exp_1000 == 3 {
            for value in values {
                *value /= 1_000_000_000.0
            }
            return "GB";
        }
        for value in values {
            *value /= 1_000_000_000_000.0;
        }
        "TB"
    }
    fn scale_throughputs(
        &self,
        typical_value: f64,
        throughput: &Throughput,
        values: &mut [f64],
    ) -> &'static str {
        self.scale_values(typical_value, values)
    }
    fn scale_for_machines(&self, values: &mut [f64]) -> &'static str {
        "B"
    }
}

criterion_group! {
    name = memory;
    config = Criterion::default().with_measurement(AllocationMeasurement).sample_size(10);
    targets = automerge_paper_memory
}

criterion_group! {
    name = runtime;
    config = Criterion::default().sample_size(10);
    targets = automerge_paper_runtime
}

criterion_main!(runtime, memory);