#![feature(int_log)]

mod apu;
mod audio;
mod cartridge;
mod cpu;
mod fuzzing_state;
mod input;
mod ppu;
mod screen;
mod state;

use apu::Apu;
use ppu::Ppu;

use crate::cpu::Cpu;
use crate::fuzzing_state::{FuzzingInputState, FuzzingState};
use crate::input::ConsoleAction::SoftReset;
use crate::input::{ConsoleAction, FuzzingInput};
use crate::screen::Screen;
use minifb::{ScaleMode, Window, WindowOptions};
use priority_queue::double_priority_queue::DoublePriorityQueue;
use priority_queue::priority_queue::PriorityQueue;
use std::collections::HashSet;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::{mpsc, Arc, Mutex};
use std::{fs, thread};

// The number of cpu instances to spawn..
const NUM_THREADS: usize = 28;

// The number of frames to fuzz and process
// A small number exploits the current point more at the expense of
// large exploration - and vice versa.
const FRAMES_TO_CONSIDER: usize = 500;

// Same input should generate the same output...
// (I make no guarantee of that at the moment)
const RNG_SEED: u32 = 0x1334357;

// If set to a low number, this disables start presses after the given frame
// Useful for some games where pausing does nothing to advance the game...
const DISABLE_START_PRESSES_AFTER: usize = 100;

// The rate at which seed inputs become corrupted..
const MUTATION_RATE: f64 = 0.10;

// The rate at which seed inputs may become soft resets..
const MUTATION_RATE_SOFT_RESET: f64 = 0.00;

// The number of cases to fuzz from a given seed input at each consideration point..
const SEED_CASES: usize = 10;

// Only add the original seed case top the queue (hammer a single state)
const SINGLE_SHOT: bool = true;

// Fast forward the seed state to a given frame
const PLAY_FROM: usize = 1900;

