//! Generates Hackers-themed benchmark proto binaries using the roto builder API.
//!
//! Run this once to create the data files that `hackers_bench` loads:
//!
//! ```sh
//! cargo run --release --bin gen_bench_data -- --preset tiny
//! cargo run --release --bin gen_bench_data -- --preset small
//! cargo run --release --bin gen_bench_data -- --preset medium
//! cargo run --release --bin gen_bench_data -- --preset large
//! cargo run --release --bin gen_bench_data -- --preset huge
//!
//! # Custom: ~50 MB — 500 ops × 100 KB stolen_data each
//! cargo run --release --bin gen_bench_data -- \
//!     --ops 500 --stolen-kb 100 --output data/bench/50mb.pb
//! ```
//!
//! Files land in `data/bench/` by default.

use clap::Parser;
use roto_benches::hackers::{
    CampaignBuilder, ConnectionBuilder, HackerBuilder, OperationBuilder, ToolBuilder, WormBuilder,
};
use std::io::{self, Write};
use std::path::Path;

// =============================================================================
// CLI
// =============================================================================

#[derive(Parser)]
#[command(
    name = "gen_bench_data",
    about = "Generate Hackers-themed proto binaries for benchmarks"
)]
struct Args {
    /// Output file.  Defaults to data/bench/<preset>.pb when --preset is used.
    #[arg(short, long)]
    output: Option<String>,

    /// Named size preset: tiny | small | medium | large | huge
    #[arg(short, long)]
    preset: Option<String>,

    /// Number of Operation messages.
    #[arg(long, default_value_t = 100)]
    ops: usize,

    /// Kilobytes of random stolen_data padding per Operation.
    #[arg(long, default_value_t = 0)]
    stolen_kb: usize,

    /// Hacker crew members per Operation.
    #[arg(long, default_value_t = 3)]
    crew: usize,

    /// RNG seed.
    #[arg(long, default_value_t = 42)]
    seed: u64,
}

// =============================================================================
// Minimal xorshift64 RNG
// =============================================================================

struct Rng(u64);

impl Rng {
    fn new(seed: u64) -> Self {
        Self(if seed == 0 { 0xdeadbeef_cafebabe } else { seed })
    }
    fn next(&mut self) -> u64 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        self.0
    }
    fn below(&mut self, n: usize) -> usize {
        (self.next() as usize) % n
    }
    fn range(&mut self, lo: u64, hi: u64) -> u64 {
        lo + self.next() % (hi - lo)
    }
    fn bool(&mut self) -> bool {
        self.next() & 1 == 0
    }
    fn pick<'a, T>(&mut self, s: &'a [T]) -> &'a T {
        &s[self.below(s.len())]
    }
    fn bytes(&mut self, n: usize) -> Vec<u8> {
        (0..n).map(|_| self.next() as u8).collect()
    }
}

// =============================================================================
// Flavour text
// =============================================================================

const HANDLES: &[&str] = &[
    "Zero Cool",
    "Acid Burn",
    "Phantom Phreak",
    "Cereal Killer",
    "Lord Nikon",
    "The Plague",
    "Crash Override",
];
const REAL_NAMES: &[&str] = &[
    "Dade Murphy",
    "Kate Libby",
    "Richard Gill",
    "Emmanuel Goldstein",
    "Paul Cook",
    "Eugene Belford",
];
const EXPLOITS: &[&str] = &[
    "buffer overflow",
    "stack smash",
    "heap spray",
    "race condition",
    "SQL injection",
    "CSRF",
    "XSS",
    "RCE",
    "privesc",
    "kernel panic",
];
const TOOL_NAMES: &[&str] = &[
    "nmap",
    "metasploit",
    "netcat",
    "tcpdump",
    "Wireshark",
    "sqlmap",
    "Burp Suite",
    "hashcat",
    "john",
];
const CORPS: &[&str] = &[
    "ELLINGSON MINERAL",
    "Cyberdelia",
    "The Gibson",
    "Prism BBS",
    "Elite BBS",
    "CRT Systems",
];
const LOG_LINES: &[&str] = &[
    "Hack the planet!",
    "Mess with the best, die like the rest.",
    "I'm in.",
    "They're tracing us.",
    "It's a Unix system!  I know this!",
    "You are elite.",
    "Garbage file accessed.",
];
const WORM_TARGETS: &[&str] = &[
    "ELLINGSON MINERAL",
    "Cyberdelia",
    "The Gibson",
    "Prism",
    "CRT BBS",
];
const OP_NAMES: &[&str] = &[
    "OPERATION HACK THE PLANET",
    "GIBSON BREACH",
    "ELLINGSON STING",
    "WORM UNLEASHED",
    "PHANTOM ACCESS",
    "DA VINCI",
];

// =============================================================================
// Message builders using the generated roto::hackers API
// =============================================================================

