/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::num::NonZeroU64;
use std::sync::Arc;
use std::time::Duration;

use async_runtime::block_unless_interrupted;
use cpython::*;
use cpython_async::PyFuture;
use cpython_async::TStream;
use cpython_ext::convert::Serde;
use cpython_ext::ExtractInner;
use cpython_ext::ExtractInnerRef;
use cpython_ext::PyCell;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use edenapi::Builder;
use edenapi::EdenApi;
use edenapi_ext::check_files;
use edenapi_ext::download_files;
use edenapi_ext::upload_snapshot;
use edenapi_types::AlterSnapshotRequest;
use edenapi_types::AlterSnapshotResponse;
use edenapi_types::AnyFileContentId;
use edenapi_types::BlameResult;
use edenapi_types::CommitGraphEntry;
use edenapi_types::CommitGraphSegmentsEntry;
use edenapi_types::CommitHashLookupResponse;
use edenapi_types::CommitHashToLocationResponse;
use edenapi_types::CommitId;
use edenapi_types::CommitIdScheme;
use edenapi_types::CommitKnownResponse;
use edenapi_types::CommitLocationToHashResponse;
use edenapi_types::CommitRevlogData;
use edenapi_types::CommitTranslateIdResponse;
use edenapi_types::EphemeralPrepareResponse;
use edenapi_types::FetchSnapshotRequest;
use edenapi_types::FetchSnapshotResponse;
use edenapi_types::FileResponse;
use edenapi_types::FileSpec;
use edenapi_types::FileType;
use edenapi_types::HgChangesetContent;
use edenapi_types::HgMutationEntryContent;
use edenapi_types::HistoryEntry;
use edenapi_types::Key;
use edenapi_types::LandStackResponse;
use edenapi_types::SnapshotRawData;
use edenapi_types::TreeAttributes;
use edenapi_types::TreeEntry;
use edenapi_types::UploadSnapshotResponse;
use edenapi_types::UploadToken;
use futures::TryStreamExt;
use minibytes::Bytes;
use pyconfigloader::config;
use pyrevisionstore::edenapifilestore;
use pyrevisionstore::edenapitreestore;
use revisionstore::EdenApiFileStore;
use revisionstore::EdenApiTreeStore;
use types::HgId;
use types::RepoPathBuf;

use crate::pyext::EdenApiPyExt;
use crate::stats::stats;
use crate::util::to_path;

