use libafl::prelude::UsesInput;
use libafl_qemu::CPUArchState;
use libafl_qemu::Emulator;
use libafl_qemu::FastSnapshot;
use libafl_qemu::QemuExecutor;
use libafl_qemu::QemuHelper;
use libafl_qemu::QemuHelperTuple;
use libafl::{executors::ExitKind, inputs::Input, observers::ObserversTuple, state::HasMetadata};
use libafl_qemu::QemuHooks;

use libafl_qemu::{
    emu,
};
// TODO be thread-safe maybe with https://amanieu.github.io/thread_local-rs/thread_local/index.html
#[derive(Debug)]
pub struct QemuStateRestoreHelper {
    has_snapshot: bool,
    use_snapshot: bool,
    saved_cpu_states: Vec<CPUArchState>,
    fastsnap: Option<FastSnapshot>
}

impl QemuStateRestoreHelper {
    #[must_use]
    pub fn new() -> Self {
        Self {
            has_snapshot: false,
            use_snapshot: true,
            saved_cpu_states: vec![],
            fastsnap: None
        }
    }
}

impl Default for QemuStateRestoreHelper {
    fn default() -> Self {
        Self::new()
    }
}

impl<S> QemuHelper<S> for QemuStateRestoreHelper
where
    S: UsesInput,
{
    const HOOKS_DO_SIDE_EFFECTS: bool = true;

    fn init_hooks<QT>(&self, _hooks: &QemuHooks<'_, QT, S>)
    where
        QT: QemuHelperTuple<S>,
    {
    }

    fn first_exec<QT>(&self, _hooks: &QemuHooks<'_, QT, S>)
    where
        QT: QemuHelperTuple<S>,
    {
    }

    fn post_exec<OT>(&mut self, emulator: &Emulator, _input: &S::Input, _observers: &mut OT, _exit_kind: &mut ExitKind) {
        // unsafe { println!("snapshot post {}",emu::icount_get_raw()) };
    }

    fn pre_exec(&mut self, emulator: &Emulator, _input: &S::Input) {
        // only restore in pre-exec, to preserve the post-execution state for inspection
        #[cfg(feature = "snapshot_restore")]
        {
            #[cfg(feature = "snapshot_fast")]
            match self.fastsnap {
                Some(s) => emulator.restore_fast_snapshot(s),
                None => {self.fastsnap = Some(emulator.create_fast_snapshot(true));},
            }
            #[cfg(not(feature = "snapshot_fast"))]
            if !self.has_snapshot {
                emulator.save_snapshot("Start", true);
                self.has_snapshot = true;
            }
            else
            {
                emulator.load_snapshot("Start", true);
            }
        }
        #[cfg(not(feature = "snapshot_restore"))]
        if !self.has_snapshot {
            self.saved_cpu_states = (0..emulator.num_cpus())
                .map(|i| emulator.cpu_from_index(i).save_state())
                .collect();
            self.has_snapshot = true;
        } else {
            for (i, s) in self.saved_cpu_states.iter().enumerate() {
                emulator.cpu_from_index(i).restore_state(s);
            }
        }

        // unsafe { println!("snapshot pre {}",emu::icount_get_raw()) };
    }
}