//! Discovery of common node sets

use std::collections::HashSet;

use hg::{discovery::PartialDiscovery as CorePartialDiscovery, Revision};

use pyo3::{

    intern, pyclass, pymethods,

    types::{PyAnyMethods, PyDict, PyModule, PyModuleMethods, PyTuple},

    Bound, Py, PyAny, PyObject, PyResult, Python,

};

use pyo3_sharedref::SharedByPyObject;

use crate::{

    exceptions::GraphError,

    revision::{rev_pyiter_collect, PyRevision},

    revlog::PySharedIndex,

    utils::{new_submodule, py_rust_index_to_graph},

};

#[pyclass]

struct PartialDiscovery {

    inner: SharedByPyObject<CorePartialDiscovery<PySharedIndex>>,

    idx: SharedByPyObject<PySharedIndex>,

}

#[pymethods]

impl PartialDiscovery {

    #[pyo3(signature = (repo, targetheads, respectsize, randomize=true))]

    #[new]

    fn new(

        py: Python,

        repo: &Bound<'_, PyAny>,

        targetheads: &Bound<'_, PyAny>,

        respectsize: bool,

        randomize: bool,

    ) -> PyResult<Self> {

        let index = repo

            .getattr(intern!(py, "changelog"))?

            .getattr(intern!(py, "index"))?;

        let cloned_index = py_rust_index_to_graph(&index.clone())?;

        let index = py_rust_index_to_graph(&index)?;

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let target_heads = {

            let borrowed_idx = unsafe { index.try_borrow(py)? };

            rev_pyiter_collect(targetheads, &*borrowed_idx)?

        };

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let lazy_disco = unsafe {

            index.map(py, |idx| {

                CorePartialDiscovery::new(

                    idx,

                    target_heads,

                    respectsize,

                    randomize,

                )

            })

        };

        Ok(Self {

            inner: lazy_disco,

            idx: cloned_index,

        })

    }

    fn addcommons(

        &mut self,

        py: Python,

        commons: &Bound<'_, PyAny>,

    ) -> PyResult<PyObject> {

        let commons = self.pyiter_to_vec(commons)?;

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };

        inner

            .add_common_revisions(commons)

            .map_err(GraphError::from_hg)?;

        Ok(py.None())

    }

    fn addmissings(

        &mut self,

        py: Python,

        missings: &Bound<'_, PyAny>,

    ) -> PyResult<PyObject> {

        let missings = self.pyiter_to_vec(missings)?;

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };

        inner

            .add_missing_revisions(missings)

            .map_err(GraphError::from_hg)?;

        Ok(py.None())

    }

    fn addinfo(

        &mut self,

        py: Python,

        sample: &Bound<'_, PyAny>,

    ) -> PyResult<PyObject> {

        let mut missing: Vec<Revision> = vec![];

        let mut common: Vec<Revision> = vec![];

        for info in sample.try_iter()? {

            // info is a pair (Revision, bool)

            let info = info?;

            let info = info.downcast::<PyTuple>()?;

            let rev: PyRevision = info.get_item(0)?.extract()?;

            // This is fine since we're just using revisions as integers

            // for the purposes of discovery

            let rev = Revision(rev.0);

            let known: bool = info.get_item(1)?.extract()?;

            if known {

                common.push(rev);

            } else {

                missing.push(rev);

            }

        }

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };

        inner

            .add_common_revisions(common)

            .map_err(GraphError::from_hg)?;

        inner

            .add_missing_revisions(missing)

            .map_err(GraphError::from_hg)?;

        Ok(py.None())

    }

    fn hasinfo(&self, py: Python<'_>) -> PyResult<bool> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let inner = unsafe { self.inner.try_borrow(py)? };

        Ok(inner.has_info())

    }

    fn iscomplete(&self, py: Python<'_>) -> PyResult<bool> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let inner = unsafe { self.inner.try_borrow(py)? };

        Ok(inner.is_complete())

    }

    fn stats(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let inner = unsafe { self.inner.try_borrow(py)? };

        let stats = inner.stats();

        let as_dict = PyDict::new(py);

        as_dict.set_item("undecided", stats.undecided)?;

        Ok(as_dict.unbind())

    }

    fn commonheads(&self, py: Python<'_>) -> PyResult<HashSet<PyRevision>> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let inner = unsafe { self.inner.try_borrow(py)? };

        let common_heads =

            inner.common_heads().map_err(GraphError::from_hg)?;

        Ok(common_heads.into_iter().map(Into::into).collect())

    }

    fn takefullsample(

        &mut self,

        py: Python,

        _headrevs: &Bound<'_, PyAny>,

        size: usize,

    ) -> PyResult<Py<PyTuple>> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };

        let sample =

            inner.take_full_sample(size).map_err(GraphError::from_hg)?;

        let as_pyrevision = sample.into_iter().map(|rev| PyRevision(rev.0));

        Ok(PyTuple::new(py, as_pyrevision)?.unbind())

    }

    fn takequicksample(

        &mut self,

        py: Python,

        headrevs: &Bound<'_, PyAny>,

        size: usize,

    ) -> PyResult<Py<PyTuple>> {

        let revs = self.pyiter_to_vec(headrevs)?;

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };

        let sample = inner

            .take_quick_sample(revs, size)

            .map_err(GraphError::from_hg)?;

        let as_pyrevision = sample.into_iter().map(|rev| PyRevision(rev.0));

        Ok(PyTuple::new(py, as_pyrevision)?.unbind())

    }

}

impl PartialDiscovery {

    /// Convert a Python iterator of revisions into a vector

    fn pyiter_to_vec(

        &self,

        iter: &Bound<'_, PyAny>,

    ) -> PyResult<Vec<Revision>> {

        // Safety: we don't leak any reference derived form the "faked" one in

        // `SharedByPyObject`

        let index = unsafe { self.idx.try_borrow(iter.py())? };

        rev_pyiter_collect(iter, &*index)

    }

}

pub fn init_module<'py>(

    py: Python<'py>,

    package: &str,

) -> PyResult<Bound<'py, PyModule>> {

    let m = new_submodule(py, package, "discovery")?;

    m.add_class::<PartialDiscovery>()?;

    Ok(m)

}