fn main() -> Result<(), String> {
    let argv = std::env::args().collect::<Vec<String>>();

    if argv.len() < 3 {
        println!("usage: fuzz <rom> <fm2 seed input>")
    }

    // The rom file...mostly tested with Super Mario Bros
    let rom_filename = argv[1].to_string();
    // The tas file to use as seed input e.g. happylee-supermariobros,warped.fm2
    let fm2_file = argv[2].to_string();

    // Seed for our Random Number Generator...
    let mut rng = Rng::init(RNG_SEED);

    // Mutex for launching windows others we get double frees in the
    // windowing library...(yup)
    let mutex = Arc::new(Mutex::new(0));

    let mut seed_input = FuzzingInput::load(fm2_file.as_str());
    seed_input.disable_start_after(DISABLE_START_PRESSES_AFTER);

    let mut starting_state = FuzzingState::default(rom_filename.clone());

    let mut previous_states: HashSet<usize> = HashSet::new();

    if PLAY_FROM > 0 {
        println!("Playing {} frames prelude...", PLAY_FROM);
        let (mut cpu, _) = starting_state.load_state(rom_filename.clone());
        play_frames(&mut cpu, PLAY_FROM, &seed_input);
        starting_state = FuzzingState::save_state(cpu, PLAY_FROM);
        for addr in starting_state.cpu_state.addresses_fetched.iter() {
            if previous_states.insert(*addr) {}
        }
    }

    let mut fuzzing_queue: PriorityQueue<FuzzingInputState, u64> = PriorityQueue::new();
    fuzzing_queue.push(
        FuzzingInputState(seed_input.clone(), starting_state.clone()),
        10000,
    );
    for _i in 1..NUM_THREADS {
        let mut mutated_input = seed_input.clone();
        mutated_input.mutate_from(&mut rng, PLAY_FROM, FRAMES_TO_CONSIDER);
        fuzzing_queue.push(
            FuzzingInputState(mutated_input.clone(), starting_state.clone()),
            rng.next() as u64,
        );
    }

    let mut result_channels = vec![];
    for i in 0..NUM_THREADS {
        let i = i.clone();
        let mutex = mutex.clone();
        let rom_filename = rom_filename.clone();
        let (tx_inputs, rx_input): (
            Sender<(FuzzingInput, FuzzingState)>,
            Receiver<(FuzzingInput, FuzzingState)>,
        ) = mpsc::channel();
        let (tx_results, rx_results): (
            Sender<(FuzzingInput, FuzzingState)>,
            Receiver<(FuzzingInput, FuzzingState)>,
        ) = mpsc::channel();
        result_channels.push((tx_inputs, rx_results));
        thread::spawn(move || {
            run_game(i as isize, mutex, rom_filename, rx_input, tx_results);
        });
    }

    let mut loop_count = 0;

    loop {
        println!("Prospective Cases: {}", fuzzing_queue.len());

        for i in 0..NUM_THREADS {
            if fuzzing_queue.is_empty() == false {
                let (state, score) = fuzzing_queue.pop().unwrap();
                // Send to thread for process...
                result_channels[i]
                    .0
                    .send((state.0.clone(), state.1.clone()))
                    .expect("error sending new fuzzing input to thread...");
            } else {
                result_channels[i]
                    .0
                    .send((seed_input.clone(), starting_state.clone()))
                    .expect("error sending new fuzzing input to thread");
            }
        }

        // Gather results from each thread
        for i in 0..NUM_THREADS {
            match result_channels[i].1.recv() {
                Ok((fuzzing_input, fuzzing_state)) => {
                    /*                    let mut lowest_similarity = u64::MAX;
                    for (_j, state) in previous_states.iter().enumerate() {
                        let similiarty =
                            hamming::distance(&state, &fuzzing_state.cpu_state.mem);

                        if similiarty < lowest_similarity {
                            lowest_similarity = similiarty
                        }
                    }

                    previous_states.insert(fuzzing_state.cpu_state.mem.clone());*/
                    let mut lowest_similarity = 0u64;
                    for addr in fuzzing_state.cpu_state.addresses_fetched.iter() {
                        if previous_states.insert(*addr) {
                            lowest_similarity += 1;
                        }
                    }

                    if i == 3 && loop_count == 3 {
                        fuzzing_input.export(
                            "megaman_interestng.fm2".parse().unwrap(),
                            fuzzing_state.frames,
                        );
                    }

                    if SINGLE_SHOT {
                        if lowest_similarity > 0 {
                            println!("{}  ", lowest_similarity);
                            for i in 0..lowest_similarity {
                                let mut mutated_input: FuzzingInput = fuzzing_input.clone();
                                mutated_input.mutate_from(
                                    &mut rng,
                                    starting_state.frames,
                                    FRAMES_TO_CONSIDER,
                                );
                                // Add the mutated input and the regular input to the queue...
                                fuzzing_queue.push(
                                    FuzzingInputState(
                                        mutated_input.clone(),
                                        starting_state.clone(),
                                    ),
                                    lowest_similarity,
                                );
                            }
                        }
                        continue;
                    } else {
                        // Seed Thread
                        if fuzzing_input.mutated == false {
                            // Seed input handling...
                            fuzzing_queue.push(
                                FuzzingInputState(fuzzing_input.clone(), fuzzing_state.clone()),
                                10000,
                            );

                            // Add seed cases..
                            for _i in 0..SEED_CASES {
                                let mut mutated_input: FuzzingInput = fuzzing_input.clone();
                                mutated_input.mutate_from(
                                    &mut rng,
                                    fuzzing_state.frames,
                                    FRAMES_TO_CONSIDER,
                                );
                                // Add the mutated input and the regular input to the queue...
                                fuzzing_queue.push(
                                    FuzzingInputState(mutated_input.clone(), fuzzing_state.clone()),
                                    lowest_similarity,
                                );
                            }
                        } else if lowest_similarity > 0 {
                            // Only add to the queue if strictly better...
                            fuzzing_queue.push(
                                FuzzingInputState(fuzzing_input.clone(), fuzzing_state.clone()),
                                lowest_similarity,
                            );

                            let mut mutated_input: FuzzingInput = fuzzing_input.clone();
                            mutated_input.mutate_from(
                                &mut rng,
                                fuzzing_state.frames,
                                FRAMES_TO_CONSIDER,
                            );
                            // Add the mutated input and the regular input to the queue...
                            fuzzing_queue.push(
                                FuzzingInputState(mutated_input.clone(), fuzzing_state.clone()),
                                lowest_similarity,
                            );
                        }
                    }
                }
                _ => {}
            }
        }

        if SINGLE_SHOT {
            // Add seed cases..
            for _i in 0..SEED_CASES {
                let mut mutated_input: FuzzingInput = seed_input.clone();
                mutated_input.mutate_from(&mut rng, starting_state.frames, FRAMES_TO_CONSIDER);
                // Add the mutated input and the regular input to the queue...
                fuzzing_queue.push(
                    FuzzingInputState(mutated_input.clone(), starting_state.clone()),
                    rng.next() as u64,
                );
            }
        }
        loop_count += 1;
    }
}

