use pyo3::{
    exceptions::{PyIndexError, PyKeyError, PyStopIteration, PyTypeError},
    prelude::*,
    types::{PyDict, PyType},
};

use crate::backing::{
    containers::PyItem,
    indexed::{IndexedBacking, IndexedBinaryHeap},
};

#[pyclass]
pub struct IndexedQueue {
    backing: Box<dyn IndexedBacking<PyItem, f64>>,
}

// TODO: Too much of this is just too similar to the other queues
// Normally I'd solve this with inheritance, but I'm not so familiar with rust trait default implementations yet
// In my limited understanding they don't seem to play nice with the Pyo3 bindings
#[pymethods]
impl IndexedQueue {
    #[new]
    #[pyo3(signature = (items=None))]
    fn new(items: Option<Py<PyAny>>) -> PyResult<Self> {
        if let Some(py_object) = items {
            Python::with_gil(|py| Self::from_any(py_object.bind(py))).map(|vec| Self {
                backing: Box::new(IndexedBinaryHeap::from_iter(vec)),
            })
        } else {
            Ok(Self {
                backing: Box::new(IndexedBinaryHeap::new()),
            })
        }
    }

    fn __setitem__(mut self_: PyRefMut<'_, Self>, key: Py<PyAny>, value: f64) {
        self_.backing.set(PyItem::new(key), value)
    }

    fn __delitem__(mut self_: PyRefMut<'_, Self>, key: Py<PyAny>) -> PyResult<()> {
        match self_.backing.remove(PyItem::new(key)) {
            Some(_) => Ok(()),
            None => Err(PyErr::new::<PyKeyError, _>(
                "Key was not contained in queue",
            )),
        }
    }

    fn __contains__(self_: PyRef<'_, Self>, key: Py<PyAny>) -> bool {
        self_.backing.contains(&PyItem::new(key))
    }

    /// Enables generic typing
    #[classmethod]
    fn __class_getitem__(cls_: Bound<'_, PyType>, _key: Py<PyAny>) -> Bound<'_, PyType> {
        cls_
    }

    fn __len__(self_: PyRef<'_, Self>) -> usize {
        self_.backing.len()
    }

    fn __iter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
        self_
    }

    fn __next__(mut self_: PyRefMut<'_, Self>) -> PyResult<Py<PyAny>> {
        if let Some(item) = self_.backing.pop() {
            Ok(item.data())
        } else {
            Err(PyErr::new::<PyStopIteration, _>(()))
        }
    }

    fn pop(mut self_: PyRefMut<'_, Self>) -> PyResult<Py<PyAny>> {
        if let Some(item) = self_.backing.pop() {
            Ok(item.data())
        } else {
            Err(PyErr::new::<PyIndexError, _>(()))
        }
    }
}

impl<'py> IndexedQueue {
    /// Tries to a Python object into a vector suitable for ingestion into the backing
    fn from_any(object: &Bound<'py, PyAny>) -> PyResult<Vec<(PyItem, f64)>> {
        if let Ok(vec) = object.extract::<Vec<(Py<PyAny>, f64)>>() {
            Ok(Self::from_vec(vec))
        } else if object.is_instance_of::<PyDict>() {
            if let Ok(dict) = object.downcast::<PyDict>() {
                Self::from_dict(dict)
            } else {
                Err(PyErr::new::<PyTypeError, _>(
                    "Argument claimed to be a dict but wasn't",
                ))
            }
        } else {
            Err(PyErr::new::<PyTypeError, _>(
                "Argument was not a properly-formed dict, list, or tuple",
            ))
        }
    }

    /// Converts a vector of Python objects and priorities into a vector of items
    fn from_vec(list: Vec<(Py<PyAny>, f64)>) -> Vec<(PyItem, f64)> {
        list.into_iter()
            .map(|(data, priority)| (PyItem::new(data), priority))
            .collect()
    }

    /// Converts a Python dictionary into a vector of items
    fn from_dict(dict: &Bound<'py, PyDict>) -> PyResult<Vec<(PyItem, f64)>> {
        if let Ok(items) = dict
            .into_iter()
            .map(|(data, priority)| match priority.extract::<f64>() {
                Ok(value) => Ok((PyItem::new(data.unbind()), value)),
                Err(err) => Err(err),
            })
            .collect::<Result<Vec<_>, _>>()
        {
            Ok(items)
        } else {
            Err(PyErr::new::<PyTypeError, _>("Dict keys were not numbers"))
        }
    }
}