use std::fs::File;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use supermachine::internal::{
    vmm::pool::WarmPool,
    RunOptions,
    VmResources,
};

#[derive(Clone, Debug)]
struct Row {
    iteration: usize,
    worker_restore_us: u128,
    reset_vsock_us: u128,
    remap_cow_us: u128,
    load_meta_us: u128,
    restore_snapshot_us: u128,
    ram_copy_us: u128,
    gic_restore_us: u128,
    vcpu_restore_us: u128,
    vtimer_offset_us: u128,
    mmio_restore_us: u128,
    listener_restore_us: u128,
    library_done_us: u128,
    first_200_us: i128,
    port: Option<u16>,
    error: Option<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut snapshot = None;
    let mut iterations = 20usize;
    let mut warmup = 1usize;
    let mut settle_ms = 10u64;
    let mut out_path = None;

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--snapshot" => snapshot = args.next(),
            "--iterations" => {
                iterations = args
                    .next()
                    .ok_or("--iterations requires a value")?
                    .parse()?
            }
            "--warmup" => warmup = args.next().ok_or("--warmup requires a value")?.parse()?,
            "--settle-ms" => {
                settle_ms = args.next().ok_or("--settle-ms requires a value")?.parse()?
            }
            "--out" => out_path = args.next(),
            _ => return Err(format!("unknown argument: {arg}").into()),
        }
    }

    let snapshot = snapshot.ok_or(
        "usage: warm_pool_bench --snapshot SNAPSHOT [--iterations N] [--warmup N] [--settle-ms MS] [--out PATH]",
    )?;
    let out_path = out_path.unwrap_or_else(|| "/tmp/supermachine-pool-warm-library.json".to_string());

    let resources = VmResources::from_snapshot(&snapshot).with_cow_restore(true);
    let pool = WarmPool::start(resources, RunOptions::default())?;

    for _ in 0..warmup {
        let t0 = Instant::now();
        let restored = pool.restore_timeout(&snapshot, Duration::from_secs(5))?;
        if let Some(port) = restored.host_port {
            let _ = wait_http_200(port, t0, Duration::from_secs(5));
        }
        if settle_ms > 0 {
            std::thread::sleep(Duration::from_millis(settle_ms));
        }
    }

    let mut rows = Vec::with_capacity(iterations);
    for i in 1..=iterations {
        let t0 = Instant::now();
        let restored = pool.restore_timeout(&snapshot, Duration::from_secs(5))?;
        let library_done_us = t0.elapsed().as_micros();
        let (first_200_us, error) = match restored.host_port {
            Some(port) => match wait_http_200(port, t0, Duration::from_secs(5)) {
                Some(us) => (us as i128, None),
                None => (-1, Some("first_200_timeout".to_string())),
            },
            None => (-1, Some("missing_host_port".to_string())),
        };

        println!(
            "iter={i:02} worker={}us done={:.3}ms first_200={:.3}ms reset={}us remap={}us meta={}us state={}us ram={}us gic={}us vcpu={}us vtimer={}us mmio={}us listen={}us port={}",
            restored.restore_us,
            library_done_us as f64 / 1000.0,
            first_200_us as f64 / 1000.0,
            restored.timings.reset_vsock_us,
            restored.timings.remap_cow_us,
            restored.timings.load_meta_us,
            restored.timings.restore_snapshot_us,
            restored.timings.ram_copy_us,
            restored.timings.gic_restore_us,
            restored.timings.vcpu_restore_us,
            restored.timings.vtimer_offset_us,
            restored.timings.mmio_restore_us,
            restored.timings.listener_restore_us,
            restored
                .host_port
                .map(|p| p.to_string())
                .unwrap_or_else(|| "?".to_string())
        );

        rows.push(Row {
            iteration: i,
            worker_restore_us: restored.restore_us,
            reset_vsock_us: restored.timings.reset_vsock_us,
            remap_cow_us: restored.timings.remap_cow_us,
            load_meta_us: restored.timings.load_meta_us,
            restore_snapshot_us: restored.timings.restore_snapshot_us,
            ram_copy_us: restored.timings.ram_copy_us,
            gic_restore_us: restored.timings.gic_restore_us,
            vcpu_restore_us: restored.timings.vcpu_restore_us,
            vtimer_offset_us: restored.timings.vtimer_offset_us,
            mmio_restore_us: restored.timings.mmio_restore_us,
            listener_restore_us: restored.timings.listener_restore_us,
            library_done_us,
            first_200_us,
            port: restored.host_port,
            error,
        });

        if settle_ms > 0 {
            std::thread::sleep(Duration::from_millis(settle_ms));
        }
    }

    let report = pool.shutdown()?;
    write_json(
        &out_path,
        &snapshot,
        iterations,
        warmup,
        settle_ms,
        report.warm_restores,
        &rows,
    )?;
    println!("wrote {out_path}");
    Ok(())
}

fn wait_http_200(port: u16, t0: Instant, timeout: Duration) -> Option<u128> {
    let deadline = Instant::now() + timeout;
    while Instant::now() < deadline {
        if http_probe(port) {
            return Some(t0.elapsed().as_micros());
        }
        std::thread::sleep(Duration::from_millis(2));
    }
    None
}

fn http_probe(port: u16) -> bool {
    let Ok(mut stream) = TcpStream::connect_timeout(
        &std::net::SocketAddr::from(([127, 0, 0, 1], port)),
        Duration::from_millis(200),
    ) else {
        return false;
    };
    let _ = stream.set_read_timeout(Some(Duration::from_millis(200)));
    let _ = stream.write_all(b"GET / HTTP/1.1\r\nHost: supermachine\r\nConnection: close\r\n\r\n");
    let mut buf = [0u8; 256];
    match stream.read(&mut buf) {
        Ok(n) => buf[..n].starts_with(b"HTTP/1.1 200") || buf[..n].starts_with(b"HTTP/1.0 200"),
        Err(_) => false,
    }
}