pub struct Rng {
    state: u32,
}

impl Rng {
    pub fn init(state: u32) -> Rng {
        Rng { state }
    }

    pub fn next(&mut self) -> u32 {
        self.state ^= self.state << 13;
        self.state ^= self.state >> 17;
        self.state ^= self.state << 5;
        return self.state;
    }
}

fn run_game(
    i: isize,
    mutex: Arc<Mutex<u8>>,
    filename: String,
    refresh_inputs: Receiver<(FuzzingInput, FuzzingState)>,
    send_results: Sender<(FuzzingInput, FuzzingState)>,
) -> FuzzingState {
    let mut window = match mutex.lock() {
        Ok(_x) => {
            let mut window = Window::new(
                "NesFuzz",
                256,
                240,
                WindowOptions {
                    resize: true,
                    borderless: true,
                    scale_mode: ScaleMode::UpperLeft,
                    ..WindowOptions::default()
                },
            )
            .unwrap_or_else(|x| panic!("Unable to create window {}", x));

            window.set_position(10 + ((i % 7) * 260), (240 * (i / 7)) + 100);
            window
        }
        _ => {
            panic!("unable to acquire mutex")
        }
    };

    let mut screen = Screen::new();

    loop {
        match refresh_inputs.recv() {
            Err(_x) => {}
            Ok((fuzzing_input, fuzzing_state)) => {
                let mut new_frames = 0;

                // Initialize hardware components
                let (mut cpu, mut frames) = fuzzing_state.load_state(filename.clone());
                let mut prev_frame = frames;
                while new_frames < FRAMES_TO_CONSIDER {
                    // step CPU: perform 1 cpu instruction, getting back number of clock cycles it took

                    let cpu_cycles = cpu.step();
                    // clock PPU three times for every CPU cycle
                    for _ in 0..cpu_cycles * 3 {
                        let (pixel, end_of_frame) = cpu.ppu.clock();
                        match pixel {
                            Some((x, y, color)) => screen.plot_pixel(
                                x,
                                y,
                                color.0 as u32,
                                color.1 as u32,
                                color.2 as u32,
                            ),
                            None => (),
                        };
                        if end_of_frame {
                            screen.draw_string(10, 210, &*format!("Frames: {}", new_frames));
                            window
                                .update_with_buffer(&screen.display, 256, 240)
                                .unwrap();
                            frames += 1;
                            new_frames += 1;
                        }
                    }
                    // Checking for Inputs...
                    match fuzzing_input.get_frame_input(frames) {
                        Some(input) => {
                            match input.console_action {
                                ConsoleAction::None => {}
                                ConsoleAction::SoftReset => {
                                    println!("soft reset");
                                    cpu.soft_reset();
                                    frames += 1;
                                }
                            }
                            //if cpu.strobe {
                            cpu.button_states = input.player_1_input;
                            // FIXME PLayer 2 doesn't play nicely with some games (e.g. mario)
                            // So to enable player 2 controls you also have to uncomment the
                            // bus in cpu/mod.rs
                            cpu.button_states2 = input.player_2_input;
                            //}
                        }
                        _ => {}
                    };
                }

                send_results
                    .send((fuzzing_input.clone(), FuzzingState::save_state(cpu, frames)))
                    .expect("error sending result to thread");
            }
        }
    }
}

fn play_frames(cpu: &mut Cpu, num_frames: usize, fuzzing_input: &FuzzingInput) {
    let mut played_frames = 0;
    let mut prev_frame = 0;
    while played_frames < num_frames {
        // step CPU: perform 1 cpu instruction, getting back number of clock cycles it took

        // A the start of each new frame set our input for that frame...
        if prev_frame == 0 || prev_frame != played_frames {
            match fuzzing_input.get_frame_input(played_frames) {
                Some(input) => {
                    match input.console_action {
                        ConsoleAction::None => {}
                        ConsoleAction::SoftReset => {
                            cpu.soft_reset();
                            played_frames += 1;
                        }
                    }

                    cpu.button_states = input.player_1_input;
                    cpu.button_states2 = input.player_2_input;
                }
                _ => {
                    cpu.button_states = 0;
                    cpu.button_states2 = 0;
                }
            };
            prev_frame = played_frames;
        }

        let cpu_cycles = cpu.step();

        // clock PPU three times for every CPU cycle
        for _ in 0..cpu_cycles * 3 {
            let (_, end_of_frame) = cpu.ppu.clock();
            if end_of_frame {
                prev_frame = played_frames;
                played_frames += 1;
            }
        }
    }
}