use crate::{PlatformDispatcher, TaskLabel};
use async_task::Runnable;
use backtrace::Backtrace;
use collections::{HashMap, HashSet, VecDeque};
use parking::{Parker, Unparker};
use parking_lot::Mutex;
use rand::prelude::*;
use std::{
    future::Future,
    ops::RangeInclusive,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::{Duration, Instant},
};
use util::post_inc;

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
struct TestDispatcherId(usize);

#[doc(hidden)]
pub struct TestDispatcher {
    id: TestDispatcherId,
    state: Arc<Mutex<TestDispatcherState>>,
    parker: Arc<Mutex<Parker>>,
    unparker: Unparker,
}

struct TestDispatcherState {
    random: StdRng,
    foreground: HashMap<TestDispatcherId, VecDeque<Runnable>>,
    background: Vec<Runnable>,
    deprioritized_background: Vec<Runnable>,
    delayed: Vec<(Duration, Runnable)>,
    start_time: Instant,
    time: Duration,
    is_main_thread: bool,
    next_id: TestDispatcherId,
    allow_parking: bool,
    waiting_hint: Option<String>,
    waiting_backtrace: Option<Backtrace>,
    deprioritized_task_labels: HashSet<TaskLabel>,
    block_on_ticks: RangeInclusive<usize>,
}

impl TestDispatcher {
    pub fn new(random: StdRng) -> Self {
        let (parker, unparker) = parking::pair();
        let state = TestDispatcherState {
            random,
            foreground: HashMap::default(),
            background: Vec::new(),
            deprioritized_background: Vec::new(),
            delayed: Vec::new(),
            time: Duration::ZERO,
            start_time: Instant::now(),
            is_main_thread: true,
            next_id: TestDispatcherId(1),
            allow_parking: false,
            waiting_hint: None,
            waiting_backtrace: None,
            deprioritized_task_labels: Default::default(),
            block_on_ticks: 0..=1000,
        };

        TestDispatcher {
            id: TestDispatcherId(0),
            state: Arc::new(Mutex::new(state)),
            parker: Arc::new(Mutex::new(parker)),
            unparker,
        }
    }

    pub fn advance_clock(&self, by: Duration) {
        let new_now = self.state.lock().time + by;
        loop {
            self.run_until_parked();
            let state = self.state.lock();
            let next_due_time = state.delayed.first().map(|(time, _)| *time);
            drop(state);
            if let Some(due_time) = next_due_time {
                if due_time <= new_now {
                    self.state.lock().time = due_time;
                    continue;
                }
            }
            break;
        }
        self.state.lock().time = new_now;
    }

    pub fn simulate_random_delay(&self) -> impl 'static + Send + Future<Output = ()> {
        struct YieldNow {
            pub(crate) count: usize,
        }

        impl Future for YieldNow {
            type Output = ();

            fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
                if self.count > 0 {
                    self.count -= 1;
                    cx.waker().wake_by_ref();
                    Poll::Pending
                } else {
                    Poll::Ready(())
                }
            }
        }

        YieldNow {
            count: self.state.lock().random.gen_range(0..10),
        }
    }

    pub fn tick(&self, background_only: bool) -> bool {
        let mut state = self.state.lock();

        while let Some((deadline, _)) = state.delayed.first() {
            if *deadline > state.time {
                break;
            }
            let (_, runnable) = state.delayed.remove(0);
            state.background.push(runnable);
        }

        let foreground_len: usize = if background_only {
            0
        } else {
            state
                .foreground
                .values()
                .map(|runnables| runnables.len())
                .sum()
        };
        let background_len = state.background.len();

        let runnable;
        let main_thread;
        if foreground_len == 0 && background_len == 0 {
            let deprioritized_background_len = state.deprioritized_background.len();
            if deprioritized_background_len == 0 {
                return false;
            }
            let ix = state.random.gen_range(0..deprioritized_background_len);
            main_thread = false;
            runnable = state.deprioritized_background.swap_remove(ix);
        } else {
            main_thread = state.random.gen_ratio(
                foreground_len as u32,
                (foreground_len + background_len) as u32,
            );
            if main_thread {
                let state = &mut *state;
                runnable = state
                    .foreground
                    .values_mut()
                    .filter(|runnables| !runnables.is_empty())
                    .choose(&mut state.random)
                    .unwrap()
                    .pop_front()
                    .unwrap();
            } else {
                let ix = state.random.gen_range(0..background_len);
                runnable = state.background.swap_remove(ix);
            };
        };

        let was_main_thread = state.is_main_thread;
        state.is_main_thread = main_thread;
        drop(state);
        runnable.run();
        self.state.lock().is_main_thread = was_main_thread;

        true
    }

    pub fn deprioritize(&self, task_label: TaskLabel) {
        self.state
            .lock()
            .deprioritized_task_labels
            .insert(task_label);
    }

    pub fn run_until_parked(&self) {
        while self.tick(false) {}
    }

    pub fn parking_allowed(&self) -> bool {
        self.state.lock().allow_parking
    }

    pub fn allow_parking(&self) {
        self.state.lock().allow_parking = true
    }

    pub fn forbid_parking(&self) {
        self.state.lock().allow_parking = false
    }

    pub fn set_waiting_hint(&self, msg: Option<String>) {
        self.state.lock().waiting_hint = msg
    }

    pub fn waiting_hint(&self) -> Option<String> {
        self.state.lock().waiting_hint.clone()
    }

    pub fn start_waiting(&self) {
        self.state.lock().waiting_backtrace = Some(Backtrace::new_unresolved());
    }

    pub fn finish_waiting(&self) {
        self.state.lock().waiting_backtrace.take();
    }

    pub fn waiting_backtrace(&self) -> Option<Backtrace> {
        self.state.lock().waiting_backtrace.take().map(|mut b| {
            b.resolve();
            b
        })
    }

    pub fn rng(&self) -> StdRng {
        self.state.lock().random.clone()
    }

    pub fn set_block_on_ticks(&self, range: std::ops::RangeInclusive<usize>) {
        self.state.lock().block_on_ticks = range;
    }

    pub fn gen_block_on_ticks(&self) -> usize {
        let mut lock = self.state.lock();
        let block_on_ticks = lock.block_on_ticks.clone();
        lock.random.gen_range(block_on_ticks)
    }
}

impl Clone for TestDispatcher {
    fn clone(&self) -> Self {
        let id = post_inc(&mut self.state.lock().next_id.0);
        Self {
            id: TestDispatcherId(id),
            state: self.state.clone(),
            parker: self.parker.clone(),
            unparker: self.unparker.clone(),
        }
    }
}

impl PlatformDispatcher for TestDispatcher {
    fn is_main_thread(&self) -> bool {
        self.state.lock().is_main_thread
    }

    fn now(&self) -> Instant {
        let state = self.state.lock();
        state.start_time + state.time
    }

    fn dispatch(&self, runnable: Runnable, label: Option<TaskLabel>) {
        {
            let mut state = self.state.lock();
            if label.map_or(false, |label| {
                state.deprioritized_task_labels.contains(&label)
            }) {
                state.deprioritized_background.push(runnable);
            } else {
                state.background.push(runnable);
            }
        }
        self.unparker.unpark();
    }

    fn dispatch_on_main_thread(&self, runnable: Runnable) {
        self.state
            .lock()
            .foreground
            .entry(self.id)
            .or_default()
            .push_back(runnable);
        self.unparker.unpark();
    }

    fn dispatch_after(&self, duration: std::time::Duration, runnable: Runnable) {
        let mut state = self.state.lock();
        let next_time = state.time + duration;
        let ix = match state.delayed.binary_search_by_key(&next_time, |e| e.0) {
            Ok(ix) | Err(ix) => ix,
        };
        state.delayed.insert(ix, (next_time, runnable));
    }
    fn park(&self, _: Option<std::time::Duration>) -> bool {
        self.parker.lock().park();
        true
    }

    fn unparker(&self) -> Unparker {
        self.unparker.clone()
    }

    fn as_test(&self) -> Option<&TestDispatcher> {
        Some(self)
    }
}