fn write_json(
    out_path: &str,
    snapshot: &str,
    iterations: usize,
    warmup: usize,
    settle_ms: u64,
    warm_restores: u64,
    rows: &[Row],
) -> std::io::Result<()> {
    let mut out = File::create(out_path)?;
    writeln!(out, "{{")?;
    writeln!(out, "  \"benchmark\": \"supermachine_pool_warm_library\",")?;
    writeln!(out, "  \"generated_at\": \"{}\",", generated_at())?;
    writeln!(out, "  \"snapshot\": \"{}\",", json_escape(snapshot))?;
    writeln!(out, "  \"iterations\": {iterations},")?;
    writeln!(out, "  \"warmup\": {warmup},")?;
    writeln!(out, "  \"settle_ms\": {settle_ms},")?;
    writeln!(out, "  \"warm_restores\": {warm_restores},")?;
    write_summary(&mut out, rows)?;
    writeln!(out, "  \"rows\": [")?;
    for (idx, row) in rows.iter().enumerate() {
        if idx > 0 {
            writeln!(out, ",")?;
        }
        write!(
            out,
            "    {{\"iteration\":{},\"worker_restore_us\":{},\"reset_vsock_us\":{},\"remap_cow_us\":{},\"load_meta_us\":{},\"restore_snapshot_us\":{},\"ram_copy_us\":{},\"gic_restore_us\":{},\"vcpu_restore_us\":{},\"vtimer_offset_us\":{},\"mmio_restore_us\":{},\"listener_restore_us\":{},\"library_done_us\":{},\"first_200_us\":{},\"port\":{}",
            row.iteration,
            row.worker_restore_us,
            row.reset_vsock_us,
            row.remap_cow_us,
            row.load_meta_us,
            row.restore_snapshot_us,
            row.ram_copy_us,
            row.gic_restore_us,
            row.vcpu_restore_us,
            row.vtimer_offset_us,
            row.mmio_restore_us,
            row.listener_restore_us,
            row.library_done_us,
            row.first_200_us,
            row.port
                .map(|p| p.to_string())
                .unwrap_or_else(|| "null".to_string())
        )?;
        if let Some(error) = row.error.as_ref() {
            write!(out, ",\"error\":\"{}\"", json_escape(error))?;
        }
        write!(out, "}}")?;
    }
    writeln!(out, "\n  ]")?;
    writeln!(out, "}}")?;
    Ok(())
}

fn write_summary(out: &mut File, rows: &[Row]) -> std::io::Result<()> {
    let worker: Vec<u128> = rows.iter().map(|r| r.worker_restore_us).collect();
    let done: Vec<u128> = rows.iter().map(|r| r.library_done_us).collect();
    let first: Vec<u128> = rows
        .iter()
        .filter_map(|r| (r.first_200_us >= 0).then_some(r.first_200_us as u128))
        .collect();

    writeln!(out, "  \"summary\": {{")?;
    write_percentiles(out, "worker_restore_us", &worker, true)?;
    write_percentiles(out, "library_done_us", &done, true)?;
    write_percentiles(out, "first_200_us", &first, false)?;
    writeln!(out, ",")?;
    writeln!(out, "    \"phase_us\": {{")?;
    write_percentiles(
        out,
        "reset_vsock_us",
        &rows.iter().map(|r| r.reset_vsock_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "remap_cow_us",
        &rows.iter().map(|r| r.remap_cow_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "load_meta_us",
        &rows.iter().map(|r| r.load_meta_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "restore_snapshot_us",
        &rows
            .iter()
            .map(|r| r.restore_snapshot_us)
            .collect::<Vec<_>>(),
        false,
    )?;
    writeln!(out, ",")?;
    write_percentiles(
        out,
        "ram_copy_us",
        &rows.iter().map(|r| r.ram_copy_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "gic_restore_us",
        &rows.iter().map(|r| r.gic_restore_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "vcpu_restore_us",
        &rows.iter().map(|r| r.vcpu_restore_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "vtimer_offset_us",
        &rows.iter().map(|r| r.vtimer_offset_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "mmio_restore_us",
        &rows.iter().map(|r| r.mmio_restore_us).collect::<Vec<_>>(),
        true,
    )?;
    write_percentiles(
        out,
        "listener_restore_us",
        &rows
            .iter()
            .map(|r| r.listener_restore_us)
            .collect::<Vec<_>>(),
        false,
    )?;
    writeln!(out, "    }}")?;
    writeln!(out, "  }},")?;
    Ok(())
}

fn write_percentiles(
    out: &mut File,
    key: &str,
    values: &[u128],
    trailing_comma: bool,
) -> std::io::Result<()> {
    let comma = if trailing_comma { "," } else { "" };
    writeln!(
        out,
        "    \"{key}\": {{\"min\":{},\"p50\":{},\"p95\":{},\"max\":{}}}{comma}",
        percentile(values, 0.0),
        percentile(values, 0.5),
        percentile(values, 0.95),
        percentile(values, 1.0)
    )
}

fn percentile(values: &[u128], q: f64) -> u128 {
    if values.is_empty() {
        return 0;
    }
    let mut sorted = values.to_vec();
    sorted.sort_unstable();
    let idx = ((sorted.len() - 1) as f64 * q + 0.5) as usize;
    sorted[idx]
}

fn generated_at() -> String {
    let secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    secs.to_string()
}

fn json_escape(s: &str) -> String {
    s.replace('\\', "\\\\").replace('"', "\\\"")
}