// Python wrapper around an EdenAPI client.
//
// This is basically just FFI boilerplate. The actual functionality
// is implemented as the default implementations of the methods in
// the `EdenApiPyExt` trait.
py_class!(pub class client |py| {
    data inner: Arc<dyn EdenApi>;

    def __new__(
        _cls,
        config: config,
        correlator: Option<String> = None,
        reponame: Option<String> = None,
    ) -> PyResult<client> {
        let config = config.get_cfg(py);
        let inner = Builder::from_config(&config)
            .map_pyerr(py)?
            .correlator(correlator)
            .repo_name(reponame)
            .build()
            .map_pyerr(py)?;

        client::create_instance(py, inner)
    }

    /// The URL that is intended to match `paths.default`.
    def url(&self) -> PyResult<Option<String>> {
        Ok(self.inner(py).url())
    }

    def health(&self) -> PyResult<PyDict> {
        self.inner(py).as_ref().health_py(py)
    }

    def capabilities(&self) -> PyResult<Vec<String>> {
        let client = self.inner(py).as_ref();
        let caps = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    client.capabilities().await
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;
        Ok(caps)
    }

    def files(
        &self,
        keys: Vec<(PyPathBuf, Serde<HgId>)>
    ) -> PyResult<TStream<anyhow::Result<Serde<FileResponse>>>> {
        self.inner(py).as_ref().files_py(py, keys)
    }

    def filesattrs(
        &self,
        spec: Serde<Vec<FileSpec>>,
    ) -> PyResult<TStream<anyhow::Result<Serde<FileResponse>>>> {
        let api = self.inner(py).as_ref();
        let entries = py
            .allow_threads(|| block_unless_interrupted(api.files_attrs(spec.0)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(entries.map_ok(Serde).map_err(Into::into).into())
    }

    def history(
        &self,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        length: Option<u32> = None
    ) -> PyResult<TStream<anyhow::Result<Serde<HistoryEntry>>>> {
        self.inner(py).as_ref().history_py(py, keys, length)
    }

    def storetrees(
        &self,
        store: PyObject,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        attributes: Option<Serde<TreeAttributes>> = None
    ) -> PyResult<stats> {
        self.inner(py).as_ref().storetrees_py(py, store, keys, attributes.map(|a| a.0))
    }

    def trees(
        &self,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        attributes: Option<Serde<TreeAttributes>> = None
    ) -> PyResult<(TStream<anyhow::Result<Serde<TreeEntry>>>, PyFuture)> {
        self.inner(py).as_ref().trees_py(py, keys, attributes.map(|a| a.0))
    }

    /// commitdata(nodes: [bytes]) -> [(node: bytes, data: bytes)], stats
    ///
    /// Fetch commit data in raw HG format (sorted([p1, p2]) + text).
    def commitdata(
        &self,
        nodes: Serde<Vec<HgId>>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<CommitRevlogData>>>, PyFuture)> {
        self.inner(py).as_ref().commit_revlog_data_py(py, nodes.0)
    }

    /// bookmarks([name]) -> {name: node|None}
    ///
    /// Resolve remote bookmarks.
    def bookmarks(
        &self,
        bookmarks: Vec<String>
    ) -> PyResult<PyDict> {
        self.inner(py).as_ref().bookmarks_py(py, bookmarks)
    }

    /// setbookmark(bookmark, to, from, pushvars)
    ///
    /// Create, delete, or move a bookmark.
    def setbookmark(
        &self,
        bookmark: String,
        to: Serde<Option<HgId>>,
        from: Serde<Option<HgId>>,
        pushvars: Vec<(String, String)> = Vec::new(),
    ) -> PyResult<bool> {
        self.inner(py).as_ref().set_bookmark_py(py, bookmark, to.0, from.0, pushvars)
    }

    /// land_stack(bookmark, head, base, pushvars) -> {old_hgids: [node], new_hgids: [node]}
    ///
    /// Land a stack of already-uploaded commits by rebasing onto a bookmark and updating the bookmark.
    def landstack(
        &self,
        bookmark: String,
        head: Serde<HgId>,
        base: Serde<HgId>,
        pushvars: Vec<(String, String)> = Vec::new(),
    ) -> PyResult<Serde<LandStackResponse>> {
        self.inner(py).as_ref().land_stack_py(py, bookmark, head.0, base.0, pushvars)
    }

    /// hashlookup(hexprefix) -> [{'request': {'InclusiveRange': (start_node, end_node)},
    ///                            'hgids': [node]}]
    ///
    /// Lookup commit hashes by hex prefixes.
    def hashlookup(
        &self,
        hash_prefixes: Vec<String>
    ) -> PyResult<Serde<Vec<CommitHashLookupResponse>>> {
        self.inner(py).as_ref().hash_lookup_py(py, hash_prefixes)
    }

    def filestore(
        &self
    ) -> PyResult<edenapifilestore> {
        let edenapi = self.extract_inner(py);
        let store = EdenApiFileStore::new(edenapi);

        edenapifilestore::new(py, store)
    }

    def treestore(
        &self
    ) -> PyResult<edenapitreestore> {
        let edenapi = self.extract_inner(py);
        let store = EdenApiTreeStore::new(edenapi);

        edenapitreestore::new(py, store)
    }

    /// commitlocationtohash(requests: [(bytes, u64, u64)) ->
    ///   [(location: (descendant: bytes, distance: u64), count: u64, hgids: [bytes])]
    ///
    /// Fetch the hash(es) of a location in the commit graph.
    /// A request is a tuple (descendant, distance, count) where descendant is a known hgid,
    /// distance represents how many parents we traverse from descendant to the desired commit and
    /// count represents the number of ancestors from the location that we want.
    def commitlocationtohash(
        &self,
        requests: Serde<Vec<(HgId, u64, u64)>>,
    ) -> PyResult<Serde<Vec<CommitLocationToHashResponse>>> {
        self.inner(py).as_ref().commit_location_to_hash_py(py, requests.0)
    }

    /// commithashtolocation(master_heads: [bytes], hghds: [bytes]) ->
    ///   [(hgid: bytes, location: (descendant: bytes, distance: u64))]
    ///
    /// Fetch the location in the commit graph of a given hash.
    /// WARNING. Only hashes of ancestors of master are supported.
    /// It is necessary pass in the hashes that Segmented Changelog tracks in the master group in
    /// order for the server to be able to construct a valid location for the client.
    /// Hashes that cannot be found will be missing from the returned list.
    def commithashtolocation(
        &self,
        master_heads: Serde<Vec<HgId>>,
        hgids: Serde<Vec<HgId>>,
    ) -> PyResult<Serde<Vec<CommitHashToLocationResponse>>> {
        self.inner(py).as_ref().commit_hash_to_location_py(py, master_heads.0, hgids.0)
    }

    /// commitknown(nodes: [bytes]) -> [{'hgid': bytes, 'known': Result[bool]}]
    def commitknown(&self, hgids: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<CommitKnownResponse>>>
    {
        self.inner(py).as_ref().commit_known_py(py, hgids.0)
    }

    /// commitgraph2(heads: [bytes], common: [bytes]) -> [{'hgid': bytes, 'parents': [bytes], 'is_draft': Option[bool]}]
    def commitgraph2(&self, heads: Serde<Vec<HgId>>, common: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<CommitGraphEntry>>>
    {
        self.inner(py).as_ref().commit_graph2_py(py, heads.0, common.0)
    }

    /// commitgraphsegments(heads: [bytes], common: [bytes]) -> [{'head': bytes, 'base': bytes, 'parents': [bytes]}]
    def commitgraphsegments(&self, heads: Serde<Vec<HgId>>, common: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<CommitGraphSegmentsEntry>>>
    {
        self.inner(py).as_ref().commit_graph_segments_py(py, heads.0, common.0)
    }

    /// clonedata() -> PyCell
    def clonedata(&self) -> PyResult<PyCell> {
        self.inner(py).as_ref().clone_data_py(py)
    }

    /// pullfastforwardmaster(old_master: Bytes, new_master: Bytes) -> PyCell
    def pullfastforwardmaster(&self, old_master: Serde<HgId>, new_master: Serde<HgId>)
        -> PyResult<PyCell>
    {
        self.inner(py).as_ref().pull_fast_forward_master_py(py, old_master.0, new_master.0)
    }

    /// pulllazy(common: [node], missing: [node]) -> PyCell
    def pulllazy(&self, common: Serde<Vec<HgId>>, missing: Serde<Vec<HgId>>) -> PyResult<PyCell> {
        self.inner(py).as_ref().pull_lazy_py(py, common.0, missing.0)
    }

    /// lookup_file_contents(content_ids: [bytes])
    def lookup_file_contents(&self, content_ids: Vec<PyBytes>)
        -> PyResult<Serde<Vec<(usize, UploadToken)>>>
    {
        self.inner(py).as_ref().lookup_file_contents(py, content_ids)
    }

    /// lookup_commits(nodes: [bytes])
    def lookup_commits(&self, nodes: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<(usize, UploadToken)>>>
    {
        self.inner(py).as_ref().lookup_commits(py, nodes.0)
    }

    /// lookup_filenodes(filenodes: [bytes])
    def lookup_filenodes(&self, hgids: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<(usize, UploadToken)>>>
    {
        self.inner(py).as_ref().lookup_filenodes(py, hgids.0)
    }

    /// lookup_trees(trees: [bytes])
    def lookup_trees(&self, hgids: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<(usize, UploadToken)>>>
    {
        self.inner(py).as_ref().lookup_trees(py, hgids.0)
    }

    /// lookup_filenodes_and_trees(filenodes: [bytes], trees: [bytes])
    def lookup_filenodes_and_trees(&self, filenodes: Serde<Vec<HgId>>, trees: Serde<Vec<HgId>>)
        -> PyResult<Serde<Vec<(usize, UploadToken)>>>
    {
        self.inner(py).as_ref().lookup_filenodes_and_trees(py, filenodes.0, trees.0)
    }

    /// Upload file contents only
    def uploadfileblobs(
        &self,
        store: PyObject,
        keys: Vec<(
            PyPathBuf,   /* path */
            Serde<HgId>, /* hgid */
        )>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        self.inner(py).as_ref().uploadfileblobs_py(py, store, keys)
    }

    /// Upload file contents and hg filenodes
    def uploadfiles(
        &self,
        store: PyObject,
        keys: Vec<(
            PyPathBuf,     /* path */
            Serde<HgId>,   /* hgid */
            Serde<HgId>,   /* p1 */
            Serde<HgId>,   /* p2 */
        )>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        self.inner(py).as_ref().uploadfiles_py(py, store, keys)
    }

    /// Upload trees
    def uploadtrees(
        &self,
        items: Vec<(
            Serde<HgId>,  /* hgid */
            Serde<HgId>,  /* p1 */
            Serde<HgId>,  /* p2 */
            PyBytes,      /* data */
        )>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        self.inner(py).as_ref().uploadtrees_py(py, items)
    }

    /// Upload changesets
    /// This method sends a single request, batching should be done outside.
    def uploadchangesets(
        &self,
        changesets: Serde<Vec<(
            HgId,               /* hgid (node_id) */
            HgChangesetContent  /* changeset content */
        )>>,
        mutations: Vec<Serde<HgMutationEntryContent>>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        self.inner(py).as_ref().uploadchangesets_py(py, changesets.0, mutations)
    }

    def uploadsnapshot(
        &self,
        data: Serde<SnapshotRawData>,
        custom_duration_secs: Option<u64>,
        copy_from_bubble_id: Option<u64>,
        use_bubble: Option<u64>,
        labels: Option<Vec<String>>,
    ) -> PyResult<Serde<UploadSnapshotResponse>> {
        let api = self.inner(py).as_ref();
        let copy_from_bubble_id = copy_from_bubble_id.and_then(NonZeroU64::new);
        let use_bubble = use_bubble.and_then(NonZeroU64::new);
        py.allow_threads(|| {
            block_unless_interrupted(upload_snapshot(
                api,
                data.0,
                custom_duration_secs,
                copy_from_bubble_id,
                use_bubble,
                labels,
            ))
        })
        .map_pyerr(py)?
        .map_pyerr(py)
        .map(Serde)
    }

    /// Fetch snapshot information
    def fetchsnapshot(
        &self,
        data: Serde<FetchSnapshotRequest>,
    ) -> PyResult<Serde<FetchSnapshotResponse>> {
        self.inner(py).as_ref().fetchsnapshot_py(py, data)
    }

    /// Alter the properties of an existing snapshot
    def altersnapshot(
        &self,
        data: Serde<AlterSnapshotRequest>,
    ) -> PyResult<Serde<AlterSnapshotResponse>> {
        self.inner(py).as_ref().altersnapshot_py(py, data)
    }

    /// Downloads files from given upload tokens to given paths
    def downloadfiles(
        &self,
        root: Serde<RepoPathBuf>,
        // (path to download, content id)
        files: Vec<(PyPathBuf, Serde<UploadToken>, Serde<FileType>)>,
    ) -> PyResult<bool> {
        let api = self.inner(py).as_ref();
        let files = files
            .into_iter()
            .map(|(p, token, tp)| Ok((to_path(py, &p)?, token.0, tp.0)))
            .collect::<Result<Vec<_>, PyErr>>()?;
        py.allow_threads(|| block_unless_interrupted(download_files(api, &root.0, files)))
            .map_pyerr(py)?
            .map_pyerr(py)
            .map(|_| true)
    }

    /// Checks which files differ from
    def checkfiles(
        &self,
        root: Serde<RepoPathBuf>,
        files: Vec<(PyPathBuf, Serde<UploadToken>, Serde<FileType>)>
    ) -> PyResult<Vec<PyPathBuf>> {
        let files = files
            .into_iter()
            .map(|(p, token, tp)| Ok((to_path(py, &p)?, token.0, tp.0)))
            .collect::<Result<Vec<_>, PyErr>>()?;
        py.allow_threads(|| block_unless_interrupted(check_files(&root.0, files)))
            .map_pyerr(py)?
            .map_pyerr(py)
            .map(|v| v.into_iter().map(Into::into).collect())
    }

    /// Download file from given upload token to memory
    def downloadfiletomemory(
        &self,
        token: Serde<UploadToken>
    ) -> PyResult<PyBytes> {
        self.inner(py).as_ref().downloadfiletomemory_py(py, token)
    }

    def ephemeralprepare(&self, custom_duration: Option<u64>, labels: Option<Vec<String>>)
        -> PyResult<TStream<anyhow::Result<Serde<EphemeralPrepareResponse>>>>
    {
        let api = self.inner(py).as_ref();
        let entries = py
            .allow_threads(|| block_unless_interrupted(api.ephemeral_prepare(custom_duration.map(Duration::from_secs), labels)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(entries.map_ok(Serde).map_err(Into::into).into())
    }

    /// uploadfilecontents(data: [(AnyFileContentId, Bytes)], bubbleid: int | None)
    /// -> Iterable[UploadToken]
    def uploadfilecontents(
        &self,
        data: Serde<Vec<(AnyFileContentId, Bytes)>>,
        bubbleid: Option<u64> = None
    ) -> PyResult<TStream<anyhow::Result<Serde<UploadToken>>>> {
        let api = self.inner(py).as_ref();
        let bubble_id = bubbleid.and_then(NonZeroU64::new);
        let entries = py
            .allow_threads(|| block_unless_interrupted(api.process_files_upload(data.0, bubble_id, None)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(entries.map_ok(Serde).map_err(Into::into).into())
    }

    def commitmutations(
        &self,
        commits: Serde<Vec<HgId>>,
    ) -> PyResult<Serde<Vec<HgMutationEntryContent>>> {
        let api = self.inner(py).as_ref();
        py.allow_threads(|| block_unless_interrupted(api.commit_mutations(commits.0)))
            .map_pyerr(py)?
            .map_pyerr(py)
            .map(|responses| Serde(responses.into_iter().map(|r| r.mutation).collect()))
    }

    def committranslateids(
        &self,
        commits: Serde<Vec<CommitId>>,
        scheme: Serde<CommitIdScheme>,
        fromrepo: Option<String> = None,
        torepo: Option<String> = None,
    ) -> PyResult<TStream<anyhow::Result<Serde<CommitTranslateIdResponse>>>> {
        let api = self.inner(py).as_ref();
        let responses = py.allow_threads(|| block_unless_interrupted(api.commit_translate_id(commits.0, scheme.0, fromrepo, torepo)))
            .map_pyerr(py)?
            .map_pyerr(py)?;
        Ok(responses.entries.map_ok(Serde).map_err(Into::into).into())
    }

    def blame(
        &self,
        files: Serde<Vec<Key>>,
    ) -> PyResult<TStream<anyhow::Result<Serde<BlameResult>>>> {
        let api = self.inner(py).as_ref();
        let blames = py.allow_threads(|| block_unless_interrupted(api.blame(files.0)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(blames.map_ok(Serde).map_err(Into::into).into())
    }
});

impl ExtractInnerRef for client {
    type Inner = Arc<dyn EdenApi>;

    fn extract_inner_ref<'a>(&'a self, py: Python<'a>) -> &'a Self::Inner {
        self.inner(py)
    }
}

impl client {
    pub fn from_edenapi(py: Python, client: Arc<dyn EdenApi>) -> PyResult<Self> {
        Self::create_instance(py, client)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! pyedenapi - Python bindings for the Rust `edenapi` crate

#![allow(non_camel_case_types)]

use cpython::*;

mod client;
mod pyext;
mod pytypes;
mod stats;
mod util;

pub use client::client as PyClient;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "edenapi"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<client::client>(py)?;
    m.add_class::<stats::stats>(py)?;
    m.add(py, "correlator", py_fn!(py, correlator()))?;
    Ok(m)
}

fn correlator(_py: Python) -> PyResult<String> {
    Ok(String::from(edenapi::DEFAULT_CORRELATOR.as_str()))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::sync::Arc;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use async_runtime::block_unless_interrupted;
use cpython::*;
use cpython_async::PyFuture;
use cpython_async::TStream;
use cpython_ext::convert::Serde;
use cpython_ext::PyCell;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use dag_types::Location;
use dag_types::VertexName;
use edenapi::EdenApi;
use edenapi::EdenApiError;
use edenapi::Response;
use edenapi::Stats;
use edenapi_ext::calc_contentid;
use edenapi_types::AlterSnapshotRequest;
use edenapi_types::AlterSnapshotResponse;
use edenapi_types::AnyFileContentId;
use edenapi_types::AnyId;
use edenapi_types::CommitGraphEntry;
use edenapi_types::CommitGraphSegmentsEntry;
use edenapi_types::CommitHashLookupResponse;
use edenapi_types::CommitHashToLocationResponse;
use edenapi_types::CommitKnownResponse;
use edenapi_types::CommitLocationToHashRequest;
use edenapi_types::CommitLocationToHashResponse;
use edenapi_types::CommitRevlogData;
use edenapi_types::EdenApiServerError;
use edenapi_types::FetchSnapshotRequest;
use edenapi_types::FetchSnapshotResponse;
use edenapi_types::FileResponse;
use edenapi_types::HgChangesetContent;
use edenapi_types::HgFilenodeData;
use edenapi_types::HgMutationEntryContent;
use edenapi_types::HistoryEntry;
use edenapi_types::IndexableId;
use edenapi_types::LandStackResponse;
use edenapi_types::LookupResult;
use edenapi_types::TreeAttributes;
use edenapi_types::TreeEntry;
use edenapi_types::UploadHgChangeset;
use edenapi_types::UploadToken;
use futures::prelude::*;
use futures::stream;
use hgstore::separate_metadata;
use progress_model::ProgressBar;
use pyrevisionstore::as_legacystore;
use revisionstore::HgIdMutableDeltaStore;
use revisionstore::StoreKey;
use revisionstore::StoreResult;
use types::HgId;

use crate::pytypes::PyStats;
use crate::stats::stats;
use crate::util::as_deltastore;
use crate::util::meta_to_dict;
use crate::util::to_contentid;
use crate::util::to_keys;
use crate::util::to_keys_with_parents;
use crate::util::to_trees_upload_items;

/// Extension trait allowing EdenAPI methods to be called from Python code.
///
/// One nice benefit of making this a trait instead of directly implementing
/// the methods inside a `py_class!` macro invocation is that tools like
/// `rustfmt` can still parse the code.
pub trait EdenApiPyExt: EdenApi {
    fn health_py(&self, py: Python) -> PyResult<PyDict> {
        let meta = block_unless_interrupted(self.health())
            .map_pyerr(py)?
            .map_pyerr(py)?;
        meta_to_dict(py, &meta)
    }

    fn files_py(
        &self,
        py: Python,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
    ) -> PyResult<TStream<anyhow::Result<Serde<FileResponse>>>> {
        let keys = to_keys(py, &keys)?;
        let entries = py
            .allow_threads(|| block_unless_interrupted(self.files(keys)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(entries.map_ok(Serde).map_err(Into::into).into())
    }

    fn history_py(
        &self,
        py: Python,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        length: Option<u32>,
    ) -> PyResult<TStream<anyhow::Result<Serde<HistoryEntry>>>> {
        let keys = to_keys(py, &keys)?;
        let entries = py
            .allow_threads(|| block_unless_interrupted(self.history(keys, length)))
            .map_pyerr(py)?
            .map_pyerr(py)?
            .entries;
        Ok(entries.map_ok(Serde).map_err(Into::into).into())
    }

    fn storetrees_py(
        &self,
        py: Python,
        store: PyObject,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        attributes: Option<TreeAttributes>,
    ) -> PyResult<stats> {
        let keys = to_keys(py, &keys)?;
        let store = as_deltastore(py, store)?;

        let stats = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let prog = ProgressBar::register_new(
                        "Downloading trees over HTTP",
                        keys.len() as u64,
                        "trees",
                    );
                    let response = self.trees(keys, attributes).await?;
                    write_trees(response, store, prog).await
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        stats::new(py, stats)
    }

    fn trees_py(
        &self,
        py: Python,
        keys: Vec<(PyPathBuf, Serde<HgId>)>,
        attributes: Option<TreeAttributes>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<TreeEntry>>>, PyFuture)> {
        let keys = to_keys(py, &keys)?;
        let (trees, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self.trees(keys, attributes).await?;
                    Ok::<_, EdenApiError>((response.entries, response.stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let trees_py = trees.map(|t| match t {
            Ok(Ok(t)) => Ok(Serde(t)),
            Ok(Err(e)) => Err(e.into()),
            Err(e) => Err(e.into()),
        });
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((trees_py.into(), stats_py))
    }

    fn commit_revlog_data_py(
        &self,
        py: Python,
        nodes: Vec<HgId>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<CommitRevlogData>>>, PyFuture)> {
        let (commits, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self.commit_revlog_data(nodes).await?;
                    let commits = response.entries;
                    let stats = response.stats;
                    Ok::<_, EdenApiError>((commits, stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let commits_py = commits.map_ok(Serde).map_err(Into::into);
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((commits_py.into(), stats_py))
    }

    fn bookmarks_py(&self, py: Python, bookmarks: Vec<String>) -> PyResult<PyDict> {
        let response = py
            .allow_threads(|| block_unless_interrupted(self.bookmarks(bookmarks)))
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let bookmarks = PyDict::new(py);
        for entry in response {
            bookmarks.set_item(py, entry.bookmark, entry.hgid.map(|id| id.to_hex()))?;
        }
        Ok(bookmarks)
    }

    fn set_bookmark_py(
        &self,
        py: Python,
        bookmark: String,
        to: Option<HgId>,
        from: Option<HgId>,
        pushvars: Vec<(String, String)>,
    ) -> PyResult<bool> {
        py.allow_threads(|| {
            block_unless_interrupted(async move {
                self.set_bookmark(bookmark, to, from, pushvars.into_iter().collect())
                    .await?;
                Ok::<(), EdenApiError>(())
            })
        })
        .map_pyerr(py)?
        .map_pyerr(py)?;

        Ok(true)
    }

    fn land_stack_py(
        &self,
        py: Python,
        bookmark: String,
        head: HgId,
        base: HgId,
        pushvars: Vec<(String, String)>,
    ) -> PyResult<Serde<LandStackResponse>> {
        let response = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self
                        .land_stack(bookmark, head, base, pushvars.into_iter().collect())
                        .await?;
                    Ok::<_, EdenApiError>(response)
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        Ok(Serde(response))
    }

    fn hash_lookup_py(
        &self,
        py: Python,
        hash_prefixes: Vec<String>,
    ) -> PyResult<Serde<Vec<CommitHashLookupResponse>>> {
        let items = py
            .allow_threads(|| block_unless_interrupted(self.hash_prefixes_lookup(hash_prefixes)))
            .map_pyerr(py)?
            .map_pyerr(py)?;
        Ok(Serde(items))
    }

    fn commit_location_to_hash_py(
        &self,
        py: Python,
        requests: Vec<(HgId, u64, u64)>,
    ) -> PyResult<Serde<Vec<CommitLocationToHashResponse>>> {
        let requests = requests
            .into_iter()
            .map(|(hgid, distance, count)| {
                let location = Location::new(hgid, distance);
                CommitLocationToHashRequest { location, count }
            })
            .collect();
        let responses = py
            .allow_threads(|| block_unless_interrupted(self.commit_location_to_hash(requests)))
            .map_pyerr(py)?
            .map_pyerr(py)?;

        Ok(Serde(responses))
    }

    fn commit_hash_to_location_py(
        &self,
        py: Python,
        master_heads: Vec<HgId>,
        hgids: Vec<HgId>,
    ) -> PyResult<Serde<Vec<CommitHashToLocationResponse>>> {
        let responses = py
            .allow_threads(|| {
                block_unless_interrupted(self.commit_hash_to_location(master_heads, hgids))
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        Ok(Serde(responses))
    }

    fn commit_known_py(
        &self,
        py: Python,
        hgids: Vec<HgId>,
    ) -> PyResult<Serde<Vec<CommitKnownResponse>>> {
        let responses = py
            .allow_threads(|| block_unless_interrupted(self.commit_known(hgids)))
            .map_pyerr(py)?
            .map_pyerr(py)?;
        Ok(Serde(responses))
    }

    fn commit_graph2_py(
        &self,
        py: Python,
        heads: Vec<HgId>,
        common: Vec<HgId>,
    ) -> PyResult<Serde<Vec<CommitGraphEntry>>> {
        let responses = py
            .allow_threads(|| block_unless_interrupted(self.commit_graph(heads, common)))
            .map_pyerr(py)?
            .map_pyerr(py)?;

        Ok(Serde(responses))
    }

    fn commit_graph_segments_py(
        &self,
        py: Python,
        heads: Vec<HgId>,
        common: Vec<HgId>,
    ) -> PyResult<Serde<Vec<CommitGraphSegmentsEntry>>> {
        let responses = py
            .allow_threads(|| block_unless_interrupted(self.commit_graph_segments(heads, common)))
            .map_pyerr(py)?
            .map_pyerr(py)?;

        Ok(Serde(responses))
    }

    /// Get the "CloneData" serialized using mincode.
    fn clone_data_py(&self, py: Python) -> PyResult<PyCell> {
        let data = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    self.clone_data().await.map(|data| {
                        data.convert_vertex(|hgid| VertexName(hgid.as_ref().to_vec().into()))
                    })
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;
        PyCell::new(py, data)
    }

    /// Get pll data for master pull fast path
    fn pull_fast_forward_master_py(
        &self,
        py: Python,
        old_master: HgId,
        new_master: HgId,
    ) -> PyResult<PyCell> {
        let data = {
            py.allow_threads(|| {
                block_unless_interrupted(async move {
                    match self.pull_fast_forward_master(old_master, new_master).await {
                        Err(e) => Err(e),
                        Ok(data) => Ok(data.convert_vertex(|hgid| {
                            VertexName(hgid.into_byte_array().to_vec().into())
                        })),
                    }
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?
        };
        PyCell::new(py, data)
    }

    fn pull_lazy_py(&self, py: Python, common: Vec<HgId>, missing: Vec<HgId>) -> PyResult<PyCell> {
        let data = {
            py.allow_threads(|| {
                block_unless_interrupted(async move {
                    match self.pull_lazy(common, missing).await {
                        Err(e) => Err(e),
                        Ok(data) => Ok(data.convert_vertex(|hgid| {
                            VertexName(hgid.into_byte_array().to_vec().into())
                        })),
                    }
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?
        };
        PyCell::new(py, data)
    }

    fn lookup_file_contents(
        &self,
        py: Python,
        ids: Vec<PyBytes>,
    ) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        self.lookup_py(
            py,
            ids.into_iter()
                .map(|id| {
                    AnyId::AnyFileContentId(AnyFileContentId::ContentId(to_contentid(py, &id)))
                })
                .collect(),
        )
    }

    fn lookup_commits(
        &self,
        py: Python,
        nodes: Vec<HgId>,
    ) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        self.lookup_py(py, nodes.into_iter().map(AnyId::HgChangesetId).collect())
    }

    fn lookup_filenodes(
        &self,
        py: Python,
        ids: Vec<HgId>,
    ) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        self.lookup_py(py, ids.into_iter().map(AnyId::HgFilenodeId).collect())
    }

    fn lookup_trees(
        &self,
        py: Python,
        ids: Vec<HgId>,
    ) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        self.lookup_py(py, ids.into_iter().map(AnyId::HgTreeId).collect())
    }

    fn lookup_filenodes_and_trees(
        &self,
        py: Python,
        filenodes_ids: Vec<HgId>,
        trees_ids: Vec<HgId>,
    ) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        self.lookup_py(
            py,
            filenodes_ids
                .into_iter()
                .map(AnyId::HgFilenodeId)
                .chain(trees_ids.into_iter().map(AnyId::HgTreeId))
                .collect(),
        )
    }

    fn lookup_py(&self, py: Python, ids: Vec<AnyId>) -> PyResult<Serde<Vec<(usize, UploadToken)>>> {
        let responses = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    // Deduplicates input, keeping indexes
                    let indexable_ids: HashMap<IndexableId, Vec<usize>> = ids
                        .iter()
                        .cloned()
                        .map(|id| IndexableId {
                            id,
                            bubble_id: None,
                        })
                        .enumerate()
                        .fold(HashMap::new(), |mut map, (idx, id)| {
                            map.entry(id).or_default().push(idx);
                            map
                        });
                    self.lookup_batch(
                        indexable_ids.iter().map(|(id, _)| id.id.clone()).collect(),
                        None,
                        None,
                    )
                    .await
                    .map(|responses| {
                        responses
                            .into_iter()
                            .flat_map(|r| match r.result {
                                LookupResult::Present(token) => indexable_ids
                                    .get(&token.indexable_id())
                                    .cloned()
                                    .unwrap_or_default()
                                    .into_iter()
                                    .map(|idx| (idx, token.clone()))
                                    .collect(),
                                LookupResult::NotPresent(_) => vec![],
                            })
                            .collect()
                    })
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;
        Ok(Serde(responses))
    }

    fn uploadfileblobs_py(
        &self,
        py: Python,
        store: PyObject,
        keys: Vec<(PyPathBuf /* path */, Serde<HgId> /* hgid */)>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        let keys = to_keys(py, &keys)?;
        let store = as_legacystore(py, store)?;
        let downcast_error = "incorrect upload token, failed to downcast 'token.data.id' to 'AnyId::AnyFileContentId::ContentId' type";

        // Preupload LFS blobs
        store
            .upload(
                &keys
                    .iter()
                    .map(|key| StoreKey::hgid(key.clone()))
                    .collect::<Vec<_>>(),
            )
            .map_pyerr(py)?;

        let (content_ids, mut data): (Vec<_>, Vec<_>) = keys
            .into_iter()
            .map(|key| {
                let content = store.get_file_content(&key).map_pyerr(py)?;
                match content {
                    Some(v) => {
                        let content_id = calc_contentid(&v);
                        Ok((content_id, (content_id, v)))
                    }
                    None => Err(format_err!(
                        "failed to fetch file content for the key '{}'",
                        key
                    ))
                    .map_pyerr(py),
                }
            })
            // fail the entire operation if content is missing for some key because this is unexpected
            .collect::<Result<Vec<_>, _>>()?
            .into_iter()
            .unzip();

        // Deduplicate upload data
        let mut uniques = BTreeSet::new();
        data.retain(|(content_id, _)| uniques.insert(*content_id));
        let data = data
            .into_iter()
            .map(|(content_id, data)| (AnyFileContentId::ContentId(content_id), data))
            .collect();

        let (responses, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self.process_files_upload(data, None, None).await?;
                    let file_content_tokens = response
                        .entries
                        .try_collect::<Vec<_>>()
                        .await?
                        .into_iter()
                        .map(|token| {
                            let content_id = match token.data.id {
                                AnyId::AnyFileContentId(AnyFileContentId::ContentId(id)) => id,
                                _ => bail!(EdenApiError::Other(format_err!(downcast_error))),
                            };
                            Ok((content_id, token))
                        })
                        .collect::<Result<BTreeMap<_, _>, _>>()?;
                    let tokens = content_ids.into_iter().map(move |cid|
                        Result::<_, EdenApiError>::Ok(
                            file_content_tokens
                            .get(&cid)
                            .ok_or_else(|| EdenApiError::Other(format_err!("unexpected error: upload token is missing for ContentId({})", cid)))?.clone(),
                        )
                    );
                    Ok::<_, EdenApiError>((stream::iter(tokens), response.stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let responses_py = responses.map_ok(Serde).map_err(Into::into);
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((responses_py.into(), stats_py))
    }

    fn uploadfiles_py(
        &self,
        py: Python,
        store: PyObject,
        keys: Vec<(
            PyPathBuf,   /* path */
            Serde<HgId>, /* hgid */
            Serde<HgId>, /* p1 */
            Serde<HgId>, /* p2 */
        )>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        let keys = to_keys_with_parents(py, &keys)?;
        let store = as_legacystore(py, store)?;

        // Preupload LFS blobs
        store
            .upload(
                &keys
                    .iter()
                    .map(|(key, _)| StoreKey::hgid(key.clone()))
                    .collect::<Vec<_>>(),
            )
            .map_pyerr(py)?;

        let (mut upload_data, filenodes_data): (Vec<_>, Vec<_>) = keys
            .into_iter()
            .map(|(key, parents)| {
                let content = store.get(StoreKey::hgid(key.clone())).map_pyerr(py)?;
                match content {
                    StoreResult::Found(raw_content) => {
                        let raw_content = raw_content.into();
                        let (raw_data, copy_from) =
                            separate_metadata(&raw_content).map_pyerr(py)?;
                        let content_id = calc_contentid(&raw_data);
                        Ok((
                            (content_id, raw_data),
                            (key.hgid, content_id, parents, copy_from),
                        ))
                    }
                    _ => Err(EdenApiError::Other(format_err!(
                        "failed to fetch file content for the key '{}'",
                        key
                    )))
                    .map_pyerr(py),
                }
            })
            // fail the entire operation if content is missing for some key (filenode) because this is unexpected
            .collect::<Result<Vec<_>, _>>()?
            .into_iter()
            .unzip();

        // Deduplicate upload data
        let mut uniques = BTreeSet::new();
        upload_data.retain(|(content_id, _)| uniques.insert(*content_id));
        let upload_data = upload_data
            .into_iter()
            .map(|(content_id, data)| (AnyFileContentId::ContentId(content_id), data))
            .collect();

        let (responses, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let downcast_error = "incorrect upload token, failed to downcast 'token.data.id' to 'AnyId::AnyFileContentId::ContentId' type";
                    // upload file contents first, receiving upload tokens
                    let file_content_tokens = self
                        .process_files_upload(upload_data, None, None)
                        .await?
                        .entries
                        .try_collect::<Vec<_>>()
                        .await?
                        .into_iter()
                        .map(|token| {
                            let content_id = match token.data.id {
                                AnyId::AnyFileContentId(AnyFileContentId::ContentId(id)) => id,
                                _ => bail!(EdenApiError::Other(format_err!(downcast_error))),
                            };
                            Ok((content_id, token))
                        })
                        .collect::<Result<BTreeMap<_, _>, _>>()?;

                    // build the list of HgFilenodeData for upload
                    let filenodes_data = filenodes_data.into_iter().map(|(node_id, content_id, parents, copy_from)| {
                        let file_content_upload_token = file_content_tokens.get(&content_id).ok_or_else(|| EdenApiError::Other(format_err!("unexpected error: upload token is missing for ContentId({})", content_id)))?.clone();
                        Ok(HgFilenodeData {
                            node_id,
                            parents,
                            file_content_upload_token,
                            metadata: copy_from.to_vec(),
                        })
                    }).collect::<Result<Vec<_>, EdenApiError>>()?;

                    // upload hg filenodes
                    let response = self
                        .upload_filenodes_batch(filenodes_data)
                        .await?;

                    Ok::<_, EdenApiError>((response.entries, response.stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let responses_py = responses.map_ok(|r| Serde(r.token)).map_err(Into::into);
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((responses_py.into(), stats_py))
    }

    /// Upload trees
    fn uploadtrees_py(
        &self,
        py: Python,
        items: Vec<(
            Serde<HgId>, /* hgid */
            Serde<HgId>, /* p1 */
            Serde<HgId>, /* p2 */
            PyBytes,     /* data */
        )>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        let items = to_trees_upload_items(py, &items)?;
        let (responses, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self.upload_trees_batch(items).await?;
                    Ok::<_, EdenApiError>((response.entries, response.stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let responses_py = responses.map_ok(|r| Serde(r.token)).map_err(Into::into);
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((responses_py.into(), stats_py))
    }

    /// Upload changesets
    /// This method sends a single request, batching should be done outside.
    fn uploadchangesets_py(
        &self,
        py: Python,
        changesets: Vec<(
            HgId,               /* hgid (node_id) */
            HgChangesetContent, /* changeset content */
        )>,
        mutations: Vec<Serde<HgMutationEntryContent>>,
    ) -> PyResult<(TStream<anyhow::Result<Serde<UploadToken>>>, PyFuture)> {
        let changesets = changesets
            .into_iter()
            .map(|(node_id, content)| UploadHgChangeset {
                node_id,
                changeset_content: content,
            })
            .collect();
        let mutations = mutations.into_iter().map(|m| m.0).collect();
        let (responses, stats) = py
            .allow_threads(|| {
                block_unless_interrupted(async move {
                    let response = self.upload_changesets(changesets, mutations).await?;
                    Ok::<_, EdenApiError>((response.entries, response.stats))
                })
            })
            .map_pyerr(py)?
            .map_pyerr(py)?;

        let responses_py = responses.map_ok(|r| Serde(r.token)).map_err(Into::into);
        let stats_py = PyFuture::new(py, stats.map_ok(PyStats))?;
        Ok((responses_py.into(), stats_py))
    }

    fn fetchsnapshot_py(
        &self,
        py: Python,
        data: Serde<FetchSnapshotRequest>,
    ) -> PyResult<Serde<FetchSnapshotResponse>> {
        py.allow_threads(|| {
            block_unless_interrupted(async move {
                let cs_id = data.0.cs_id;
                self.fetch_snapshot(data.0)
                    .await?
                    .entries
                    .next()
                    .await
                    .with_context(|| format_err!("Failed to find snapshot {}", cs_id))?
            })
        })
        .map_pyerr(py)?
        .map_pyerr(py)
        .map(Serde)
    }

    fn altersnapshot_py(
        &self,
        py: Python,
        data: Serde<AlterSnapshotRequest>,
    ) -> PyResult<Serde<AlterSnapshotResponse>> {
        py.allow_threads(|| {
            block_unless_interrupted(async move {
                let cs_id = data.0.cs_id;
                self.alter_snapshot(data.0)
                    .await?
                    .entries
                    .next()
                    .await
                    .with_context(|| format_err!("Failed to alter snapshot {}", cs_id))?
            })
        })
        .map_pyerr(py)?
        .map_pyerr(py)
        .map(Serde)
    }

    fn downloadfiletomemory_py(&self, py: Python, token: Serde<UploadToken>) -> PyResult<PyBytes> {
        py.allow_threads(|| block_unless_interrupted(self.download_file(token.0)))
            .map_pyerr(py)?
            .map_pyerr(py)
            .map(|data| PyBytes::new(py, &data))
    }
}

impl<T: EdenApi + ?Sized> EdenApiPyExt for T {}

async fn write_trees(
    mut response: Response<Result<TreeEntry, EdenApiServerError>>,
    store: Arc<dyn HgIdMutableDeltaStore>,
    prog: Arc<ProgressBar>,
) -> Result<Stats, EdenApiError> {
    while let Some(Ok(entry)) = response.entries.try_next().await? {
        store.add_tree(&entry)?;
        prog.increase_position(1);
    }
    response.stats.await
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! Python types that implements `ToPyObject`.

use cpython::*;
use edenapi::Stats;

use crate::stats::stats;

/// Converts `Stats` to Python `stats`.
pub struct PyStats(pub Stats);

impl ToPyObject for PyStats {
    type ObjectType = stats;

    fn to_py_object(&self, py: Python) -> Self::ObjectType {
        stats::new(py, self.0.clone()).unwrap()
    }

    fn into_py_object(self, py: Python) -> Self::ObjectType {
        stats::new(py, self.0).unwrap()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::*;
use edenapi::Stats;

py_class!(pub class stats |py| {
    data stats: Stats;

    def __str__(&self) -> PyResult<String> {
        Ok(self.stats(py).to_string())
    }

    def downloaded(&self) -> PyResult<usize> {
        Ok(self.stats(py).downloaded)
    }

    def uploaded(&self) -> PyResult<usize> {
        Ok(self.stats(py).uploaded)
    }

    def requests(&self) -> PyResult<usize> {
        Ok(self.stats(py).requests)
    }

    def time_in_seconds(&self) -> PyResult<f64> {
        Ok(self.stats(py).time_in_seconds())
    }

    def time_in_millis(&self) -> PyResult<usize> {
        Ok(self.stats(py).time.as_millis() as usize)
    }

    def latency_in_millis(&self) -> PyResult<usize> {
        Ok(self.stats(py).latency.as_millis() as usize)
    }

    def bytes_per_second(&self) -> PyResult<f64> {
        Ok(self.stats(py).bytes_per_second())
    }
});

impl stats {
    pub fn new(py: Python, stats: Stats) -> PyResult<Self> {
        Self::create_instance(py, stats)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use cpython::*;
use cpython_ext::convert::Serde;
use cpython_ext::ExtractInner;
use cpython_ext::PyPath;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use edenapi::ResponseMeta;
use edenapi_types::ContentId;
use edenapi_types::UploadTreeEntry;
use pyrevisionstore::mutabledeltastore;
use revisionstore::HgIdMutableDeltaStore;
use types::HgId;
use types::Key;
use types::Parents;
use types::RepoPathBuf;

pub fn to_path(py: Python, name: &PyPath) -> PyResult<RepoPathBuf> {
    name.to_repo_path()
        .map_pyerr(py)
        .map(|path| path.to_owned())
}

pub fn to_contentid(py: Python, content_id: &PyBytes) -> ContentId {
    let mut bytes = [0u8; ContentId::len()];
    bytes.copy_from_slice(&content_id.data(py)[0..ContentId::len()]);
    ContentId::from(bytes)
}

pub fn to_key(py: Python, path: &PyPath, hgid: HgId) -> PyResult<Key> {
    let path = to_path(py, path)?;
    Ok(Key::new(path, hgid))
}

pub fn to_key_with_parents(
    py: Python,
    path: &PyPath,
    hgid: HgId,
    p1: HgId,
    p2: HgId,
) -> PyResult<(Key, Parents)> {
    let path = to_path(py, path)?;
    Ok((Key::new(path, hgid), Parents::new(p1, p2)))
}

pub fn to_trees_upload_item(
    py: Python,
    hgid: HgId,
    p1: HgId,
    p2: HgId,
    data: &PyBytes,
) -> PyResult<UploadTreeEntry> {
    Ok(UploadTreeEntry {
        node_id: hgid,
        data: data.data(py).to_vec(),
        parents: Parents::new(p1, p2),
    })
}

pub fn to_keys<'a>(
    py: Python,
    keys: impl IntoIterator<Item = &'a (PyPathBuf, Serde<HgId>)>,
) -> PyResult<Vec<Key>> {
    keys.into_iter()
        .map(|(path, hgid)| to_key(py, path, hgid.0))
        .collect()
}

pub fn to_keys_with_parents<'a>(
    py: Python,
    keys: impl IntoIterator<Item = &'a (PyPathBuf, Serde<HgId>, Serde<HgId>, Serde<HgId>)>,
) -> PyResult<Vec<(Key, Parents)>> {
    keys.into_iter()
        .map(|(path, hgid, p1, p2)| to_key_with_parents(py, path, hgid.0, p1.0, p2.0))
        .collect()
}

pub fn to_trees_upload_items<'a>(
    py: Python,
    items: impl IntoIterator<Item = &'a (Serde<HgId>, Serde<HgId>, Serde<HgId>, PyBytes)>,
) -> PyResult<Vec<UploadTreeEntry>> {
    items
        .into_iter()
        .map(|(hgid, p1, p2, data)| to_trees_upload_item(py, hgid.0, p1.0, p2.0, data))
        .collect()
}

pub fn as_deltastore(py: Python, store: PyObject) -> PyResult<Arc<dyn HgIdMutableDeltaStore>> {
    Ok(store.extract::<mutabledeltastore>(py)?.extract_inner(py))
}

pub fn meta_to_dict(py: Python, meta: &ResponseMeta) -> PyResult<PyDict> {
    let dict = PyDict::new(py);
    dict.set_item(py, "version", format!("{:?}", &meta.version))?;
    dict.set_item(py, "status", meta.status.as_u16())?;
    dict.set_item(py, "server", &meta.server)?;
    dict.set_item(py, "request_id", &meta.request_id)?;
    dict.set_item(py, "tw_task_handle", &meta.tw_task_handle)?;
    dict.set_item(py, "tw_task_version", &meta.tw_task_version)?;
    dict.set_item(py, "tw_canary_id", &meta.tw_canary_id)?;
    dict.set_item(py, "server_load", &meta.server_load)?;
    dict.set_item(py, "content_length", &meta.content_length)?;
    dict.set_item(py, "content_encoding", &meta.content_encoding)?;
    Ok(dict)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::*;
use cpython_ext::error;

py_exception!(error, CertificateError);
py_exception!(error, CommitLookupError, exc::KeyError);
py_exception!(error, ConfigError);
py_exception!(error, FetchError, exc::KeyError);
py_exception!(error, HttpError);
py_exception!(error, IndexedLogError);
py_exception!(error, InvalidRepoPath);
py_exception!(error, LockContendedError);
py_exception!(error, MetaLogError);
py_exception!(error, NeedSlowPathError);
py_exception!(error, NonUTF8Path);
py_exception!(error, WorkingCopyError);
py_exception!(error, RepoInitError);
py_exception!(error, UncategorizedNativeError);
py_exception!(error, TlsError);

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "error"].join(".");
    let m = PyModule::new(py, &name)?;

    m.add(py, "CertificateError", py.get_type::<CertificateError>())?;
    m.add(py, "CommitLookupError", py.get_type::<CommitLookupError>())?;
    m.add(py, "FetchError", py.get_type::<FetchError>())?;
    m.add(py, "HttpError", py.get_type::<HttpError>())?;
    m.add(py, "IndexedLogError", py.get_type::<IndexedLogError>())?;
    m.add(py, "InvalidRepoPath", py.get_type::<InvalidRepoPath>())?;
    m.add(
        py,
        "LockContendedError",
        py.get_type::<LockContendedError>(),
    )?;
    m.add(py, "MetaLogError", py.get_type::<MetaLogError>())?;
    m.add(py, "NeedSlowPathError", py.get_type::<NeedSlowPathError>())?;
    m.add(
        py,
        "UncategorizedNativeError",
        py.get_type::<UncategorizedNativeError>(),
    )?;
    m.add(py, "WorkingCopyError", py.get_type::<WorkingCopyError>())?;
    m.add(py, "RepoInitError", py.get_type::<RepoInitError>())?;
    m.add(py, "ConfigError", py.get_type::<ConfigError>())?;
    m.add(py, "NonUTF8Path", py.get_type::<NonUTF8Path>())?;
    m.add(py, "TlsError", py.get_type::<TlsError>())?;

    register_error_handlers();

    Ok(m)
}

fn register_error_handlers() {
    fn specific_error_handler(py: Python, mut e: &error::Error) -> Option<PyErr> {
        // We care about concrete errors, so peel away anyhow contextual layers.
        while let Some(inner) = e.downcast_ref::<error::Error>() {
            e = inner;
        }

        // Extract inner io::Error out.
        // Why does Python need the low-level IOError? It doesn't have to.
        // Consider:
        // - Only expose high-level errors to Python with just enough
        //   information that Python can consume. Python no longer handles
        //   IOError directly.
        // - Gain more explicit control about error types exposed to
        //   Python. This means dropping anyhow.
        if let Some(e) = e.downcast_ref::<std::io::Error>() {
            return Some(cpython_ext::error::translate_io_error(py, e));
        }

        if let Some(revlogindex::Error::Corruption(e)) = e.downcast_ref::<revlogindex::Error>() {
            if let revlogindex::errors::CorruptionError::Io(e) = e.as_ref() {
                return Some(cpython_ext::error::translate_io_error(py, e));
            }
        }

        let mut dag_error = None;
        if let Some(e) = e.downcast_ref::<dag::Error>() {
            dag_error = Some(e);
        } else if let Some(hgcommits::Error::Dag(e)) = e.downcast_ref::<hgcommits::Error>() {
            dag_error = Some(e);
        }

        if let Some(e) = dag_error {
            match e {
                dag::Error::Backend(ref backend_error) => match backend_error.as_ref() {
                    dag::errors::BackendError::Io(e) => {
                        return Some(cpython_ext::error::translate_io_error(py, &e));
                    }
                    dag::errors::BackendError::Other(e) => return specific_error_handler(py, e),
                    _ => {}
                },
                dag::Error::VertexNotFound(_) | dag::Error::IdNotFound(_) => {
                    return Some(PyErr::new::<CommitLookupError, _>(
                        py,
                        cpython_ext::Str::from(e.to_string()),
                    ));
                }
                dag::Error::NeedSlowPath(_) => {
                    return Some(PyErr::new::<NeedSlowPathError, _>(
                        py,
                        cpython_ext::Str::from(e.to_string()),
                    ));
                }
                _ => {}
            }
        }

        if let Some(e) = e.downcast_ref::<repolock::LockError>() {
            match e {
                repolock::LockError::Contended(repolock::LockContendedError {
                    contents, ..
                }) => {
                    return Some(PyErr::new::<LockContendedError, _>(
                        py,
                        cpython_ext::Str::from(contents.clone()),
                    ));
                }
                repolock::LockError::Io(e) => {
                    return Some(cpython_ext::error::translate_io_error(py, &e.to_io_err()));
                }
                _ => {}
            };
        }

        if e.is::<indexedlog::Error>() {
            Some(PyErr::new::<IndexedLogError, _>(
                py,
                cpython_ext::Str::from(format!("{:?}", e)),
            ))
        } else if e.is::<metalog::Error>() {
            Some(PyErr::new::<MetaLogError, _>(
                py,
                cpython_ext::Str::from(format!("{:?}", e)),
            ))
        } else if e.is::<configmodel::Error>() {
            Some(PyErr::new::<ConfigError, _>(
                py,
                cpython_ext::Str::from(format!("{:?}", e)),
            ))
        } else if matches!(
            e.downcast_ref::<dag::Error>(),
            Some(dag::Error::NeedSlowPath(_))
        ) {
            Some(PyErr::new::<NeedSlowPathError, _>(
                py,
                cpython_ext::Str::from(e.to_string()),
            ))
        } else if matches!(
            e.downcast_ref::<dag::Error>(),
            Some(dag::Error::VertexNotFound(_)) | Some(dag::Error::IdNotFound(_))
        ) {
            Some(PyErr::new::<CommitLookupError, _>(
                py,
                cpython_ext::Str::from(e.to_string()),
            ))
        } else if e.is::<repo::errors::InitError>() {
            Some(PyErr::new::<RepoInitError, _>(
                py,
                cpython_ext::Str::from(e.to_string()),
            ))
        } else if e.is::<repo::errors::InvalidWorkingCopy>() {
            Some(PyErr::new::<WorkingCopyError, _>(
                py,
                cpython_ext::Str::from(e.to_string()),
            ))
        } else if e.is::<cpython_ext::Error>() {
            Some(PyErr::new::<NonUTF8Path, _>(
                py,
                cpython_ext::Str::from(format!("{:?}", e)),
            ))
        } else if e.is::<types::path::ParseError>() {
            Some(PyErr::new::<InvalidRepoPath, _>(
                py,
                cpython_ext::Str::from(format!("{:?}", e)),
            ))
        } else if let Some(e) = e.downcast_ref::<edenapi::EdenApiError>() {
            match e {
                edenapi::EdenApiError::Http(http_client::HttpClientError::Tls(
                    http_client::TlsError { source: e, .. },
                )) => Some(PyErr::new::<TlsError, _>(
                    py,
                    cpython_ext::Str::from(e.to_string()),
                )),
                _ => Some(PyErr::new::<HttpError, _>(
                    py,
                    cpython_ext::Str::from(e.to_string()),
                )),
            }
        } else if e.is::<auth::MissingCerts>() {
            Some(PyErr::new::<CertificateError, _>(
                py,
                cpython_ext::Str::from(format!("{}", e)),
            ))
        } else if e.is::<auth::X509Error>() {
            Some(PyErr::new::<CertificateError, _>(
                py,
                cpython_ext::Str::from(format!("{}", e)),
            ))
        } else if let Some(e) = e.downcast_ref::<revisionstore::scmstore::KeyFetchError>() {
            use revisionstore::scmstore::KeyFetchError::*;
            if let Other(ref e) = e {
                specific_error_handler(py, e)
            } else {
                Some(PyErr::new::<FetchError, _>(
                    py,
                    cpython_ext::Str::from(format!("{}", e)),
                ))
            }
        } else if let Some(e) = e.downcast_ref::<types::errors::NetworkError>() {
            // If we don't handle inner error specifically, default to
            // HttpError which will trigger the network doctor.
            specific_error_handler(py, &e.0).or_else(|| {
                Some(PyErr::new::<HttpError, _>(
                    py,
                    cpython_ext::Str::from(e.0.to_string()),
                ))
            })
        } else if let Some(e) = e.downcast_ref::<cpython_ext::PyErr>() {
            Some(e.clone(py).into())
        } else {
            None
        }
    }

    fn fallback_error_handler(py: Python, e: &error::Error) -> Option<PyErr> {
        Some(PyErr::new::<UncategorizedNativeError, _>(
            py,
            format!("{:?}", e),
        ))
    }

    error::register("010-specific", specific_error_handler);
    error::register("999-fallback", fallback_error_handler);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use configmodel::Config;
use cpython::*;
use cpython_ext::convert::ImplInto;
use cpython_ext::error::ResultPyErrExt;
use cpython_ext::ExtractInner;
use cpython_ext::PyNone;
use pydag::commits::commits;
use pyedenapi::PyClient;
use pymetalog::metalog;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "pull"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(
        py,
        "clone",
        py_fn!(
            py,
            clone(
                config: ImplInto<Arc<dyn Config>>,
                edenapi: &PyClient,
                metalog: metalog,
                commits: &commits,
                bookmarks: Vec<String>,
            )
        ),
    )?;

    Ok(m)
}

fn clone(
    py: Python,
    config: ImplInto<Arc<dyn Config>>,
    edenapi: &PyClient,
    metalog: metalog,
    commits: &commits,
    bookmarks: Vec<String>,
) -> PyResult<PyNone> {
    let client = edenapi.extract_inner(py);
    let commits = commits.get_inner(py);
    let mut commits = commits.write();
    let meta = metalog.metalog_rwlock(py);
    let mut meta = meta.write();
    exchange::clone(&config.into(), client, &mut meta, &mut commits, bookmarks).map_pyerr(py)?;
    Ok(PyNone)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::io;

use cpython::*;
use cpython_ext::PyNone;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "fail"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(py, "failpoint", py_fn!(py, failpoint(name: &str)))?;
    Ok(m)
}

fn failpoint(py: Python, name: &str) -> PyResult<PyNone> {
    if let Some(e) = fail::eval(name, |_| fail_error(name)) {
        Err(cpython_ext::error::translate_io_error(py, &e).into())
    } else {
        Ok(PyNone)
    }
}

fn fail_error(name: &str) -> io::Error {
    io::Error::new(
        io::ErrorKind::Other,
        format!("failpoint '{}' set by FAILPOINTS", name),
    )
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::*;
use cpython_ext::PyPath;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "fs"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(py, "canonicalize", py_fn!(py, canonicalize(path: &str)))?;
    m.add(py, "fstype", py_fn!(py, fstype(path: &PyPath)))?;
    Ok(m)
}

fn canonicalize(py: Python, path: &str) -> PyResult<String> {
    let canonical_path = std::fs::canonicalize(path).map_pyerr(py)?;
    match canonical_path.to_str() {
        Some(s) => Ok(s.to_string()),
        None => {
            let message = format!("canonicalized({}) is non-utf-8", path);
            Err(PyErr::new::<exc::ValueError, _>(py, message))
        }
    }
}

fn fstype(py: Python, path: &PyPath) -> PyResult<Str> {
    let fstype = fsinfo::fstype(path).map_pyerr(py)?;
    Ok(fstype.to_string().into())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! Support `ImplInto` from cpython-ext.

use std::sync::Arc;

use cpython::*;
use cpython_ext::convert::register_into;
use storemodel::ReadFileContents;
use storemodel::TreeStore;

use crate::gitstore;

pub(crate) fn register(py: Python) {
    register_into(py, |py, g: gitstore| g.to_dyn_treestore(py));
    register_into(py, |py, g: gitstore| g.to_read_file_contents(py));
}

impl gitstore {
    fn to_dyn_treestore(&self, py: Python) -> Arc<dyn TreeStore + Send + Sync> {
        let store = self.inner(py);
        store.clone() as Arc<_>
    }

    fn to_read_file_contents(
        &self,
        py: Python,
    ) -> Arc<dyn ReadFileContents<Error = anyhow::Error> + Send + Sync> {
        let store = self.inner(py).clone();
        store as Arc<_>
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::sync::Arc;

use ::gitstore::git2;
use ::gitstore::GitStore;
use cpython::*;
use cpython_ext::convert::Serde;
use cpython_ext::PyPath;
use cpython_ext::ResultPyErrExt;
use storemodel::types::HgId;

mod impl_into;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "gitstore"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<gitstore>(py)?;
    impl_into::register(py);
    Ok(m)
}

py_class!(pub class gitstore |py| {
    data inner: Arc<GitStore>;

    def __new__(_cls, gitdir: &PyPath) -> PyResult<Self> {
        let store = GitStore::open(gitdir.as_path()).map_pyerr(py)?;
        Self::create_instance(py, Arc::new(store))
    }

    /// readobj(node, kind="any") -> bytes.
    /// Read a git object of the given type.
    def readobj(&self, node: Serde<HgId>, kind: &str = "any") -> PyResult<PyBytes> {
        let kind = str_to_object_type(py, kind)?;
        let data = self.inner(py).read_obj(node.0, kind).map_pyerr(py)?;
        Ok(PyBytes::new(py, &data))
    }

    /// readobjsize(node, kind="any") -> int.
    /// Read a git object size without reading its full content.
    def readobjsize(&self, node: Serde<HgId>, kind: &str = "any") -> PyResult<usize> {
        let kind = str_to_object_type(py, kind)?;
        let size = self.inner(py).read_obj_size(node.0, kind).map_pyerr(py)?;
        Ok(size)
    }

    /// writeobj(kind, data) -> node.
    /// Write object to the store. Not buffered in memory.
    /// Returns the SHA1 hash.
    def writeobj(&self, kind: &str, data: PyBytes) -> PyResult<Serde<HgId>> {
        let kind = str_to_object_type(py, kind)?;
        let node = self.inner(py).write_obj(kind, data.data(py)).map_pyerr(py)?;
        Ok(Serde(node))
    }
});

fn str_to_object_type(py: Python, kind: &str) -> PyResult<git2::ObjectType> {
    match git2::ObjectType::from_str(kind) {
        Some(v) => Ok(v),
        None => Err(PyErr::new::<exc::ValueError, _>(
            py,
            format!("invalid kind: {}", kind),
        )),
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::io::Write;

use cpython::*;
use cpython_ext::PyNone;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "hgmetrics"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(
        py,
        "incrementcounter",
        py_fn!(py, increment_counter(key: &str, value: usize = 1)),
    )?;
    m.add(py, "summarize", py_fn!(py, summarize()))?;

    m.add(
        py,
        "samplingcategory",
        py_fn!(py, sampling_category(key: &str)),
    )?;
    m.add(
        py,
        "appendsamples",
        py_fn!(py, append_samples(data: String)),
    )?;

    Ok(m)
}

fn increment_counter(_py: Python, key: &str, value: usize) -> PyResult<PyNone> {
    hg_metrics::increment_counter(key, value);
    Ok(PyNone)
}

fn summarize(_py: Python) -> PyResult<HashMap<Str, usize>> {
    let summary = hg_metrics::summarize();
    let summary: HashMap<Str, usize> = summary.into_iter().map(|(k, v)| (k.into(), v)).collect();
    Ok(summary)
}

fn sampling_category(_py: Python, key: &str) -> PyResult<Option<String>> {
    Ok(match sampling::CONFIG.get() {
        Some(Some(sc)) => sc.category(key).map(|cat| cat.to_string()),
        _ => None,
    })
}

fn append_samples(py: Python, data: String) -> PyResult<PyNone> {
    if let Some(Some(sc)) = sampling::CONFIG.get() {
        let mut file = sc.file();
        file.write_all(data.as_bytes()).map_pyerr(py)?;
        file.write_all(b"\0").map_pyerr(py)?;
    }
    Ok(PyNone)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::*;
use cpython_ext::PyNone;
use hgtime::HgTime;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "hgtime"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(py, "parse", py_fn!(py, parse(content: &str)))?;
    m.add(py, "parserange", py_fn!(py, parserange(content: &str)))?;
    m.add(
        py,
        "setnowfortesting",
        py_fn!(py, setnowfortesting(content: &str)),
    )?;

    // The Rust library indirectly uses libc for timezone handling.
    // Expose tzset() from libc so chg server can update timezone
    // environments when $TZ gets changed.
    m.add(py, "tzset", py_fn!(py, pytzset()))?;

    // Initialize. Useful on Windows.
    tzset();
    Ok(m)
}

fn parse(_py: Python, content: &str) -> PyResult<Option<(i64, i32)>> {
    Ok(HgTime::parse(content).map(|t| (t.unixtime, t.offset)))
}

fn parserange(_py: Python, content: &str) -> PyResult<Option<((i64, i32), (i64, i32))>> {
    Ok(HgTime::parse_range(content).map(|r| {
        (
            (r.start.unixtime, r.start.offset),
            (r.end.unixtime, r.end.offset),
        )
    }))
}

fn setnowfortesting(_py: Python, content: &str) -> PyResult<PyNone> {
    if let Some(time) = HgTime::parse(content) {
        time.set_as_now_for_testing();
    }
    Ok(PyNone)
}

fn pytzset(_: Python) -> PyResult<Option<i32>> {
    Ok(tzset())
}

/// Initialize timezone setting by reading the `TZ` environment variable.
/// Return the timezone offset in seconds (without daylight savings).
///
/// On Windows, the daylight saving handling is currently broken if `TZ`
/// is set. But settings from the Windows Control Panel (without `TZ`)
/// should work as expected.
fn tzset() -> Option<i32> {
    #[cfg(unix)]
    {
        use std::os::raw::c_long;
        extern "C" {
            // See https://www.gnu.org/software/libc/manual/html_node/Time-Zone-Functions.html#Time-Zone-Functions
            fn tzset();
            static timezone: c_long;
        }
        unsafe { tzset() };
        return Some(unsafe { timezone } as i32);
    }
    #[cfg(windows)]
    {
        use std::os::raw::c_int;
        use std::os::raw::c_long;
        extern "C" {
            // See https://docs.microsoft.com/en-us/cpp/c-runtime-library/daylight-dstbias-timezone-and-tzname?view=vs-2019
            fn _tzset();
            fn _get_timezone(seconds: *mut c_long) -> c_int;
        }
        unsafe { _tzset() };
        let mut timezone = 0;
        unsafe { _get_timezone(&mut timezone) };
        // The time-0.1.42 crate calls FileTimeToSystemTime on Windows which
        // does not respect $TZ. Ideally it should call functions from MSVC
        // runtime instead. To reduce issues, use `set_default_offset` to
        // tell `hgtime` to use a fixed timezone offset. Note: it does not
        // really work for complex timezones with daylight savings. But it
        // is good enough for TZ=GMT cases (which is used by tests).
        if std::env::var_os("TZ").is_some() {
            hgtime::set_default_offset(timezone);
        }
        return Some(timezone as i32);
    }
    #[allow(unreachable_code)]
    None
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

extern crate identity as rsident;

use cpython::*;
use cpython_ext::error::Result;
use cpython_ext::error::ResultPyErrExt;
use cpython_ext::PyPathBuf;
use rsident::Identity;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "identity"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<identity>(py)?;
    m.add(py, "all", py_fn!(py, all()))?;
    m.add(py, "default", py_fn!(py, default()))?;
    m.add(py, "envvar", py_fn!(py, try_env_var(suffix: PyString)))?;
    m.add(py, "resetdefault", py_fn!(py, reset_default()))?;
    m.add(py, "sniffroot", py_fn!(py, sniff_root(path: PyPathBuf)))?;
    m.add(py, "sniffdir", py_fn!(py, sniff_dir(path: PyPathBuf)))?;

    Ok(m)
}

py_class!(pub class identity |py| {
    data ident: Identity;

    def dotdir(&self) -> PyResult<String> {
        Ok(self.ident(py).dot_dir().to_string())
    }

    def __str__(&self) -> PyResult<String> {
        Ok(format!("{}", self.ident(py)))
    }

    def cliname(&self) -> PyResult<String> {
        Ok(self.ident(py).cli_name().to_string())
    }

    def configrepofile(&self) -> PyResult<String> {
        Ok(self.ident(py).config_repo_file().to_string())
    }

    def userconfigpaths(&self) -> PyResult<Vec<PyPathBuf>> {
        self.ident(py).user_config_paths().iter().map(|p| p.as_path().try_into()).collect::<Result<Vec<PyPathBuf>>>().map_pyerr(py)
    }

    def systemconfigpath(&self) -> PyResult<Option<PyPathBuf>> {
        match self.ident(py).system_config_path() {
            None => Ok(None),
            Some(p) => p.try_into().map(|p| Some(p)).map_pyerr(py),
        }
    }

    def productname(&self) -> PyResult<String> {
        Ok(self.ident(py).product_name().to_string())
    }

    def longproductname(&self) -> PyResult<String> {
        Ok(self.ident(py).long_product_name().to_string())
    }
});

fn sniff_root(py: Python, path: PyPathBuf) -> PyResult<Option<(PyPathBuf, identity)>> {
    Ok(match rsident::sniff_root(path.as_path()).map_pyerr(py)? {
        None => None,
        Some((path, ident)) => Some((
            path.try_into().map_pyerr(py)?,
            identity::create_instance(py, ident)?,
        )),
    })
}

fn sniff_dir(py: Python, path: PyPathBuf) -> PyResult<Option<identity>> {
    Ok(match rsident::sniff_dir(path.as_path()).map_pyerr(py)? {
        None => None,
        Some(ident) => Some(identity::create_instance(py, ident)?),
    })
}

fn try_env_var(py: Python, suffix: PyString) -> PyResult<Option<String>> {
    rsident::env_var(suffix.to_string(py)?.as_ref())
        .transpose()
        .map_pyerr(py)
}

fn default(py: Python) -> PyResult<identity> {
    identity::create_instance(py, rsident::default())
}

fn reset_default(_py: Python) -> PyResult<PyNone> {
    rsident::reset_default();
    Ok(PyNone)
}

fn all(py: Python) -> PyResult<Vec<identity>> {
    rsident::all()
        .iter()
        .map(|id| identity::create_instance(py, id.clone()))
        .collect()
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cell::RefCell;
use std::ops::Bound;

use cpython::*;
use cpython_ext::PyNone;
use cpython_ext::PyPath;
use cpython_ext::ResultPyErrExt;
use indexedlog::log::IndexDef;
use indexedlog::log::IndexOutput;
use indexedlog::log::LogLookupIter;
use indexedlog::log::LogRangeIter;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "indexedlog"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<Log>(py)?;
    m.add_class::<OpenOptions>(py)?;
    Ok(m)
}

py_class!(class OpenOptions |py| {
    data defs: RefCell<Vec<IndexDef>>;

    def __new__(_cls) -> PyResult<Self> {
        Self::create_instance(py, RefCell::new(Vec::new()))
    }

    /// index_def(name: str, func: Callable[[bytes], List[bytes | range]], lag_threshold=None)
    ///
    /// Append an index definition defined by `func`.
    /// `func` takes a byte slice (entry to insert)
    def index_def(&self, name: String, func: PyObject, lag_threshold: Option<u64> = None) -> PyResult<Self> {
        let mut defs = self.defs(py).borrow_mut();
        let func = to_index_func(py, func);
        let mut def = IndexDef::new(name, func);
        if let Some(lag) = lag_threshold {
            def = def.lag_threshold(lag);
        }
        defs.push(def);
        Ok(self.clone_ref(py))
    }

    /// open(path, crate=True) -> Log
    def open(&self, dir: &PyPath, create: bool = true) -> PyResult<Log> {
        let dir = dir.as_path();
        let defs = self.defs(py).borrow().clone();
        let opts = indexedlog::log::OpenOptions::new().create(create).index_defs(defs);
        let log = opts.open(dir).map_pyerr(py)?;
        Log::create_instance(py, RefCell::new(log))
    }
});

py_class!(class Log |py| {
    data log: RefCell<indexedlog::log::Log>;

    def __new__(_cls, path: &PyPath) -> PyResult<Self> {
        let index_defs = Vec::new();
        let log = indexedlog::log::Log::open(path, index_defs).map_pyerr(py)?;
        Self::create_instance(py, RefCell::new(log))
    }

    /// entries(skip=0, take=sys.maxsize, dirty=False) -> List[memoryview]
    ///
    /// Get all entries in the Log. The entries are wrapped in a zero-copy
    /// `Bytes` type that requires `asref()` to get the underlying bytes
    /// as a memoryview.
    def entries(&self, skip: usize=0, take: usize=usize::MAX, dirty: bool=false) -> PyResult<Vec<pybytes::Bytes>> {
        let log = self.log(py).borrow();
        let iter = if dirty { log.iter_dirty() } else { log.iter() };
        let items: Vec<&[u8]> = iter.skip(skip).take(take).collect::<Result<Vec<_>, _>>().map_pyerr(py)?;
        let items: Vec<pybytes::Bytes> = items.into_iter().map(|s| {
            pybytes::Bytes::from_bytes(py, log.slice_to_bytes(s))
        }).collect::<Result<_, _>>()?;
        Ok(items)
    }

    /// Append an entry to the Log.
    def append(&self, data: PyBytes) -> PyResult<PyNone> {
        let mut log = self.log(py).borrow_mut();
        let data = data.data(py);
        log.append(data).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// lookup(index_id: int, key: bytes) -> List[bytes]
    def lookup(&self, index_id: usize, key: PyBytes) -> PyResult<Vec<PyBytes>> {
        let log = self.log(py).borrow();
        let key = key.data(py);
        let iter = log.lookup(index_id, key).map_pyerr(py)?;
        let result = lookup_iter_to_vec(py, iter).map_pyerr(py)?;
        Ok(result)
    }

    /// lookup_prefix(index_id: int, prefix: bytes) -> List[Tuple[bytes, List[bytes]]]
    def lookup_prefix(&self, index_id: usize, prefix: PyBytes) -> PyResult<Vec<(PyBytes, Vec<PyBytes>)>> {
        let log = self.log(py).borrow();
        let prefix = prefix.data(py);
        let iter = log.lookup_prefix(index_id, prefix).map_pyerr(py)?;
        let result = range_iter_to_vec(py, iter).map_pyerr(py)?;
        Ok(result)
    }

    /// lookup_prefix_hex(index_id: int, hex_prefix: bytes) -> List[Tuple[bytes, List[bytes]]]
    def lookup_prefix_hex(&self, index_id: usize, hex_prefix: PyBytes) -> PyResult<Vec<(PyBytes, Vec<PyBytes>)>> {
        let log = self.log(py).borrow();
        let hex_prefix = hex_prefix.data(py);
        let iter = log.lookup_prefix_hex(index_id, hex_prefix).map_pyerr(py)?;
        let result = range_iter_to_vec(py, iter).map_pyerr(py)?;
        Ok(result)
    }

    /// lookup_range(index_id: int, start: bytes, end: bytes, start_inclusive: bool=True, end_inclusive: bool=False) -> List[Tuple[bytes, List[bytes]]]
    def lookup_range(&self, index_id: usize, start: Option<PyBytes> = None, end: Option<PyBytes> = None, start_inclusive: bool = true, end_inclusive: bool = false) -> PyResult<Vec<(PyBytes, Vec<PyBytes>)>> {
        let log = self.log(py).borrow();
        let range = (to_bound(py, start.as_ref(), start_inclusive), to_bound(py, end.as_ref(), end_inclusive));
        let iter = log.lookup_range(index_id, range).map_pyerr(py)?;
        let result = range_iter_to_vec(py, iter).map_pyerr(py)?;
        Ok(result)
    }

    /// Write pending changes to disk and pick up changes from disk.
    def sync(&self) -> PyResult<PyNone> {
        let mut log = self.log(py).borrow_mut();
        log.sync().map_pyerr(py)?;
        Ok(PyNone)
    }
});

fn to_bound<'a>(py: Python, key: Option<&'a PyBytes>, inclusive: bool) -> Bound<&'a [u8]> {
    match key {
        None => Bound::Unbounded,
        Some(key) => {
            let key = key.data(py);
            match inclusive {
                true => Bound::Included(key),
                false => Bound::Excluded(key),
            }
        }
    }
}

fn to_index_func(
    py: Python,
    func: PyObject,
) -> impl Fn(&[u8]) -> Vec<IndexOutput> + Send + Sync + 'static {
    let func = func.clone_ref(py);
    move |data: &[u8]| -> Vec<IndexOutput> {
        let gil_guard = Python::acquire_gil();
        let py = gil_guard.python();
        let data = PyBytes::new(py, data);
        let result = func.call(py, (data,), None).unwrap();
        let result: Vec<PyObject> = result.extract(py).unwrap();
        result
            .into_iter()
            .map(|obj| {
                if let Ok(obj) = obj.extract::<PyBytes>(py) {
                    let out: &[u8] = obj.data(py);
                    return IndexOutput::Owned(out.to_vec().into_boxed_slice());
                } else if let Ok((start, end)) = obj.extract::<(u64, u64)>(py) {
                    return IndexOutput::Reference(start..end);
                } else {
                    panic!("python index function returned unknown value")
                }
            })
            .collect()
    }
}

fn lookup_iter_to_vec(py: Python, iter: LogLookupIter) -> indexedlog::Result<Vec<PyBytes>> {
    let result = iter.collect::<Result<Vec<_>, _>>()?;
    let result = result
        .into_iter()
        .map(|item| PyBytes::new(py, item))
        .collect();
    Ok(result)
}

fn range_iter_to_vec(
    py: Python,
    iter: LogRangeIter,
) -> indexedlog::Result<Vec<(PyBytes, Vec<PyBytes>)>> {
    let result = iter.collect::<Result<Vec<_>, _>>()?;
    let result = result
        .into_iter()
        .map(
            |(full_key, iter)| -> indexedlog::Result<(PyBytes, Vec<PyBytes>)> {
                Ok((PyBytes::new(py, &full_key), lookup_iter_to_vec(py, iter)?))
            },
        )
        .collect::<Result<Vec<_>, _>>()?;
    Ok(result)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::cell::Cell;
use std::cell::RefCell;

use clidispatch::io::IO as RustIO;
use cpython::*;
use cpython_ext::wrap_rust_write;
use cpython_ext::PyNone;
use cpython_ext::ResultPyErrExt;
use pyconfigloader::config as PyConfig;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "io"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<IO>(py)?;
    m.add_class::<styler>(py)?;
    m.add(
        py,
        "shouldcolor",
        py_fn!(py, should_color(config: PyConfig)),
    )?;
    Ok(m)
}

py_class!(class IO |py| {
    data closed: Cell<bool>;

    @staticmethod
    def main() -> PyResult<IO> {
        Self::create_instance(py, Cell::new(false))
    }

    /// Start the stream pager.
    def start_pager(&self, config: PyConfig) -> PyResult<PyNone> {
        self.check_closed(py)?;
        let io = RustIO::main().map_pyerr(py)?;
        let config = &config.get_cfg(py);
        io.start_pager(config).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Test if the pager is active.
    def is_pager_active(&self) -> PyResult<bool> {
        let io = RustIO::main().map_pyerr(py)?;
        Ok(io.is_pager_active())
    }

    /// Write to pager's main buffer. Text should be in utf-8.
    def write(&self, bytes: PyBytes) -> PyResult<PyNone> {
        self.check_closed(py)?;
        let io = RustIO::main().map_pyerr(py)?;
        io.write(bytes.data(py)).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Write to pager's stderr buffer. Text should be in utf-8.
    def write_err(&self, bytes: PyBytes) -> PyResult<PyNone> {
        self.check_closed(py)?;
        let io = RustIO::main().map_pyerr(py)?;
        io.write_err(bytes.data(py)).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Set the progress content.
    def set_progress(&self, message: &str) -> PyResult<PyNone> {
        self.check_closed(py)?;
        let io = RustIO::main().map_pyerr(py)?;
        io.set_progress_str(message).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Wait for the pager to end.
    def wait_pager(&self) -> PyResult<PyNone> {
        self.closed(py).set(true);
        let io = RustIO::main().map_pyerr(py)?;
        io.wait_pager().map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Return the output stream as a Python object with "write" method.
    def output(&self) -> PyResult<PyObject> {
        let io = RustIO::main().map_pyerr(py)?;
        Ok(wrap_rust_write(py, io.output())?.into_object())
    }

    /// Return the error stream as a Python object with "write" method.
    def error(&self) -> PyResult<PyObject> {
        let io = RustIO::main().map_pyerr(py)?;
        Ok(wrap_rust_write(py, io.error())?.into_object())
    }

    /// Flush pending changes.
    def flush(&self) -> PyResult<PyNone> {
        let io = RustIO::main().map_pyerr(py)?;
        io.flush().map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Disable progress output.
    def disable_progress(&self, disabled: bool = true) -> PyResult<PyNone> {
        let io = RustIO::main().map_pyerr(py)?;
        io.disable_progress(disabled).map_pyerr(py)?;
        Ok(PyNone)
    }
});

impl IO {
    fn check_closed(&self, py: Python) -> PyResult<PyNone> {
        if self.closed(py).get() {
            Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "pager was closed",
            ))
            .map_pyerr(py)
        } else {
            Ok(PyNone)
        }
    }
}

py_class!(class styler |py| {
    data inner: RefCell<termstyle::Styler>;

    def __new__(_cls, colors: u32) -> PyResult<styler> {
        let level = if colors >= 16777216 {
            termstyle::ColorLevel::TrueColor
        } else if colors >= 256 {
            termstyle::ColorLevel::TwoFiftySix
        } else {
            termstyle::ColorLevel::Sixteen
        };

        Self::create_instance(py, RefCell::new(termstyle::Styler::from_level(level).map_pyerr(py)?))
    }

    def renderbytes(&self, style: &str, text: &str)  -> PyResult<PyBytes> {
        Ok(PyBytes::new(
            py,
            self.inner(py).borrow_mut().render_bytes(style, text)
                .map_pyerr(py)?
                .as_ref(),
        ))
    }
});

fn should_color(py: Python, config: PyConfig) -> PyResult<bool> {
    let config = &config.get_cfg(py);
    Ok(termstyle::should_color(
        config,
        &RustIO::main().map_pyerr(py)?.output(),
    ))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::cell::Cell;

use cpython::*;
use cpython_ext::PyNone;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "lock"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<pathlock>(py)?;
    Ok(m)
}

py_class!(class pathlock |py| {
    data lock: Cell<Option<repolock::LockHandle>>;

    @classmethod def trylock(_cls, dir: PyPathBuf, name: String, contents: String) -> PyResult<pathlock> {
        Self::create_instance(
            py,
            Cell::new(Some(
                repolock::try_lock(dir.as_path(), &name, contents.as_bytes()).map_pyerr(py)?,
            )),
        )
    }

    def unlock(&self) -> PyResult<PyNone> {
        if let Some(mut f) = self.lock(py).replace(None) {
            f.unlock().map_pyerr(py)?;
            Ok(PyNone)
        } else {
            Err(PyErr::new::<exc::ValueError, _>(py, "lock is already unlocked"))
        }
    }
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use cpython::*;
use cpython_ext::allocate_pybytes;
use cpython_ext::vec_to_pyobj;
use cpython_ext::ResultPyErrExt;
use cpython_ext::SimplePyBuf;
use lz4_pyframe::compress;
use lz4_pyframe::compresshc;
use lz4_pyframe::decompress_into;
use lz4_pyframe::decompress_size;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "lz4"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(py, "compress", py_fn!(py, compress_py(data: PyObject)))?;
    m.add(py, "compresshc", py_fn!(py, compresshc_py(data: PyObject)))?;
    m.add(py, "decompress", py_fn!(py, decompress_py(data: PyObject)))?;
    Ok(m)
}

fn compress_py(py: Python, data: PyObject) -> PyResult<PyObject> {
    let data = SimplePyBuf::new(py, &data);
    compress(data.as_ref())
        .map_pyerr(py)
        .map(|bytes| vec_to_pyobj(py, bytes))
}

fn compresshc_py(py: Python, data: PyObject) -> PyResult<PyObject> {
    let data = SimplePyBuf::new(py, &data);
    compresshc(data.as_ref())
        .map_pyerr(py)
        .map(|bytes| vec_to_pyobj(py, bytes))
}

fn decompress_py(py: Python, data: PyObject) -> PyResult<PyObject> {
    let data = SimplePyBuf::new(py, &data);
    let data = data.as_ref();
    let size = decompress_size(data).map_pyerr(py)?;
    let (obj, slice) = allocate_pybytes(py, size);
    decompress_into(data, slice).map_pyerr(py).map(move |_| obj)
}