use clap::{Parser, Args};
use rand::{Rng, distributions::{Distribution, Uniform}};

use std::io::{self, Read, Write};

const BUF_SIZE: usize = 8192;

#[derive(Parser, Debug)]
#[clap(author, version, about)]
struct Cli {
    #[clap(short, long)]
    verbose: bool,

    /// Introduce an error every N bytes on average
    #[clap(short, long, value_name = "N", default_value_t = 1000)]
    rarity: u32,

    #[clap(flatten, help_heading = "Types of errors to cause. If none are specified individually, all will be allowed")]
    error_spec: ErrorSpec,
}

#[derive(Args, Debug, Clone, Copy)]
struct ErrorSpec {
    /// Insert extra bytes into the stream
    #[clap(short, long)]
    insert: bool,

    /// Delete bytes from the stream
    #[clap(short, long)]
    delete: bool,

    /// Alter bytes in the stream
    #[clap(short, long)]
    alter: bool,
}

enum Error {
    Insert(u8),
    Delete,
    Alter(u8),
}

impl ErrorSpec {
    fn gen(self) -> Error {
        let spec = if let ErrorSpec { insert: false, delete: false, alter: false } = self {
            ErrorSpec { insert: true, delete: true, alter: true }
        } else {
            self
        };

        let mut options_remaining =
            if spec.insert { 1 } else { 0 } +
            if spec.delete { 1 } else { 0 } +
            if spec.alter { 1 } else { 0 };

        let mut rng = rand::thread_rng();

        if spec.insert {
            if rng.gen_range(0..options_remaining) == 0 {
                return Error::Insert(rng.gen())
            }
            options_remaining -= 1;
        }
        if spec.delete {
            if rng.gen_range(0..options_remaining) == 0 {
                return Error::Delete
            }
            options_remaining -= 1;
        }
        if spec.alter {
            if rng.gen_range(0..options_remaining) == 0 {
                return Error::Alter(rng.gen())
            }
            options_remaining -= 1;
        }
        panic!("Expected 0 options remaining, got {}, this shouldn't happen", options_remaining)
    }
}

fn read_handling_interruptions<T: Read>(source: &mut T, dest: &mut [u8]) -> io::Result<usize> {
    loop {
        let result = source.read(dest);
        if let Err(e) = &result {
            if let io::ErrorKind::Interrupted = e.kind() {
                continue
            }
        }
        return result
    }
}

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

    let mut stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut buffer = [0u8; BUF_SIZE];
    let mut total = 0;
    let rarity = Uniform::from(0..cli.rarity);
    let mut rng = rand::thread_rng();

    loop {
        let fill = read_handling_interruptions(&mut stdin, &mut buffer).unwrap();
        if fill == 0 {
            break
        }
        let mut pos = 0;
        let error_indices = (0..fill).filter(|_| rarity.sample(&mut rng) == 0);
        for idx in error_indices {
            stdout.write_all(&buffer[pos..idx]).unwrap();
            if cli.verbose { eprint!("@ 0x{:08x}: ", total + idx); }
            match cli.error_spec.gen() {
                Error::Insert(b) => {
                    stdout.write_all(&[buffer[idx], b]).unwrap();
                    if cli.verbose { eprintln!("Inserted byte {:02x}", b); }
                }
                Error::Delete => {
                    if cli.verbose { eprintln!("Deleted byte {:02x}", buffer[idx]); }
                }
                Error::Alter(b) => {
                    stdout.write_all(&[b]).unwrap();
                    if cli.verbose { eprintln!("Altered byte {:02x} to {:02x}", buffer[idx], b); }
                }
            }
            pos = idx + 1;
        }
        stdout.write_all(&buffer[pos..fill]).unwrap();
        total += fill;
    }

    if cli.verbose { eprintln!("Reached EOF after {} bytes", total); }
}