use crate::{
    fuzzer::context::{Context, MutationMetadata},
    Node, Visitor,
};
use libafl::{
    corpus::Corpus,
    events::EventFirer,
    executors::Executor,
    stages::{Restartable, Stage},
    state::{HasCorpus, HasCurrentTestcase},
    Evaluator, HasMetadata,
};
use serde::Serialize;
use std::{
    cell::RefCell,
    collections::{BTreeMap, HashSet},
    marker::PhantomData,
    path::PathBuf,
    rc::Rc,
    time::{Duration, Instant},
};

#[derive(Debug)]
pub struct StatsStage<I> {
    last_run: Instant,
    out_dir: PathBuf,
    phantom: PhantomData<I>,
}

impl<I> StatsStage<I> {
    pub fn new(out_dir: PathBuf) -> Self {
        Self {
            last_run: Instant::now(),
            out_dir,
            phantom: PhantomData,
        }
    }
}

impl<E, EM, Z, S, I> Stage<E, EM, S, Z> for StatsStage<I>
where
    I: Node + Serialize,
    S: HasCurrentTestcase<I> + HasCorpus<I> + HasMetadata,
    E: Executor<EM, I, S, Z>,
    EM: EventFirer<I, S>,
    Z: Evaluator<E, EM, I, S>,
{
    fn perform(
        &mut self,
        fuzzer: &mut Z,
        executor: &mut E,
        state: &mut S,
        manager: &mut EM,
    ) -> Result<(), libafl::Error> {
        if Instant::now() - self.last_run > Duration::from_secs(5) {
            let mut metadata = state.metadata_mut::<AutarkieStats>()?;
            std::fs::write(
                self.out_dir.join("stats.json"),
                serde_json::to_string_pretty(&metadata).expect("____YR5BenN6"),
            )
            .expect("____weNooV3S");
            self.last_run = Instant::now();
        }
        Ok(())
    }
}

impl<I, S> Restartable<S> for StatsStage<I> {
    fn should_restart(&mut self, state: &mut S) -> Result<bool, libafl::Error> {
        Ok(true)
    }

    fn clear_progress(&mut self, state: &mut S) -> Result<(), libafl::Error> {
        Ok(())
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, libafl_bolts::SerdeAny, Default)]
pub struct AutarkieStats {
    mutations: BTreeMap<MutationMetadata, usize>,
}

impl AutarkieStats {
    pub fn add_new_input_mutations(&mut self, mutations: HashSet<MutationMetadata>) {
        for m in mutations {
            self.mutations
                .entry(m)
                .and_modify(|v| {
                    *v += 1;
                })
                .or_insert(1);
        }
    }
    pub fn add_new_input_mutation(&mut self, m: MutationMetadata) {
        self.mutations
            .entry(m)
            .and_modify(|v| {
                *v += 1;
            })
            .or_insert(1);
    }
}