fn gen_tool(rng: &mut Rng) -> Vec<u8> {
    let payload_n = rng.range(8, 64) as usize;
    let payload = rng.bytes(payload_n);
    let version = format!("{}.{}", rng.range(1, 9), rng.range(0, 99));
    let mut buf = vec![0u8; 512];
    ToolBuilder::builder(&mut buf)
        .name(*rng.pick(TOOL_NAMES))
        .unwrap()
        .version(&version)
        .unwrap()
        .payload(&payload)
        .unwrap()
        .is_active(rng.bool() as u64)
        .unwrap()
        .exploit_count(rng.range(0, 50) as i32)
        .unwrap()
        .finish()
        .unwrap()
        .to_vec()
}

fn gen_connection(rng: &mut Rng) -> Vec<u8> {
    let host = format!("192.168.{}.{}", rng.range(1, 254), rng.range(1, 254));
    let session_key = rng.bytes(32);
    let mut buf = vec![0u8; 256];
    ConnectionBuilder::builder(&mut buf)
        .host(&host)
        .unwrap()
        .port(rng.range(1024, 65535) as i32)
        .unwrap()
        .encrypted(rng.bool() as u64)
        .unwrap()
        .bandwidth_bps(rng.range(1200, 1_000_000_000))
        .unwrap()
        .session_key(&session_key)
        .unwrap()
        .finish()
        .unwrap()
        .to_vec()
}

fn gen_hacker(rng: &mut Rng) -> Vec<u8> {
    let tools: Vec<Vec<u8>> = (0..rng.range(1, 4)).map(|_| gen_tool(rng)).collect();
    let connection = gen_connection(rng);
    // Float (skill_level) is written as raw bytes by the generated builder
    let skill_bits = (rng.range(10, 100) as f32 / 10.0).to_bits().to_le_bytes();
    let handle = *rng.pick(HANDLES);
    let real_name = *rng.pick(REAL_NAMES);
    let n_exploits = rng.range(2, 5) as usize;
    let exploits: Vec<&str> = (0..n_exploits).map(|_| *rng.pick(EXPLOITS)).collect();
    let crew_id = rng.next();

    let mut buf = vec![0u8; 8 * 1024];
    let mut b = HackerBuilder::builder(&mut buf)
        .handle(handle)
        .unwrap()
        .real_name(real_name)
        .unwrap()
        .age(rng.range(16, 35) as i32)
        .unwrap()
        .skill_level(&skill_bits)
        .unwrap()
        .is_elite(rng.bool() as u64)
        .unwrap()
        .crew_id(crew_id)
        .unwrap();
    for e in &exploits {
        b = b.exploits(e).unwrap();
    }
    for t in &tools {
        b = b.tools(t).unwrap();
    }
    b.active_connection(&connection)
        .unwrap()
        .finish()
        .unwrap()
        .to_vec()
}

fn gen_worm(rng: &mut Rng) -> Vec<u8> {
    let payload = rng.bytes(64);
    let name = format!("da_vinci.{}", rng.range(1, 99));
    let n_targets = rng.range(1, 4) as usize;
    let targets: Vec<&str> = (0..n_targets).map(|_| *rng.pick(WORM_TARGETS)).collect();

    let mut buf = vec![0u8; 1024];
    let mut b = WormBuilder::builder(&mut buf)
        .name(&name)
        .unwrap()
        .variant(rng.range(1, 5) as i32)
        .unwrap()
        .size_bytes(rng.range(1024, 10_000_000))
        .unwrap()
        .payload(&payload)
        .unwrap()
        .polymorphic(rng.bool() as u64)
        .unwrap();
    for t in &targets {
        b = b.targets(t).unwrap();
    }
    b.finish().unwrap().to_vec()
}

fn gen_operation(rng: &mut Rng, crew_count: usize, stolen_bytes: usize) -> Vec<u8> {
    let crew: Vec<Vec<u8>> = (0..crew_count).map(|_| gen_hacker(rng)).collect();
    let worm = gen_worm(rng);
    let stolen = if stolen_bytes > 0 {
        rng.bytes(stolen_bytes)
    } else {
        vec![]
    };
    let codename = *rng.pick(OP_NAMES);
    let corp = *rng.pick(CORPS);
    let ts = 810_000_000u64 + rng.range(0, 10_000_000);
    let n_logs = rng.range(2, 6) as usize;
    let logs: Vec<&str> = (0..n_logs).map(|_| *rng.pick(LOG_LINES)).collect();

    // Operation buffer: crew + worm + stolen_data + small overhead
    let crew_size: usize = crew.iter().map(|h| h.len() + 5).sum();
    let buf_size = crew_size + worm.len() + stolen_bytes + 1024;
    let mut buf = vec![0u8; buf_size];

    let mut b = OperationBuilder::builder(&mut buf)
        .codename(codename)
        .unwrap()
        .target_corp(corp)
        .unwrap()
        .timestamp(ts)
        .unwrap()
        .successful(rng.bool() as u64)
        .unwrap();
    if stolen_bytes > 0 {
        b = b.stolen_data(&stolen).unwrap();
    }
    for h in &crew {
        b = b.crew(h).unwrap();
    }
    b = b.worm(&worm).unwrap();
    for l in &logs {
        b = b.log_entries(l).unwrap();
    }
    b.severity(rng.range(1, 10) as i32)
        .unwrap()
        .finish()
        .unwrap()
        .to_vec()
}

fn gen_campaign(
    rng: &mut Rng,
    op_count: usize,
    crew_per_op: usize,
    stolen_bytes: usize,
) -> Vec<u8> {
    let ops: Vec<Vec<u8>> = (0..op_count)
        .map(|i| {
            if i > 0 && i % 100 == 0 {
                eprintln!(
                    "  {i}/{op_count} operations ({:.1} MB in ops so far)…",
                    i * (stolen_bytes + 10_000) / 1_000_000
                );
            }
            gen_operation(rng, crew_per_op, stolen_bytes)
        })
        .collect();

    // Pre-compute campaign buffer size: for each op the wire encoding is
    // tag(1B) + varint_length(1-5B) + op_bytes
    let ops_wire_size: usize = ops
        .iter()
        .map(|o| 1 + varint_len(o.len() as u64) + o.len())
        .sum();
    let mut buf = vec![0u8; ops_wire_size + 64];

    let mut b = CampaignBuilder::builder(&mut buf)
        .name("HACK THE PLANET CAMPAIGN")
        .unwrap();
    for op in &ops {
        b = b.operations(op).unwrap();
    }
    b.total_bytes_stolen((stolen_bytes * op_count) as u64)
        .unwrap()
        .finish()
        .unwrap()
        .to_vec()
}

/// Number of bytes needed to encode `v` as a varint.
fn varint_len(mut v: u64) -> usize {
    let mut n = 1usize;
    while v >= 128 {
        v >>= 7;
        n += 1;
    }
    n
}

// =============================================================================
// Preset table
// =============================================================================

struct Preset {
    ops: usize,
    crew: usize,
    stolen_kb: usize,
}

fn resolve(name: &str) -> Option<Preset> {
    match name {
        //                              ops     crew  stolen_kb   approx size
        "tiny" => Some(Preset {
            ops: 1,
            crew: 1,
            stolen_kb: 0,
        }), // ~400 B
        "small" => Some(Preset {
            ops: 20,
            crew: 3,
            stolen_kb: 0,
        }), // ~25 KB
        "medium" => Some(Preset {
            ops: 2_000,
            crew: 3,
            stolen_kb: 0,
        }), // ~2 MB
        "large" => Some(Preset {
            ops: 200,
            crew: 3,
            stolen_kb: 500,
        }), // ~100 MB
        "huge" => Some(Preset {
            ops: 1_000,
            crew: 3,
            stolen_kb: 500,
        }), // ~500 MB
        _ => None,
    }
}

// =============================================================================
// main
// =============================================================================

fn main() {
    let mut args = Args::parse();

    if let Some(ref name) = args.preset.clone() {
        match resolve(name) {
            Some(p) => {
                args.ops = p.ops;
                args.crew = p.crew;
                args.stolen_kb = p.stolen_kb;
            }
            None => {
                eprintln!("Unknown preset '{name}'. Valid: tiny, small, medium, large, huge");
                std::process::exit(1);
            }
        }
    }

    let stolen_bytes = args.stolen_kb * 1024;
    let approx_mb = args.ops * (700 + args.crew * 300 + stolen_bytes) / 1_000_000;
    eprintln!(
        "Generating: {} ops × {} crew, {} KB stolen_data each → ~{} MB",
        args.ops, args.crew, args.stolen_kb, approx_mb
    );

    let mut rng = Rng::new(args.seed);
    let data = gen_campaign(&mut rng, args.ops, args.crew, stolen_bytes);

    eprintln!(
        "Generated {} bytes ({:.2} MB)",
        data.len(),
        data.len() as f64 / 1_000_000.0
    );

    // Default output: data/bench/<preset>.pb, or stdout if no output and no preset
    let out_path = args.output.clone().or_else(|| {
        args.preset
            .as_ref()
            .map(|p| format!("benches/data/bench/{}.pb", p))
    });

    match out_path {
        Some(ref path) => {
            if let Some(parent) = Path::new(path).parent() {
                if !parent.as_os_str().is_empty() {
                    std::fs::create_dir_all(parent).unwrap_or_else(|e| eprintln!("Warning: {e}"));
                }
            }
            std::fs::write(path, &data).unwrap_or_else(|e| {
                eprintln!("Error writing {path}: {e}");
                std::process::exit(1);
            });
            eprintln!("Saved to {path}");
        }
        None => {
            io::stdout().write_all(&data).unwrap_or_else(|e| {
                eprintln!("Error: {e}");
                std::process::exit(1);
            });
        }
    }
}
