/*
 * 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.
 */

//! A simple Git copytrace implementation to detect copies by calling git2 library.
//! NOTE: this is a temporary solution to unblock hg-over-git use case.

use std::collections::HashMap;
use std::path::Path;

use anyhow::Result;
use types::HgId;

pub struct GitCopyTrace {
    // A git repository instance. Check `GitCopyTrace::open` for more details.
    git_repo: git2::Repository,
}

impl GitCopyTrace {
    /// `open` a GitCopyTrace at `git_dir`.
    pub fn open(git_dir: &Path) -> Result<Self> {
        let git_repo = git2::Repository::open(git_dir)?;
        Ok(GitCopyTrace { git_repo })
    }

    /// Find copies/moves between old and new commits (HgId), the result is
    /// a map of new_path -> old_path
    pub fn find_copies(&self, old_id: HgId, new_id: HgId) -> Result<HashMap<String, String>> {
        let mut map: HashMap<_, _> = HashMap::new();

        let old_tree = self.find_tree(old_id)?;
        let new_tree = self.find_tree(new_id)?;

        let mut diff = self
            .git_repo
            .diff_tree_to_tree(Some(&old_tree), Some(&new_tree), None)?;
        diff.find_similar(None)?;

        diff.foreach(
            &mut |file, _progress| {
                if let git2::Delta::Renamed = file.status() {
                    let old_path = diff_file_to_string(file.old_file());
                    let new_path = diff_file_to_string(file.new_file());
                    // skip non-utf8 paths
                    if let (Some(old_path), Some(new_path)) = (old_path, new_path) {
                        map.insert(new_path.to_owned(), old_path.to_owned());
                    }
                }
                true
            },
            None,
            None,
            None,
        )?;

        Ok(map)
    }

    fn find_tree(&self, id: HgId) -> Result<git2::Tree> {
        let oid = git2::Oid::from_bytes(id.as_ref())?;
        let commit = self.git_repo.find_commit(oid)?;
        let tree = commit.tree()?;
        Ok(tree)
    }
}

fn diff_file_to_string(diff_file: git2::DiffFile) -> Option<&str> {
    let path = diff_file
        .path()
        .expect("expect a valid path for the diff file");
    path.to_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.
 */

mod copy_trace;
mod dag_copy_trace;
mod error;
mod git_copy_trace;
mod rename_finders;
mod utils;

pub use crate::copy_trace::CopyTrace;
pub use crate::dag_copy_trace::DagCopyTrace;
pub use crate::git_copy_trace::GitCopyTrace;
pub use crate::rename_finders::ContentSimilarityRenameFinder;
pub use crate::rename_finders::MetadataRenameFinder;
pub use crate::rename_finders::RenameFinder;

#[cfg(test)]
mod tests;

/// SearchDirection when searching renames.
///
/// Assuming we have a commit graph like below:
///
///  a..z # draw dag syntax
///
/// Forward means searching from a to z.
/// Backward means searching from z to a.
#[derive(Debug, PartialEq, Copy, Clone, Eq, Hash)]
pub(crate) enum SearchDirection {
    Forward,
    Backward,
}


/*
 * 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::cmp::min;
use std::collections::HashSet;
use std::iter::zip;
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use configmodel::Config;
use configmodel::ConfigExt;
use dag::Vertex;
use hg_metrics::increment_counter;
use lru_cache::LruCache;
use manifest::DiffType;
use manifest::Manifest;
use manifest_tree::Diff;
use manifest_tree::TreeManifest;
use parking_lot::Mutex;
use pathmatcher::AlwaysMatcher;
use storemodel::futures::StreamExt;
use storemodel::ReadFileContents;
use types::Key;
use types::RepoPath;
use types::RepoPathBuf;
use xdiff::edit_cost;

use crate::error::CopyTraceError;
use crate::utils::file_path_similarity;
use crate::SearchDirection;

/// Maximum rename candidate files to check
const DEFAULT_MAX_RENAME_CANDIDATES: usize = 10;
/// Content similarity threhold for rename detection. The definition of "similarity"
/// between file a and file b is: (len(a.lines()) - edit_cost(a, b)) / len(a.lines())
///   * 1.0 means exact match
///   * 0.0 means not match at all
/// The default value is 0.8.
const DEFAULT_SIMILARITY_THRESHOLD: f32 = 0.8;
/// Maximum rename edit cost determines whether we treat two files as a rename
const DEFAULT_MAX_EDIT_COST: u64 = 1000;
/// Control if MetadataRenameFinder fallbacks to content similarity finder
const DEFAULT_FALLBACK_TO_CONTENT_SIMILARITY: bool = false;
/// Default Rename cache size
const DEFAULT_RENAME_CACHE_SIZE: usize = 1000;

/// Finding rename between old and new trees (commits).
/// old_tree is a parent of new_tree
#[async_trait]
pub trait RenameFinder {
    /// Find the new path of the given old path in the new_tree
    async fn find_rename_forward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        old_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>>;

    /// Find the old path of the given new path in the old_tree
    async fn find_rename_backward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        new_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>>;
}

/// Rename finder based on the copy information in the file header metadata
pub struct MetadataRenameFinder {
    inner: RenameFinderInner,
}

/// Content similarity based Rename finder (mainly for Git repo)
pub struct ContentSimilarityRenameFinder {
    inner: RenameFinderInner,
}

/// RenameFinderInner is the base struct for MetadataRenameFinder and ContentSimilarityRenameFinder,
/// It is introduced for code reuse between those two file based rename finders.
struct RenameFinderInner {
    // Read content and rename metadata of a file
    file_reader: Arc<dyn ReadFileContents<Error = anyhow::Error> + Send + Sync>,
    // Read configs
    config: Arc<dyn Config + Send + Sync>,
    // Dir move caused rename candidates
    cache: Mutex<LruCache<CacheKey, Key>>,
}

type CacheKey = (Vertex, RepoPathBuf);

impl MetadataRenameFinder {
    pub fn new(
        file_reader: Arc<dyn ReadFileContents<Error = anyhow::Error> + Send + Sync>,
        config: Arc<dyn Config + Send + Sync>,
    ) -> Result<Self> {
        let cache_size = get_rename_cache_size(&config)?;
        let inner = RenameFinderInner {
            file_reader,
            config,
            cache: Mutex::new(LruCache::new(cache_size)),
        };
        Ok(Self { inner })
    }
}

#[async_trait]
impl RenameFinder for MetadataRenameFinder {
    async fn find_rename_forward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        old_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>> {
        let candidates = self.inner.generate_candidates(
            old_tree,
            new_tree,
            old_path,
            new_vertex,
            SearchDirection::Forward,
        )?;

        let found = self
            .inner
            .read_renamed_metadata_forward(candidates.clone(), old_path)
            .await?;

        if found.is_some() || !self.inner.get_fallback_to_content_similarity()? {
            return Ok(found);
        }

        // fallback to content similarity
        let old_path_key = self.inner.get_key_from_path(old_tree, old_path)?;
        let found = self
            .inner
            .find_similar_file(candidates, old_path_key)
            .await?;
        emit_content_similarity_fallback_metric(found.is_some());
        Ok(found)
    }

    async fn find_rename_backward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        new_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>> {
        let new_key = self.inner.get_key_from_path(new_tree, new_path)?;
        let found = self
            .inner
            .read_renamed_metadata_backward(new_key.clone())
            .await?;

        if found.is_some() || !self.inner.get_fallback_to_content_similarity()? {
            return Ok(found);
        }

        // fallback to content similarity
        let candidates = self.inner.generate_candidates(
            old_tree,
            new_tree,
            new_path,
            new_vertex,
            SearchDirection::Backward,
        )?;
        let found = self.inner.find_similar_file(candidates, new_key).await?;
        emit_content_similarity_fallback_metric(found.is_some());
        Ok(found)
    }
}

impl ContentSimilarityRenameFinder {
    pub fn new(
        file_reader: Arc<dyn ReadFileContents<Error = anyhow::Error> + Send + Sync>,
        config: Arc<dyn Config + Send + Sync>,
    ) -> Result<Self> {
        let cache_size = get_rename_cache_size(&config)?;
        let inner = RenameFinderInner {
            file_reader,
            config,
            cache: Mutex::new(LruCache::new(cache_size)),
        };
        Ok(Self { inner })
    }
}

#[async_trait]
impl RenameFinder for ContentSimilarityRenameFinder {
    async fn find_rename_forward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        old_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>> {
        self.inner
            .find_rename_in_direction(
                old_tree,
                new_tree,
                old_path,
                new_vertex,
                SearchDirection::Forward,
            )
            .await
    }

    async fn find_rename_backward(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        new_path: &RepoPath,
        new_vertex: &Vertex,
    ) -> Result<Option<RepoPathBuf>> {
        self.inner
            .find_rename_in_direction(
                old_tree,
                new_tree,
                new_path,
                new_vertex,
                SearchDirection::Backward,
            )
            .await
    }
}

impl RenameFinderInner {
    fn generate_candidates(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        path: &RepoPath,
        vertex: &Vertex,
        direction: SearchDirection,
    ) -> Result<Vec<Key>> {
        let cache_key = (vertex.clone(), path.to_owned());
        {
            let mut cache = self.cache.lock();
            if let Some(val) = cache.get_mut(&cache_key) {
                return Ok(vec![val.clone()]);
            }
        }

        let (mut added_files, mut deleted_files) =
            self.get_added_and_deleted_files(old_tree, new_tree)?;
        let batch_mv_candidates = detect_batch_move(&mut added_files, &mut deleted_files);

        if batch_mv_candidates.is_empty() {
            match direction {
                SearchDirection::Forward => {
                    select_rename_candidates(added_files, path, &self.config)
                }
                SearchDirection::Backward => {
                    select_rename_candidates(deleted_files, path, &self.config)
                }
            }
        } else {
            let mut candidates: Vec<Key> = vec![];
            let mut cache = self.cache.lock();

            match direction {
                SearchDirection::Forward => {
                    for (to, from) in batch_mv_candidates {
                        let key = (vertex.clone(), from.path);
                        cache.insert(key, to);
                    }
                }
                SearchDirection::Backward => {
                    for (to, from) in batch_mv_candidates {
                        let key = (vertex.clone(), to.path);
                        cache.insert(key, from);
                    }
                }
            }

            if let Some(val) = cache.get_mut(&cache_key) {
                candidates.push(val.to_owned());
            }
            Ok(candidates)
        }
    }

    fn get_added_and_deleted_files(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
    ) -> Result<(Vec<Key>, Vec<Key>)> {
        let mut added_files = Vec::new();
        let mut deleted_files = Vec::new();
        let matcher = AlwaysMatcher::new();
        let diff = Diff::new(old_tree, new_tree, &matcher)?;
        for entry in diff {
            let entry = entry?;
            match entry.diff_type {
                DiffType::RightOnly(file_metadata) => {
                    let path = entry.path;
                    let key = Key {
                        path,
                        hgid: file_metadata.hgid,
                    };
                    added_files.push(key);
                }
                DiffType::LeftOnly(file_metadata) => {
                    let path = entry.path;
                    let key = Key {
                        path,
                        hgid: file_metadata.hgid,
                    };
                    deleted_files.push(key);
                }
                _ => {}
            }
        }
        Ok((added_files, deleted_files))
    }

    async fn read_renamed_metadata_forward(
        &self,
        keys: Vec<Key>,
        old_path: &RepoPath,
    ) -> Result<Option<RepoPathBuf>> {
        tracing::trace!(keys_len = keys.len(), " read_renamed_metadata_forward");
        let mut renames = self.file_reader.read_rename_metadata(keys).await;
        while let Some(rename) = renames.next().await {
            let (key, rename_from_key) = rename?;
            if let Some(rename_from_key) = rename_from_key {
                if rename_from_key.path.as_repo_path() == old_path {
                    return Ok(Some(key.path));
                }
            }
        }
        Ok(None)
    }

    async fn read_renamed_metadata_backward(&self, key: Key) -> Result<Option<RepoPathBuf>> {
        let mut renames = self.file_reader.read_rename_metadata(vec![key]).await;
        if let Some(rename) = renames.next().await {
            let (_, rename_from_key) = rename?;
            return Ok(rename_from_key.map(|k| k.path));
        }
        Ok(None)
    }

    async fn find_rename_in_direction(
        &self,
        old_tree: &TreeManifest,
        new_tree: &TreeManifest,
        source_path: &RepoPath,
        new_vertex: &Vertex,
        direction: SearchDirection,
    ) -> Result<Option<RepoPathBuf>> {
        tracing::trace!(?source_path, ?direction, " find_rename_in_direction");

        let candidates =
            self.generate_candidates(old_tree, new_tree, source_path, new_vertex, direction)?;

        tracing::trace!(candidates_len = candidates.len(), " found");

        let source_tree = match direction {
            SearchDirection::Forward => old_tree,
            SearchDirection::Backward => new_tree,
        };
        let source = self.get_key_from_path(source_tree, source_path)?;

        self.find_similar_file(candidates, source).await
    }

    async fn find_similar_file(
        &self,
        keys: Vec<Key>,
        source_key: Key,
    ) -> Result<Option<RepoPathBuf>> {
        let mut source = self
            .file_reader
            .read_file_contents(vec![source_key.clone()])
            .await;
        let source_content = match source.next().await {
            None => return Err(CopyTraceError::FileNotFound(source_key.path).into()),
            Some(content_and_key) => content_and_key?.0,
        };

        let config_percentage = self.get_similarity_threshold()?;
        let config_max_edit_cost = self.get_max_edit_cost()?;
        let lines = source_content.iter().filter(|&&c| c == b'\n').count();
        let max_edit_cost = min(
            (lines as f32 * (1.0 - config_percentage)).round() as u64,
            config_max_edit_cost,
        );
        tracing::trace!(
            ?config_percentage,
            ?config_max_edit_cost,
            ?lines,
            ?max_edit_cost,
            " content similarity configs"
        );

        let mut candidates = self.file_reader.read_file_contents(keys).await;
        while let Some(candidate) = candidates.next().await {
            let (candidate_content, k) = candidate?;
            if edit_cost(&source_content, &candidate_content, max_edit_cost + 1) <= max_edit_cost {
                return Ok(Some(k.path));
            }
        }

        Ok(None)
    }

    fn get_key_from_path(&self, tree: &TreeManifest, path: &RepoPath) -> Result<Key> {
        let key = match tree.get_file(path)? {
            None => return Err(CopyTraceError::FileNotFound(path.to_owned()).into()),
            Some(file_metadata) => Key {
                path: path.to_owned(),
                hgid: file_metadata.hgid,
            },
        };
        Ok(key)
    }

    fn get_similarity_threshold(&self) -> Result<f32> {
        let v = self
            .config
            .get_opt::<f32>("copytrace", "similarity-threshold")?
            .unwrap_or(DEFAULT_SIMILARITY_THRESHOLD);
        Ok(v)
    }

    fn get_max_edit_cost(&self) -> Result<u64> {
        let v = self
            .config
            .get_opt::<u64>("copytrace", "max-edit-cost")?
            .unwrap_or(DEFAULT_MAX_EDIT_COST);
        Ok(v)
    }

    fn get_fallback_to_content_similarity(&self) -> Result<bool> {
        let v = self
            .config
            .get_opt::<bool>("copytrace", "fallback-to-content-similarity")?
            .unwrap_or(DEFAULT_FALLBACK_TO_CONTENT_SIMILARITY);
        Ok(v)
    }
}

fn select_rename_candidates(
    mut candidates: Vec<Key>,
    source_path: &RepoPath,
    config: &dyn Config,
) -> Result<Vec<Key>> {
    // It's rare that a file will be copied and renamed (multiple copies) in one commit.
    // We don't plan to support this one-to-many mapping since it will make copytrace
    // complexity increase exponentially. Here, we order the potential new files in
    // path similarity order (most similar one first), and return the first one that
    // is a copy of the old_path.
    candidates.sort_by_key(|k| {
        let path = k.path.as_repo_path();
        let score = file_path_similarity(path, source_path);
        (-score, path.to_owned())
    });
    let max_rename_candidates = config
        .get_opt::<usize>("copytrace", "max-rename-candidates")?
        .unwrap_or(DEFAULT_MAX_RENAME_CANDIDATES);
    if candidates.len() > max_rename_candidates {
        Ok(candidates.into_iter().take(max_rename_candidates).collect())
    } else {
        Ok(candidates)
    }
}

fn detect_batch_move(added_files: &mut Vec<Key>, deleted_files: &mut Vec<Key>) -> Vec<(Key, Key)> {
    if added_files.len() != deleted_files.len() {
        return vec![];
    }

    added_files.sort_unstable_by(|a, b| a.path.cmp(&b.path));
    deleted_files.sort_unstable_by(|a, b| a.path.cmp(&b.path));

    let stripped_pairs: HashSet<(String, String)> = zip(added_files.clone(), deleted_files.clone())
        .map(|(a, b)| strip_common_prefix_and_suffix(a.path.as_str(), b.path.as_str()))
        .collect();

    if stripped_pairs.len() > 1 {
        vec![]
    } else {
        zip(added_files, deleted_files)
            .map(|(a, b)| (a.clone(), b.clone()))
            .collect()
    }
}

fn strip_common_prefix_and_suffix(s1: &str, s2: &str) -> (String, String) {
    let s1_chars: Vec<char> = s1.chars().collect();
    let s2_chars: Vec<char> = s2.chars().collect();

    let mut start = 0;
    while start < s1_chars.len() && start < s2_chars.len() && s1_chars[start] == s2_chars[start] {
        start += 1;
    }

    let mut end = 0;
    while end < s1_chars.len() - start
        && end < s2_chars.len() - start
        && s1_chars[s1_chars.len() - 1 - end] == s2_chars[s2_chars.len() - 1 - end]
    {
        end += 1;
    }

    let stripped_s1: String = s1_chars[start..s1_chars.len() - end].iter().collect();
    let stripped_s2: String = s2_chars[start..s2_chars.len() - end].iter().collect();
    (stripped_s1, stripped_s2)
}

fn get_rename_cache_size(config: &dyn Config) -> Result<usize> {
    let v = config
        .get_opt::<usize>("copytrace", "rename-cache-size")?
        .unwrap_or(DEFAULT_RENAME_CACHE_SIZE);
    Ok(v)
}

fn emit_content_similarity_fallback_metric(is_found: bool) {
    let metric = if is_found {
        "copytrace_content_similarity_fallback_success"
    } else {
        "copytrace_content_similarity_fallback_failure"
    };
    increment_counter(metric, 1);
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use types::HgId;

    use super::*;

    #[test]
    fn test_select_rename_candidates() {
        let candidates: Vec<Key> = vec![
            gen_key("a/b/c.txt"),
            gen_key("a/b/c.md"),
            gen_key("a/d.txt"),
            gen_key("e.txt"),
        ];
        let source_path = &RepoPath::from_str("a/c.txt").unwrap();
        let mut config: BTreeMap<&'static str, &'static str> = Default::default();
        config.insert("copytrace.max-rename-candidates", "2");
        let config = Arc::new(config);

        let actual = select_rename_candidates(candidates, source_path, &config).unwrap();

        let expected = vec![gen_key("a/b/c.txt"), gen_key("a/d.txt")];
        assert_eq!(actual, expected)
    }

    fn gen_key(path: &str) -> Key {
        let path = RepoPath::from_str(path).unwrap().to_owned();
        let hgid = HgId::null_id().clone();
        Key { path, hgid }
    }

    #[test]
    fn test_detect_batch_move() {
        let mut added_files: Vec<Key> = vec![
            gen_key("a/b/c.txt"),
            gen_key("a/b/c.md"),
            gen_key("a/d.txt"),
        ];

        let mut deleted_files: Vec<Key> = vec![
            gen_key("b/b/c.txt"),
            gen_key("b/b/c.md"),
            gen_key("b/d.txt"),
        ];

        assert_eq!(
            detect_batch_move(&mut added_files, &mut deleted_files),
            vec![
                (gen_key("a/b/c.md"), gen_key("b/b/c.md")),
                (gen_key("a/b/c.txt"), gen_key("b/b/c.txt")),
                (gen_key("a/d.txt"), gen_key("b/d.txt")),
            ]
        );
    }

    #[test]
    fn test_detect_batch_move_with_unequal_num_of_files() {
        let mut added_files: Vec<Key> = vec![
            gen_key("a/b/c.txt"),
            gen_key("a/b/c.md"),
            gen_key("a/d.txt"),
            gen_key("a/e.txt"),
        ];

        let mut deleted_files: Vec<Key> = vec![
            gen_key("b/b/c.txt"),
            gen_key("b/b/c.md"),
            gen_key("b/d.txt"),
        ];

        assert_eq!(
            detect_batch_move(&mut added_files, &mut deleted_files),
            vec![]
        );
    }

    #[test]
    fn test_detect_batch_move_with_unmatched_basename() {
        let mut added_files: Vec<Key> = vec![
            gen_key("a/b/c.txt"),
            gen_key("a/b/c.md"),
            gen_key("a/d.txt"),
        ];

        let mut deleted_files: Vec<Key> = vec![
            gen_key("b/b/ccc.txt"),
            gen_key("b/b/c.md"),
            gen_key("b/d.txt"),
        ];

        assert_eq!(
            detect_batch_move(&mut added_files, &mut deleted_files),
            vec![]
        );
    }

    #[test]
    fn test_strip_common_prefix_and_suffix() {
        let func = strip_common_prefix_and_suffix;

        assert_eq!(func("", ""), ("".to_owned(), "".to_owned()));
        assert_eq!(func("", "a"), ("".to_owned(), "a".to_owned()));
        assert_eq!(func("a", ""), ("a".to_owned(), "".to_owned()));

        assert_eq!(func("1a2", "1b2"), ("a".to_owned(), "b".to_owned()));
        assert_eq!(func("1a22", "1b2"), ("a2".to_owned(), "b".to_owned()));

        assert_eq!(
            func("/a/1.txt", "/a/1.md"),
            ("txt".to_owned(), "md".to_owned())
        );
        assert_eq!(
            func("/a/b/1.txt", "/a/c/1.txt"),
            ("b".to_owned(), "c".to_owned())
        );

        assert_eq!(
            func("/æ–‡ä»¶/æˆ‘çš„/å¥½.txt", "/æ–‡ä»¶/ä½ çš„/å¥½.txt"),
            ("æˆ‘".to_owned(), "ä½ ".to_owned()),
        );
    }
}


/*
 * 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.
 */

//! Testing.

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

use anyhow::Result;
use async_trait::async_trait;
use dag::ops::ImportAscii;
use dag::DagAlgorithm;
use dag::MemDag;
use dag::Vertex;
use manifest::FileMetadata;
use manifest::FileType;
use manifest::Manifest;
use manifest_tree::testutil::TestStore;
use manifest_tree::TreeManifest;
use manifest_tree::TreeStore;
use storemodel::futures::stream;
use storemodel::futures::StreamExt;
use storemodel::ReadFileContents;
use storemodel::ReadRootTreeIds;
use storemodel::TreeFormat;
use tracing_test::traced_test;
use types::HgId;
use types::Key;
use types::RepoPath;
use types::RepoPathBuf;

use crate::CopyTrace;
use crate::DagCopyTrace;
use crate::MetadataRenameFinder;

#[derive(Clone)]
struct CopyTraceTestCase {
    inner: Arc<CopyTraceTestCaseInner>,
}

struct CopyTraceTestCaseInner {
    /// Commits that change trees.
    commit_to_tree: HashMap<Vertex, HgId>,
    /// In memory tree store
    tree_store: Arc<dyn TreeStore + Send + Sync>,
    /// Dag algorithm
    dagalgo: Arc<dyn DagAlgorithm + Send + Sync>,
    /// Copies info: dest -> src mapping
    copies: HashMap<Key, Key>,
    /// Config
    config: BTreeMap<&'static str, &'static str>,
}

#[derive(Debug)]
enum Change {
    Add(RepoPathBuf, FileMetadata),
    Delete(RepoPathBuf),
    Rename(RepoPathBuf, RepoPathBuf),
    Modify(RepoPathBuf, FileMetadata),
    Copy(RepoPathBuf, RepoPathBuf),
}

impl CopyTraceTestCase {
    pub async fn new(text: &str, changes: HashMap<&str, Vec<&str>>) -> Self {
        let mut mem_dag = MemDag::new();
        mem_dag
            .import_ascii_with_vertex_fn(text, vertex_from_str)
            .unwrap();

        let mut commit_to_tree: HashMap<Vertex, HgId> = Default::default();
        let mut copies: HashMap<Key, Key> = Default::default();
        let tree_store = Arc::new(TestStore::new().with_format(TreeFormat::Git));
        let changes = Change::build_changes(changes);
        let config: BTreeMap<&'static str, &'static str> = Default::default();

        // iterate through the commit graph to build trees in topo order (ascending)
        let set = mem_dag.all().await.unwrap();
        let mut iter = set.iter_rev().await.unwrap();
        while let Some(item) = iter.next().await {
            Self::build_tree(
                &mem_dag,
                tree_store.clone(),
                &mut commit_to_tree,
                &mut copies,
                item.unwrap(),
                &changes,
            )
            .await;
        }

        let inner = Arc::new(CopyTraceTestCaseInner {
            commit_to_tree,
            tree_store,
            dagalgo: Arc::new(mem_dag),
            copies,
            config,
        });
        CopyTraceTestCase { inner }
    }

    /// Create a DagCopyTrace instance
    pub async fn copy_trace(&self) -> Arc<dyn CopyTrace + Send + Sync> {
        let file_reader = Arc::new(self.clone());
        let config = Arc::new(self.inner.config.clone());
        let rename_finder = Arc::new(MetadataRenameFinder::new(file_reader, config).unwrap());

        let root_tree_reader = Arc::new(self.clone());
        let tree_store = self.inner.tree_store.clone();
        let dagalgo = self.inner.dagalgo.clone();

        let copy_trace =
            DagCopyTrace::new(root_tree_reader, tree_store, rename_finder, dagalgo).unwrap();
        Arc::new(copy_trace)
    }

    async fn build_tree(
        dag: &MemDag,
        tree_store: Arc<dyn TreeStore + Send + Sync>,
        commit_to_tree: &mut HashMap<Vertex, HgId>,
        copies: &mut HashMap<Key, Key>,
        commit: Vertex,
        changes: &HashMap<Vertex, Vec<Change>>,
    ) {
        let parents = dag.parent_names(commit.clone()).await.unwrap();
        let mut tree = match parents.first() {
            None => TreeManifest::ephemeral(tree_store.clone()),
            Some(p1) => {
                let tree_id = commit_to_tree[p1].clone();
                TreeManifest::durable(tree_store.clone(), tree_id)
            }
        };
        if let Some(val) = changes.get(&commit) {
            for v in val {
                match v {
                    Change::Add(path, metadata) => tree.insert(path.clone(), *metadata).unwrap(),
                    Change::Delete(path) => {
                        tree.remove(path).unwrap();
                    }
                    Change::Modify(path, metadata) => {
                        // `path` should exist
                        tree.get_file(path).unwrap().unwrap();
                        tree.insert(path.clone(), *metadata).unwrap();
                    }
                    Change::Rename(from_path, to_path) | Change::Copy(from_path, to_path) => {
                        let file_metadata = tree.get_file(from_path).unwrap().unwrap();
                        if let Change::Rename(_, _) = v {
                            tree.remove(from_path).unwrap();
                        }
                        tree.insert(to_path.clone(), file_metadata).unwrap();

                        // update copies mapping
                        let to_key = Key::new(to_path.clone(), file_metadata.hgid);
                        let from_key = Key::new(from_path.clone(), file_metadata.hgid);
                        copies.insert(to_key, from_key);
                    }
                }
            }
        }

        let tree_id = tree.flush().unwrap();
        commit_to_tree.insert(commit, tree_id);
    }
}

#[async_trait]
impl ReadRootTreeIds for CopyTraceTestCase {
    async fn read_root_tree_ids(&self, commits: Vec<HgId>) -> Result<Vec<(HgId, HgId)>> {
        let result = commits
            .into_iter()
            .map(|commit_id| {
                let commit = Vertex::copy_from(commit_id.as_ref());
                let tree_id = self.inner.commit_to_tree.get(&commit).unwrap().clone();
                (commit_id, tree_id)
            })
            .collect();
        Ok(result)
    }
}

#[async_trait]
impl ReadFileContents for CopyTraceTestCase {
    type Error = anyhow::Error;

    #[allow(dead_code)]
    async fn read_file_contents(
        &self,
        _keys: Vec<Key>,
    ) -> stream::BoxStream<Result<(storemodel::minibytes::Bytes, Key), Self::Error>> {
        // We will need this for computing content similarity score later.
        todo!()
    }

    async fn read_rename_metadata(
        &self,
        keys: Vec<Key>,
    ) -> stream::BoxStream<Result<(Key, Option<Key>), Self::Error>> {
        let renames: Vec<_> = {
            keys.iter()
                .map(|k| Ok((k.clone(), self.inner.copies.get(k).cloned())))
                .collect()
        };
        stream::iter(renames).boxed()
    }
}

impl Change {
    pub(crate) fn from_str(s: &str) -> Self {
        let items: Vec<_> = s.split(' ').collect();
        let change = match items[0] {
            "+" => Change::Add(to_repo_path_buf(items[1]), to_file_metadata(items[2])),
            "-" => Change::Delete(to_repo_path_buf(items[1])),
            "M" => Change::Modify(to_repo_path_buf(items[1]), to_file_metadata(items[2])),
            "->" => Change::Rename(to_repo_path_buf(items[1]), to_repo_path_buf(items[2])),
            "C" => Change::Copy(to_repo_path_buf(items[1]), to_repo_path_buf(items[2])),
            _ => unreachable!("unexpected token {}", items[0]),
        };
        return change;

        fn to_repo_path_buf(s: &str) -> RepoPathBuf {
            RepoPath::from_str(s).unwrap().to_owned()
        }

        fn to_file_metadata(s: &str) -> FileMetadata {
            FileMetadata::new(hgid_from_str(s), FileType::Regular)
        }
    }

    pub(crate) fn build_changes(changes: HashMap<&str, Vec<&str>>) -> HashMap<Vertex, Vec<Change>> {
        changes
            .into_iter()
            .map(|(k, vs)| {
                (
                    vertex_from_str(k),
                    vs.iter().map(|v| Change::from_str(v)).collect(),
                )
            })
            .collect()
    }
}

fn hgid_from_str(s: &str) -> HgId {
    let mut bytes = s.as_bytes().to_vec();
    bytes.resize(HgId::len(), 0);
    HgId::from_slice(&bytes).unwrap()
}

fn vertex_from_str(s: &str) -> Vertex {
    let mut bytes = s.as_bytes().to_vec();
    bytes.resize(HgId::len(), 0);
    Vertex::copy_from(&bytes)
}

macro_rules! assert_trace_rename {
    ($copy_trace:ident $src:tt $dst:tt, $src_path:tt -> $dst_path:tt) => {{
        let src = vertex_from_str(stringify!($src));
        let dst = vertex_from_str(stringify!($dst));
        let src_path = RepoPath::from_str(stringify!($src_path).trim_matches('"'))
            .unwrap()
            .to_owned();
        let dst_path = match stringify!($dst_path).trim_matches('"') {
            "!" => None,
            p => Some(RepoPath::from_str(p).unwrap().to_owned()),
        };

        let res = $copy_trace.trace_rename(src, dst, src_path).await.unwrap();

        assert_eq!(res, dst_path);
    }};
}

#[tokio::test]
async fn test_linear_single_rename() {
    let ascii = r#"
    C
    |
    B
    |
    A
    "#;
    let changes = HashMap::from([
        ("A", vec!["+ a 1", "+ c 2"]),
        ("B", vec!["-> a b"]),
        ("C", vec!["M b 3"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A C, a -> b);
    assert_trace_rename!(c C A, b -> a);

    assert_trace_rename!(c A B, a -> b);
    assert_trace_rename!(c B A, b -> a);

    assert_trace_rename!(c A C, d -> !);
}

#[tokio::test]
#[traced_test]
async fn test_non_linear_single_rename() {
    let ascii = r#"
    B C
    |/
    A
    |
    S
    "#;
    let changes = HashMap::from([
        ("S", vec!["+ s 0"]),
        ("A", vec!["+ a 1", "+ c 2"]),
        ("B", vec!["-> a b"]),
        ("C", vec!["M a 3"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c C B, a -> b);
    assert_trace_rename!(c B C, b -> a);
}

#[tokio::test]
#[traced_test]
async fn test_linear_multiple_renames() {
    let ascii = r#"
    Z
    :
    A
    "#;
    let changes = HashMap::from([
        ("A", vec!["+ a 1"]),
        ("B", vec!["-> a b"]),
        ("H", vec!["-> b c"]),
        ("X", vec!["-> c d"]),
        ("Z", vec!["- d"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A B, a -> b);
    assert_trace_rename!(c A H, a -> c);
    assert_trace_rename!(c A K, a -> c);
    assert_trace_rename!(c A X, a -> d);

    assert_trace_rename!(c X H, d -> c);
    assert_trace_rename!(c X C, d -> b);
    assert_trace_rename!(c X B, d -> b);
    assert_trace_rename!(c X A, d -> a);
}

#[tokio::test]
async fn test_linear_multiple_renames_with_deletes() {
    let ascii = r#"
    Z
    :
    A
    "#;
    let changes = HashMap::from([
        ("A", vec!["+ a 1"]),
        ("B", vec!["-> a b", "+ b2 2"]),
        ("H", vec!["-> b c"]),
        ("X", vec!["-> c d"]),
        ("Z", vec!["- d"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A X, a -> d);
    assert_trace_rename!(c A Z, a -> !);

    assert_trace_rename!(c Z B, b2 -> b2);
    assert_trace_rename!(c Z A, b2 -> !);
}

#[tokio::test]
async fn test_non_linear_multiple_renames() {
    let ascii = r#"
    1..10..1023
        \
         A..Z
    "#;
    let changes = HashMap::from([
        ("1", vec!["+ a 1"]),
        ("100", vec!["-> a b"]),
        ("500", vec!["-> b c"]),
        ("1000", vec!["-> c d"]),
        ("C", vec!["M a 11"]),
        ("D", vec!["-> a a2"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c Z 99, a2 -> a);
    assert_trace_rename!(c Z 100, a2 -> b);
    assert_trace_rename!(c Z 101, a2 -> b);
    assert_trace_rename!(c Z 500, a2 -> c);
    assert_trace_rename!(c Z 999, a2 -> c);
    assert_trace_rename!(c Z 1000, a2 -> d);
    assert_trace_rename!(c Z 1001, a2 -> d);
    assert_trace_rename!(c Z 1023, a2 -> d);

    assert_trace_rename!(c C 999, a -> c);
    assert_trace_rename!(c C 1023, a -> d);

    assert_trace_rename!(c 1023 B, d -> a);
    assert_trace_rename!(c 1023 C, d -> a);
    assert_trace_rename!(c 1023 D, d -> a2);
    assert_trace_rename!(c 1023 E, d -> a2);
    assert_trace_rename!(c 1023 Z, d -> a2);
}

#[tokio::test]
async fn test_non_linear_multiple_renames_with_deletes() {
    let ascii = r#"
    1..10..1023
        \
         A..Z
    "#;
    let changes = HashMap::from([
        ("1", vec!["+ a 1"]),
        ("100", vec!["-> a b"]),
        ("500", vec!["-> b c"]),
        ("1000", vec!["-> c d"]),
        ("1001", vec!["- d"]),
        ("C", vec!["M a 11"]),
        ("D", vec!["-> a a2"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c Z 1000, a2 -> d);
    assert_trace_rename!(c Z 1001, a2 -> !);
    assert_trace_rename!(c Z 1023, a2 -> !);
}

#[tokio::test]
async fn test_multiple_copies_ordering_default() {
    let ascii = r#"
    C
    :
    A
    "#;
    let changes = HashMap::from([("A", vec!["+ a 1"]), ("B", vec!["C a c", "-> a b"])]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A C, a -> b);
}

#[tokio::test]
async fn test_multiple_copies_ordering_same_basename_win() {
    let ascii = r#"
    C
    :
    A
    "#;
    let changes = HashMap::from([
        ("A", vec!["+ a 1"]),
        ("B", vec!["C a x/b", "C a z/a", "-> a b"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A C, a -> "z/a");
}

#[tokio::test]
async fn test_multiple_copies_ordering_same_directory_win() {
    let ascii = r#"
    C
    :
    A
    "#;

    let changes = HashMap::from([("A", vec!["+ a 1"]), ("B", vec!["C a x/b", "-> a b"])]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A C, a -> b);
}

#[tokio::test]
#[traced_test]
async fn test_linear_dir_move() {
    let ascii = r#"
    C
    |
    B
    |
    A
    "#;
    let changes = HashMap::from([
        ("A", vec!["+ a/1.txt 1", "+ a/2.md 2", "+ a/b/3.c 3"]),
        (
            "B",
            vec![
                "-> a/1.txt b/1.txt",
                "-> a/2.md b/2.md",
                "-> a/b/3.c b/b/3.c",
            ],
        ),
        ("C", vec!["M b/1.txt 4"]),
    ]);
    let t = CopyTraceTestCase::new(ascii, changes).await;
    let c = t.copy_trace().await;

    assert_trace_rename!(c A C, "a/1.txt" -> "b/1.txt");
    assert_trace_rename!(c A C, "a/2.md" -> "b/2.md");
    assert_trace_rename!(c C A, "b/1.txt" -> "a/1.txt");
    assert_trace_rename!(c C A, "b/2.md" -> "a/2.md");

    assert_trace_rename!(c A B, "a/b/3.c" -> "b/b/3.c");
    assert_trace_rename!(c B A, "b/b/3.c" -> "a/b/3.c");
}


/*
 * 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 types::RepoPath;

/// Contants for path similarity score. The actually number does not matter
/// here, we are more care about the ordering. Based on Git community's data,
/// ~70%+ of renames have the same basename:
/// https://github.com/git/git/blob/74cc1aa55f30ed76424a0e7226ab519aa6265061/diffcore-rename.c#L904-L907
const PATH_SAME_BASENAME_SCORE: i8 = 70;
const PATH_SAME_DIRECTORY_SCORE: i8 = 20;
const PATH_DEFAULT_SCORE: i8 = 1;

/// Computes the similarity score between file paths.
///
/// The score will be in range [0, 100]. Higher number means more similar.
/// The algorithm here is based heuristics, we're assuming the renames are
/// of following two types in most cases:
///   - moving from one directory to antoher (same basename)
///   - moving inside a directory (same directory name)
#[allow(dead_code)]
pub(crate) fn file_path_similarity(p1: &RepoPath, p2: &RepoPath) -> i8 {
    let (dir1, basename1) = match p1.split_last_component() {
        None => return 0,
        Some(val) => val,
    };
    let (dir2, basename2) = match p2.split_last_component() {
        None => return 0,
        Some(val) => val,
    };

    if basename1 == basename2 {
        return PATH_SAME_BASENAME_SCORE;
    }
    if dir1 == dir2 {
        return PATH_SAME_DIRECTORY_SCORE;
    }
    PATH_DEFAULT_SCORE
}

#[cfg(test)]
mod tests {
    use super::*;

    macro_rules! assert_file_path_similarity {
        ($p1:expr, $p2:expr, $expected:expr) => {{
            let p1 = RepoPath::from_str($p1).unwrap();
            let p2 = RepoPath::from_str($p2).unwrap();
            let actual = file_path_similarity(p1, p2);
            assert_eq!($expected, actual);
        }};
    }

    #[test]
    fn test_paths_with_same_basename() {
        assert_file_path_similarity!("a/1.txt", "1.txt", PATH_SAME_BASENAME_SCORE);
        assert_file_path_similarity!("a/1.txt", "b/1.txt", PATH_SAME_BASENAME_SCORE);
        assert_file_path_similarity!("a/b/1.txt", "b/1.txt", PATH_SAME_BASENAME_SCORE);
        assert_file_path_similarity!("a/b/1.txt", "b/c/d/1.txt", PATH_SAME_BASENAME_SCORE);
    }

    #[test]
    fn test_paths_with_same_directory() {
        assert_file_path_similarity!("1.txt", "2.txt", PATH_SAME_DIRECTORY_SCORE);
        assert_file_path_similarity!("a/1.txt", "a/2.txt", PATH_SAME_DIRECTORY_SCORE);
        assert_file_path_similarity!("a/b/1.txt", "a/b/2.txt", PATH_SAME_DIRECTORY_SCORE);
    }

    #[test]
    fn test_paths_without_same_basename_or_directory() {
        assert_file_path_similarity!("a/1.txt", "2.txt", PATH_DEFAULT_SCORE);
        assert_file_path_similarity!("a/1.txt", "b/2.txt", PATH_DEFAULT_SCORE);
    }
}


/*
 * 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 cpython::*;
use cpython_ext::ResultPyErrExt;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;

// Type to make Python able to reason about a Rust future.
//
// Unlike TStream, it does not support typed lossless Python -> Rust conversion.
py_class!(pub class future |py| {
    data inner: RefCell<Option<BoxFuture<'static, PyResult<PyObject>>>>;

    /// Resolve the future and return the resolved object.
    def wait(&self) -> PyResult<PyObject> {
        let mut inner = None;
        std::mem::swap(&mut inner, &mut self.inner(py).borrow_mut());

        match inner {
            Some(future) => py.allow_threads(|| async_runtime::block_on(future)),
            None => Err(PyErr::new::<exc::ValueError, _>(py, "future was already waited")),
        }
    }
});

impl future {
    /// Convert Rust Future to Python object.
    pub fn new<T, E, F>(py: Python, f: F) -> PyResult<Self>
    where
        T: ToPyObject,
        E: Into<anyhow::Error>,
        F: Future<Output = Result<T, E>> + Send + 'static,
    {
        let future = f.map(|v| {
            let gil = Python::acquire_gil();
            let py = gil.python();
            v.map_pyerr(py).map(|v| v.into_py_object(py).into_object())
        });
        Self::create_instance(py, RefCell::new(Some(Box::pin(future))))
    }
}


/*
 * 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)]

//! Utilities to make async <-> Python integration easier.
//!
//! The `TStream` type provides easy conversion between Rust `Stream` and Python
//! objects.
//!
//! The `PyFuture` type provides a way to export Rust `Future` to
//! Python.

mod future;
mod stream;

// Re-export.
pub use anyhow;
pub use async_runtime;
pub use cpython;
pub use cpython_ext;
pub use future::future as PyFuture;
pub use futures;
pub use stream::TStream;


/*
 * 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 cpython::*;
use cpython_ext::AnyhowResultExt;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;
use futures::stream::BoxStream;
use futures::stream::TryStreamExt;
use futures::Stream;

/// `TStream` is a thin wrapper of `Stream` from async Rust to Python.
///
/// `TStream` can be used as both input or output parameters in cpython binding
/// functions:
///
/// ```
/// # use cpython_async::{*, cpython::*, futures::*};
/// type S = TStream<anyhow::Result<Vec<u8>>>;
///
/// // Pass a stream to Python:
/// fn produce_stream(py: Python) -> PyResult<S> {
///     Ok(futures::stream::once(async { Ok(vec![1]) }).into())
/// }
///
/// // Receive a stream (iterable) from Python:
/// fn map_reverse_stream(py: Python, tstream: S) -> PyResult<S> {
///     // Use `.stream()` to extract the pure Rust stream object
///     // that implements `Stream`.
///     let stream = tstream.stream();
///     let stream = stream.map_ok(|mut x| { x.reverse(); x });
///     Ok(stream.into())
/// }
/// ```
///
/// In Python, the stream can be used as an iterator, functions accepting
/// streams also accepts iterators:
///
/// ```python,ignore
/// stream = rustmod.produce_stream()
/// stream = rustmod.map_reverse_stream(stream)
/// for value in stream:
///     print(value)
/// ```
pub struct TStream<T>(BoxStream<'static, T>);

impl<I> TStream<I> {
    /// Converts to `BoxStream` which implements the `Stream` trait.
    pub fn stream(self) -> BoxStream<'static, I> {
        self.0
    }
}

// This is convenient but prevents TStream from implementing Stream.
impl<S, I> From<S> for TStream<I>
where
    S: Stream<Item = I> + Send + 'static,
{
    fn from(s: S) -> Self {
        TStream(Box::pin(s))
    }
}

mod pytypes {
    use super::*;

    // Convert from a Python iterable to TStream.
    impl<'s, T> FromPyObject<'s> for TStream<anyhow::Result<T>>
    where
        T: for<'b> FromPyObject<'b> + Send + 'static,
    {
        fn extract(py: Python, obj: &'s PyObject) -> PyResult<Self> {
            let pyiter = obj.iter(py)?.into_object();
            let iter = itertools::unfold(pyiter, |pyiter| {
                let item = (|pyiter: &PyObject| -> PyResult<Option<T>> {
                    let gil = Python::acquire_gil();
                    let py = gil.python();
                    let mut iter = pyiter.iter(py)?;
                    if let Some(v) = iter.next() {
                        Ok(Some(v?.extract::<T>(py)?))
                    } else {
                        Ok(None)
                    }
                })(&pyiter);
                item.into_anyhow_result().transpose()
            });
            // async_runtime::iter_to_stream supports blocking `next` calls.
            // futures::stream::iter doesn't. If futures::stream::iter is used,
            // then test_nested_stream_to_and_from_python() will hang.
            let stream = async_runtime::iter_to_stream(iter);
            return Ok(stream.into());
        }
    }

    // Convert TStream to a Python object.
    impl<T> ToPyObject for TStream<anyhow::Result<T>>
    where
        T: ToPython,
    {
        type ObjectType = PyObject;

        fn to_py_object(&self, _py: Python) -> Self::ObjectType {
            panic!("bug: TStream::to_py_object should not be used");
        }

        fn into_py_object(self, py: Python) -> Self::ObjectType {
            // Erase the type. Do not convert to PyObject directly to avoid GIL cost.
            // ('py' cannot be used in the stream 'map' closure).
            let inner = self.0.map_ok(|t| Box::new(t) as Box<dyn ToPython>);
            let typename = std::any::type_name::<T>().to_string();
            let result =
                pytypes::stream::create_instance(py, RefCell::new(Some(Box::pin(inner))), typename);
            result.unwrap().into_object()
        }
    }

    /// Like `ToPyObject` but without `type Target`.
    pub trait ToPython: Send + 'static {
        fn to_py(&self, py: Python) -> PyObject;
    }

    impl<T: ToPyObject + Send + 'static> ToPython for T {
        fn to_py(&self, py: Python) -> PyObject {
            self.to_py_object(py).into_object()
        }
    }

    py_class!(pub class stream |py| {
        data state: RefCell<Option<BoxStream<'static, anyhow::Result<Box<dyn ToPython>>>>>;
        data type_name: String;

        def __iter__(&self) -> PyResult<streamiter> {
            let state = self.state(py).borrow_mut().take();
            match state {
                Some(state) => {
                    let iter = async_runtime::stream_to_iter(state);
                    streamiter::create_instance(py, RefCell::new(Box::new(iter)))
                }
                None => {
                    Err(PyErr::new::<exc::ValueError, _>(py, "stream was consumed"))
                }
            }
        }

        def typename(&self) -> PyResult<Str> {
            Ok(self.type_name(py).clone().into())
        }
    });

    py_class!(pub class streamiter |py| {
        data iter: RefCell<Box<dyn Iterator<Item = anyhow::Result<Box<dyn ToPython>>> + Send>>;

        def __next__(&self) -> PyResult<Option<PyObject>> {
            // safety: py.allow_threads runs in the same thread.
            // Its 'Send' requirement is just to disallow passing 'py'.
            struct ForceSend<T>(T);
            unsafe impl<T> Send for ForceSend<T> {}
            let iter = ForceSend(self.iter(py).borrow_mut());
            // py.allow_threads is needed because iter.next might take Python GIL.
            let next = py.allow_threads(|| {
                let mut iter = iter; // capture ForceSend into closure
                iter.0.next()
            });
            match next {
                None => Ok(None),
                Some(result) => {
                    let v = result.map_pyerr(py)?;
                    let obj = v.to_py(py);
                    Ok(Some(obj))
                }
            }
        }

        def __iter__(&self) -> PyResult<Self> {
            Ok(self.clone_ref(py))
        }
    });
}

#[cfg(test)]
mod tests {
    use futures::stream::StreamExt;

    use super::*;

    #[tokio::test]
    async fn test_stream_from_python() {
        let tstream: TStream<anyhow::Result<usize>> = with_py(|py| {
            let input = vec![3, 10, 20].into_py_object(py).into_object();
            input.extract(py).unwrap()
        });
        let mut stream = tstream.stream();
        assert_eq!(stream.next().await.transpose().unwrap(), Some(3));
        assert_eq!(stream.next().await.transpose().unwrap(), Some(10));
        assert_eq!(stream.next().await.transpose().unwrap(), Some(20));
        assert_eq!(stream.next().await.transpose().unwrap(), None);
        assert_eq!(stream.next().await.transpose().unwrap(), None);
    }

    #[test]
    fn test_stream_to_python() {
        let orig_vec = vec![5, 20, 10];
        let stream = futures::stream::iter(orig_vec.clone());
        let tstream: TStream<anyhow::Result<usize>> = stream.map(Ok).into();
        let new_vec: Vec<usize> = with_py(|py| {
            let pyobj = tstream.into_py_object(py);
            let pyiter = pyobj.iter(py).unwrap();
            pyiter
                .map(|r| with_py(|py| r.unwrap().extract::<usize>(py).unwrap()))
                .collect()
        });
        assert_eq!(new_vec, orig_vec);
    }

    #[tokio::test]
    async fn test_nested_stream_to_and_from_python() {
        let pyplus1 = with_py(|py| {
            py.run(
                r#"def plus1(stream):
                       for item in stream:
                           if item == 5:
                               raise RuntimeError("dislike this number")
                           yield item + 1"#,
                None,
                None,
            )
            .unwrap();
            py.eval("plus1", None, None).unwrap()
        });

        let stream = futures::stream::iter((5..9).rev());
        let tstream: TStream<anyhow::Result<usize>> = stream.map(Ok).into();

        let plus1 = |tstream: TStream<anyhow::Result<usize>>| -> TStream<anyhow::Result<usize>> {
            with_py(|py| {
                // TStream -> PyObject
                let arg = tstream.into_py_object(py);
                let obj: PyObject = pyplus1.call(py, (arg,), None).unwrap();
                // PyObject -> TStream
                obj.extract(py).unwrap()
            })
        };

        let tstream = plus1(tstream);
        let tstream = plus1(tstream);
        let tstream = plus1(tstream);

        let mut stream = tstream.stream();
        assert_eq!(stream.next().await.map(|v| v.unwrap()), Some(11));
        assert_eq!(stream.next().await.map(|v| v.unwrap()), Some(10));
        assert_eq!(stream.next().await.map(|v| v.unwrap()), Some(9));

        let err = stream.next().await.unwrap().unwrap_err();
        with_py(|_py| {
            assert!(format!("{}", err).contains("dislike this number"));
        });

        // The traceback includes the nested (3) functions.
        let traceback: String = with_py(|py| {
            let pyerr = Err::<u8, _>(err).map_pyerr(py).unwrap_err();
            let traceback: PyObject = pyerr.ptraceback.unwrap();
            let formatted = py
                .import("traceback")
                .unwrap()
                .call(py, "format_tb", (traceback,), None)
                .unwrap();
            formatted.extract::<Vec<String>>(py).unwrap().concat()
        });
        assert_eq!(
            traceback,
            r#"  File "<string>", line 2, in plus1
  File "<string>", line 2, in plus1
  File "<string>", line 4, in plus1
"#
        );

        assert_eq!(stream.next().await.map(|v| v.unwrap()), None);
        assert_eq!(stream.next().await.map(|v| v.unwrap()), None);
    }

    fn with_py<R>(f: impl FnOnce(Python) -> R) -> R {
        let gil = Python::acquire_gil();
        let py = gil.python();
        f(py)
    }
}


/*
 * 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::mem;

use cpython::PyObject as RustPyObject;
use cpython::Python as RustPythonGILGuard;
use ffi::PyByteArray_Size;
use ffi::PyObject;
use ffi::PyVarObject;
use ffi::Py_ssize_t;
use python3_sys as ffi;

// Ideally this struct comes from bindgen to ensure it's up-to-date
// with what the current cpython has. However that significantly
// complicates the build step.
#[repr(C)]
struct PyByteArrayObject {
    pub ob_base: PyVarObject,
    pub ob_alloc: Py_ssize_t,
    pub ob_bytes: *mut u8,
    pub ob_start: *mut u8,
    // These fields exist in cpython but we don't care about them.
    // pub ob_exports: c_int | Py_ssize_t,
}

/// Consume a `Vec<u8>`. Create a Python `bytearray` object.
/// Bytes stored are not copied.
pub fn vec_to_pyobj(py: RustPythonGILGuard<'_>, mut value: Vec<u8>) -> RustPyObject {
    let capacity = value.capacity();
    let obj = unsafe { pybytearray_from_slice(py, value.as_mut(), Some(capacity)) };
    // now cpython is responsible to free `value`.
    mem::forget(value);
    obj
}

/// Consume a `Box<[u8]>`. Create a Python `bytearray` object.
/// Bytes stored are not copied.
pub fn boxed_slice_to_pyobj(py: RustPythonGILGuard<'_>, mut value: Box<[u8]>) -> RustPyObject {
    let obj = unsafe { pybytearray_from_slice(py, value.as_mut(), None) };
    // now cpython is responsible to free `value`.
    mem::forget(value);
    obj
}

/// Similar to `PyByteArray_FromStringAndSize(value)` without copying `value`.
///
/// CPython will free `value` when releasing the python object. When `value` is
/// a prefix of a larger buffer, set `alloc_size` so CPython can free the buffer
/// properly.
///
/// To avoid double-free from both CPython and Rust, the callsite must
/// `mem::forget` the owner of `value` after this function.
unsafe fn pybytearray_from_slice(
    py: RustPythonGILGuard,
    value: &mut [u8],
    alloc_size: Option<usize>,
) -> RustPyObject {
    let ptr: *mut PyObject = ffi::PyByteArray_FromStringAndSize(std::ptr::null(), 0);
    let typed: *mut PyByteArrayObject = mem::transmute(ptr);

    // We want to rewrite ob_alloc and ob_bytes directly. Ensure that
    // they are not allocated by cpython.
    assert_eq!(
        (*typed).ob_alloc,
        0,
        "PyByteArray_FromStringAndSize(size=0) should not alloc"
    );
    assert!(
        (*typed).ob_bytes.is_null(),
        "PyByteArray_FromStringAndSize(size=0) should not alloc"
    );

    // There are no public cpython APIs to update these fields.
    // So we update them directly based on our understanding of
    // the PyByteArrayObject struct.
    (*typed).ob_alloc = alloc_size.unwrap_or_else(|| value.len()) as Py_ssize_t;
    (*typed).ob_bytes = value.as_mut_ptr();
    {
        (*typed).ob_base.ob_size = value.len() as Py_ssize_t;
        (*typed).ob_start = (*typed).ob_bytes;
    }

    // Sanity check: cpython's understanding of the bytes matches ours.
    let ptr: *mut PyObject = typed as *mut PyObject;
    assert_eq!(
        PyByteArray_Size(ptr) as usize,
        value.len(),
        "PyByteArray struct mismatch (ob_size)"
    );
    assert_eq!(
        ffi::PyByteArray_AsString(ptr) as *const u8,
        value.as_ptr(),
        "PyByteArray struct mismatch (ob_bytes)"
    );

    RustPyObject::from_owned_ptr(py, ptr)
}


/*
 * 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::*;

#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Default, Hash, Ord)]
pub struct Bytes(pub Box<[u8]>);

impl ToPyObject for Bytes {
    type ObjectType = PyBytes;

    #[inline]
    fn to_py_object(&self, py: Python) -> PyBytes {
        PyBytes::new(py, &self.0)
    }
}

impl<'source> FromPyObject<'source> for Bytes {
    fn extract(py: Python, obj: &'source PyObject) -> PyResult<Self> {
        let data = obj.cast_as::<PyBytes>(py)?.data(py);
        Ok(Bytes(data.to_vec().into_boxed_slice()))
    }
}

impl From<Box<[u8]>> for Bytes {
    fn from(v: Box<[u8]>) -> Bytes {
        Bytes(v)
    }
}

impl From<Vec<u8>> for Bytes {
    fn from(v: Vec<u8>) -> Bytes {
        Bytes(v.into_boxed_slice())
    }
}

impl From<String> for Bytes {
    fn from(s: String) -> Bytes {
        Bytes(s.into_bytes().into_boxed_slice())
    }
}

impl AsRef<[u8]> for Bytes {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}


/*
 * 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::mem;
use std::os::raw::c_char;
use std::slice;

use cpython::PyObject as RustPyObject;
use cpython::Python as RustPythonGILGuard;
use ffi::PyBytes_Type;
use ffi::PyObject;
use ffi::PyTypeObject;
use ffi::PyVarObject;
use ffi::Py_hash_t;
use ffi::Py_ssize_t;
use ffi::_PyObject_NewVar;
use python3_sys as ffi;

// From Python bytesobject.h. Must match the C definition.
#[repr(C)]
struct PyBytesObject {
    pub ob_base: PyVarObject,
    pub ob_shash: Py_hash_t,
    pub ob_sval: [c_char; 1],
}

/// Create a `PyBytes` object that have `size` bytes. Return the object and
/// its internal buffer to be written. This is useful to bypass the memcpy
/// cost creating a large `PyBytesObject`.
pub fn allocate_pybytes(py: RustPythonGILGuard<'_>, size: usize) -> (RustPyObject, &mut [u8]) {
    unsafe {
        let ptr: *mut PyVarObject = _PyObject_NewVar(
            &mut PyBytes_Type as *mut PyTypeObject,
            (size + mem::size_of::<PyBytesObject>()) as Py_ssize_t,
        );
        let mut ptr: *mut PyObject = mem::transmute(ptr);
        let typed: *mut PyBytesObject = mem::transmute(ptr);
        (*typed).ob_shash = -1; // hash not calculated
        {
            (*typed).ob_base.ob_size = size as Py_ssize_t;
        }
        // Set the first byte to '\0'. If the caller forgot to populate the
        // slice, PyBytes_AsString would still return an empty C string.
        (*typed).ob_sval[0] = 0;
        // Set the byte after the slice to '\0' If the caller does populate the
        // slice, PyBytes_AsString would return a C string with tailing '\0'.
        *((*typed).ob_sval.as_mut_ptr().offset(size as isize)) = 0;
        let pptr: *mut *mut PyObject = &mut ptr;
        ptr = *pptr;
        assert!(!ptr.is_null(), "out of memory");
        let sval: *mut u8 = mem::transmute((*typed).ob_sval.as_mut_ptr());
        let slice = slice::from_raw_parts_mut(sval, size);
        (RustPyObject::from_owned_ptr(py, ptr), slice)
    }
}


/*
 * 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::any::Any;
use std::cell::RefCell;
use std::fmt::Debug;

use ::serde::de::DeserializeOwned;
use ::serde::Serialize;
use cpython::*;

use crate::convert::Serde;
use crate::none::PyNone;

type BoxedAny = Box<dyn Any + Sync + Send + 'static>;

// PyCell allows to put arbitrary rust data and pass it between different rust functions through python code
// This allows to avoid using bincode or writing wrapper types for some basic use cases
py_class!(pub class PyCell |py| {
    data inner: RefCell<Option<BoxedAny>>;
    data fmt: Box<dyn (Fn(&BoxedAny) -> String) + Send + Sync>;
    data fn_export: Box<dyn (Fn(&BoxedAny, Python) -> PyObject) + Send + Sync>;
    data fn_import: Box<dyn (Fn(PyObject, Python) -> PyResult<BoxedAny>) + Send + Sync>;

    def __str__(&self) -> PyResult<String> {
        let str = self.inner(py).borrow().as_ref().map(|inner| {
                let fmt = self.fmt(py);
                fmt(inner)
        });
        let str = str.unwrap_or_else(||"<None>".to_string());
        Ok(str)
    }

    def export(&self) -> PyResult<Option<PyObject>> {
        let pyobj = self.inner(py).borrow().as_ref().map(|inner| {
            let export = self.fn_export(py);
            export(inner, py)
        });
        Ok(pyobj)
    }

    def import(&self, obj: PyObject) -> PyResult<PyNone> {
        let import = self.fn_import(py);
        let obj = import(obj, py)?;
        let inner = self.inner(py);
        *inner.borrow_mut() = Some(obj);
        Ok(PyNone)
    }
});

impl PyCell {
    pub fn new<T: DeserializeOwned + Serialize + Debug + Sync + Send + 'static + Sized>(
        py: Python,
        data: T,
    ) -> PyResult<Self> {
        let inner = Box::new(data) as BoxedAny;
        let fmt = |obj: &BoxedAny| {
            let obj = obj.as_ref().downcast_ref::<T>().unwrap(); // does not fail
            format!("{:?}", obj)
        };
        let export = |obj: &BoxedAny, py: Python| {
            let obj = obj.as_ref().downcast_ref::<T>().unwrap(); // does not fail
            Serde(obj).to_py_object(py)
        };
        let import = |obj: PyObject, py: Python| -> PyResult<BoxedAny> {
            let obj: Serde<T> = Serde::extract(py, &obj)?;
            Ok(Box::new(obj.0) as BoxedAny)
        };
        Self::create_instance(
            py,
            RefCell::new(Some(inner)),
            Box::new(fmt),
            Box::new(export),
            Box::new(import),
        )
    }

    pub fn take<T: Sync + Send + 'static + Sized>(&self, py: Python) -> Option<Box<T>> {
        match self.inner(py).borrow_mut().take() {
            Some(x) => x.downcast::<T>().ok(),
            None => 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.
 */

use std::any;
use std::any::Any;
use std::any::TypeId;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::Hash;
use std::hash::Hasher;
use std::ops::Deref;

use ::serde::Deserialize;
use ::serde::Serialize;
use cpython::FromPyObject;
use cpython::PyObject;
use cpython::PyType;
use cpython::Python;
use cpython::PythonObjectDowncastError;
use cpython::PythonObjectWithTypeObject;
use cpython::_detail::ffi;
use cpython::*;
use once_cell::sync::Lazy;
use parking_lot::RwLock;

/// Wrapper type. Converts between pure Rust bytes-like types and PyBytes.
///
/// The Rust type needs to implement `AsRef<[u8]>` and `From<Vec<u8>>`.
///
/// In bindings code:
/// - For input, use `v: BytesLike<MyType>` in definition, and `v.0` to extract
///   `MyType`.
/// - For output, use `-> BytesLike<MyType>` in definition, and `BytesLike(v)`
///   to construct the return value.
#[derive(Clone, Debug, Eq, Ord, Hash, PartialEq, PartialOrd)]
pub struct BytesLike<T>(pub T);

impl<T: AsRef<[u8]>> ToPyObject for BytesLike<T> {
    type ObjectType = PyBytes;

    fn to_py_object(&self, py: Python) -> Self::ObjectType {
        PyBytes::new(py, self.0.as_ref())
    }
}

impl<'s, T: From<Vec<u8>>> FromPyObject<'s> for BytesLike<T> {
    fn extract(py: Python, obj: &'s PyObject) -> PyResult<Self> {
        obj.extract::<PyBytes>(py)
            .map(|v| Self(v.data(py).to_vec().into()))
    }
}

impl<T> Deref for BytesLike<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Wrapper type. Converts between pure Rust serde types and PyObjct.
///
/// In bindings code:
/// - For input, use `v: Serde<MyType>` in definition, and `v.0` to extract
///   `MyType`.
/// - For output, use `-> Serde<MyType>` in definition, and `Serde(v)` to
///   construct the return value.
#[derive(Debug)]
pub struct Serde<T>(pub T);

impl<T: Serialize> ToPyObject for Serde<T> {
    type ObjectType = PyObject;

    fn to_py_object(&self, py: Python) -> Self::ObjectType {
        crate::ser::to_object(py, &self.0).unwrap()
    }
}

impl<'s, T> FromPyObject<'s> for Serde<T>
where
    T: for<'de> Deserialize<'de>,
{
    fn extract(py: Python, obj: &'s PyObject) -> PyResult<Self> {
        let inner = crate::de::from_object(py, obj.clone_ref(py))?;
        Ok(Self(inner))
    }
}

impl<T> Deref for Serde<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: PartialOrd> PartialOrd<Serde<T>> for Serde<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<T: Ord> Ord for Serde<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl<T: Hash> Hash for Serde<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state)
    }
}

impl<T: PartialEq> PartialEq for Serde<T> {
    fn eq(&self, other: &Self) -> bool {
        self.0.eq(&other.0)
    }
}

impl<T: Eq> Eq for Serde<T> {}

impl<T: Clone> Clone for Serde<T> {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

/// Registered functions about converting PyObject (type decided by PyType)
/// to Box<T> (type decided by TypeId).
///
/// The key of the map is (input Python type, output Rust type).
static CONVERT_FUNC_BY_TYPE: Lazy<RwLock<HashMap<(PyTypeId, TypeId), ExtractPyObjectFunc>>> =
    Lazy::new(|| Default::default());

type ExtractPyObjectFunc =
    Box<dyn (Fn(Python, &PyObject) -> Box<dyn Any + Send + Sync>) + Send + Sync>;

/// An `usize` used to identify CPython types.
#[derive(Hash, Eq, PartialEq, Ord, PartialOrd)]
struct PyTypeId(*const ffi::PyTypeObject);

// safety: PyTypeObject is a static pointer. It can be treated as an
// integer to compare.
unsafe impl Send for PyTypeId {}
unsafe impl Sync for PyTypeId {}

impl From<&'_ PyType> for PyTypeId {
    fn from(py_type: &PyType) -> Self {
        Self(py_type.as_type_ptr() as _)
    }
}

impl PyTypeId {
    fn null() -> Self {
        Self(std::ptr::null())
    }
}

/// Register a function to convert a PyObject (Python type: P)
/// to type O using the specified function.
///
/// After registration, `ImplInto<O>` can be used in Python function
/// definitions.
///
/// If P is `PyObject`, then it matches heap types (usually, objects
/// created by instantiating a pure Python `class`). It does not
/// match non-heap types such as `PyBytes`.
pub fn register_into<P, F, O>(py: Python, convert_func: F)
where
    P: PythonObjectWithTypeObject,
    P: for<'a> FromPyObject<'a>,
    F: Fn(Python, P) -> O,
    F: Send + Sync + 'static,
    O: Send + Sync + 'static,
{
    let py_type_id = if TypeId::of::<P>() == TypeId::of::<PyObject>() {
        PyTypeId::null()
    } else {
        let py_type = P::type_object(py);
        assert!(
            !is_heap_type(py, &py_type),
            "native heap type is unsupported"
        );
        PyTypeId::from(&py_type)
    };
    let output_type_id = TypeId::of::<O>();
    let func = move |py: Python, obj: &PyObject| -> Box<dyn Any + Send + Sync> {
        let obj: P = obj.extract::<P>(py).expect("PyTypeId was checked");
        let obj: O = convert_func(py, obj);
        Box::new(obj) as Box<dyn Any + Send + Sync>
    };
    CONVERT_FUNC_BY_TYPE
        .write()
        .insert((py_type_id, output_type_id), Box::new(func));
}

/// Wrapper type. Converts `PyObject` to `T`.
///
/// How to convert `PyObject` to `T` needs to be pre-registered via `register_into`.
pub struct ImplInto<T>(pub T);

impl<'s, T> FromPyObject<'s> for ImplInto<T>
where
    T: 'static,
{
    fn extract(py: Python, obj: &'s PyObject) -> PyResult<Self> {
        let py_type = obj.get_type(py);
        let py_type_id = if is_heap_type(py, &py_type) {
            PyTypeId::null()
        } else {
            PyTypeId::from(&py_type)
        };
        let output_type_id = TypeId::of::<T>();

        let table = CONVERT_FUNC_BY_TYPE.read();
        match table.get(&(py_type_id, output_type_id)) {
            Some(convert_func) => {
                let inner: Box<T> = (convert_func)(py, obj)
                    .downcast::<T>()
                    .expect("TypeId was checked");
                Ok(ImplInto(*inner))
            }
            None => {
                let expected_type_name = any::type_name::<Self>();
                Err(PythonObjectDowncastError::new(py, expected_type_name, py_type).into())
            }
        }
    }
}

impl<T> ImplInto<T> {
    /// Convert to the inner type.
    pub fn into(self) -> T {
        self.0
    }
}

/// Test whether `typeobj` is a heap type. A heap type is usually
/// defined in pure Python using the `class` keyword and has a
/// `__dict__` slot supporting `setattr`. `py_class!` types or
/// builtin types like bytes, str and are not heap types.
fn is_heap_type(_py: Python, typeobj: &PyType) -> bool {
    let type_ptr: *mut ffi::PyTypeObject = typeobj.as_type_ptr();
    // safety: _py holds GIL. The pointer is valid.
    let result = (unsafe { *type_ptr }.tp_flags & ffi::Py_TPFLAGS_HEAPTYPE) != 0;
    result
}

#[cfg(test)]
mod tests {
    use cpython::*;

    use super::*;

    py_class!(class T1 |py| {
        data v: String;
    });

    py_class!(class T2 |py| {
        data v: i32;
    });

    #[test]
    fn test_impl_into() {
        let gil = Python::acquire_gil();
        let py = gil.python();

        // Register how to convert T1, T2 to String.
        register_into(py, |py, t: T1| t.v(py).clone());
        register_into(py, |py, t: T2| t.v(py).to_string());
        register_into(py, |_py, t: PyObject| t.to_string());

        let v1 = T1::create_instance(py, "abc".to_string()).unwrap();
        let v2 = T2::create_instance(py, 123).unwrap();
        let p1: PyObject = v1.into_object();
        let p2: PyObject = v2.into_object();

        // p3 is a heap type.
        let p3: PyObject = py.eval("type('Foo', (object,), {})()", None, None).unwrap();

        // Converting to Impl<String> works without casting to T1 or T2 first.
        let s1 = p1.extract::<ImplInto<String>>(py).unwrap();
        let s2 = p2.extract::<ImplInto<String>>(py).unwrap();
        let s3 = p3.extract::<ImplInto<String>>(py).unwrap();

        assert_eq!(s1.into(), "abc");
        assert_eq!(s2.into(), "123");
        assert_eq!(
            // strip the unstable "at 0x......" part.
            s3.into().split_whitespace().next().unwrap(),
            "<__main__.Foo"
        );

        // p4 is a builtin (native, non-heap) type that is not registered.
        let p4: PyObject = PyBytes::new(py, b"Foo").into_object();
        assert!(p4.extract::<ImplInto<String>>(py).is_err());
    }
}


/*
 * 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 ::serde::de;
use ::serde::de::Visitor;
use cpython::*;

use crate::PyErr as Error;
use crate::PyNone;

type Result<T> = std::result::Result<T, Error>;

/// Deserialize from Python object.
pub fn from_object<'de, T>(py: Python, obj: PyObject) -> Result<T>
where
    T: de::Deserialize<'de>,
{
    let mut deserializer = Deserializer::new(py, obj);
    T::deserialize(&mut deserializer)
}

struct Deserializer<'gil> {
    py: Python<'gil>,
    obj: PyObject,

    /// Iterator of `obj`. Constructed on demand.
    obj_iter: Option<PyIterator<'gil>>,

    /// Used by `MapAccess`. The `iter` produces `(key, value)` at a time, the
    /// `value` is temporarily stored here for `MapAccess::next_value_seed` to
    /// pick up.
    pending_value: Option<PyObject>,
}

impl<'gil> Deserializer<'gil> {
    /// Constructs from a Python object.
    fn new(py: Python<'gil>, obj: PyObject) -> Self {
        Self {
            py,
            obj,
            obj_iter: None,
            pending_value: None,
        }
    }

    /// Returns the next item from a Python iterator (obj should be an iterable).
    fn next(&mut self) -> Result<Option<PyObject>> {
        match self.obj_iter {
            None => {
                // Convert `obj` to a Python iterator object.
                // Usually this is a no-op for a type that is already an iterator.
                //
                // Special case: for PyDict, call the "items" method first to get
                // an iterator of (key, value) instead of just keys.
                let iter = if self.extract::<PyDict>().is_ok() {
                    let items = self.obj.call_method(self.py, "items", NoArgs, None)?;
                    items.iter(self.py)?
                } else {
                    self.obj.iter(self.py)?
                };
                self.obj_iter = Some(iter);
                self.next()
            }
            Some(ref mut iter) => iter.next().transpose().map_err(Into::into),
        }
    }

    /// Attempt to convert `obj` to a given type.
    fn extract<T>(&self) -> Result<T>
    where
        for<'s> T: FromPyObject<'s>,
    {
        self.obj.extract(self.py).map_err(Into::into)
    }
}

impl<'de, 'a, 'gil> de::Deserializer<'de> for &'a mut Deserializer<'gil> {
    type Error = Error;

    fn deserialize_any<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        if self.extract::<PyNone>().is_ok() {
            v.visit_none()
        } else if self.extract::<bool>().is_ok() {
            self.deserialize_bool(v)
        } else if self.extract::<PyDict>().is_ok() {
            self.deserialize_map(v)
        } else if self.extract::<PyList>().is_ok() || self.extract::<PyTuple>().is_ok() {
            self.deserialize_seq(v)
        } else if self.extract::<PyBytes>().is_ok() {
            self.deserialize_bytes(v)
        } else if self.extract::<PyString>().is_ok() {
            self.deserialize_string(v)
        } else if self.extract::<i64>().is_ok() {
            self.deserialize_i64(v)
        } else if self.extract::<u64>().is_ok() {
            self.deserialize_u64(v)
        } else if self.extract::<f64>().is_ok() {
            self.deserialize_f64(v)
        } else {
            Err(PyErr::new::<exc::TypeError, _>(
                self.py,
                "deserialize_any encountered undecided type",
            )
            .into())
        }
    }

    // Code generation for simple types.
    // Run with: python -m cogapp -r src/de.rs
    //
    // [[[cog
    // import cog
    // for name in 'bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 string'.split():
    //     cog.out("""
    //     fn deserialize_{name}<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {{
    //         v.visit_{name}(self.extract()?)
    //     }}
    // """.format(name=name))
    // ]]]

    fn deserialize_bool<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_bool(self.extract()?)
    }

    fn deserialize_i8<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_i8(self.extract()?)
    }

    fn deserialize_i16<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_i16(self.extract()?)
    }

    fn deserialize_i32<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_i32(self.extract()?)
    }

    fn deserialize_i64<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_i64(self.extract()?)
    }

    fn deserialize_u8<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_u8(self.extract()?)
    }

    fn deserialize_u16<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_u16(self.extract()?)
    }

    fn deserialize_u32<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_u32(self.extract()?)
    }

    fn deserialize_u64<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_u64(self.extract()?)
    }

    fn deserialize_f32<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_f32(self.extract()?)
    }

    fn deserialize_f64<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_f64(self.extract()?)
    }

    fn deserialize_string<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_string(self.extract()?)
    }
    // [[[end]]]

    fn deserialize_char<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.deserialize_string(v)
    }

    fn deserialize_str<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.deserialize_string(v)
    }

    fn deserialize_bytes<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        if self.obj.get_type(self.py) == PyString::type_object(self.py) {
            let s: String = self.extract()?;
            v.visit_bytes(s.as_bytes())
        } else {
            let pybytes: PyBytes = self.extract()?;
            v.visit_bytes(pybytes.data(self.py))
        }
    }

    fn deserialize_byte_buf<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.deserialize_bytes(v)
    }

    fn deserialize_option<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        if let Ok(PyNone) = self.extract() {
            v.visit_none()
        } else {
            v.visit_some(self)
        }
    }

    fn deserialize_unit<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.extract::<PyNone>()?;
        v.visit_unit()
    }

    fn deserialize_unit_struct<V: Visitor<'de>>(self, _: &'static str, v: V) -> Result<V::Value> {
        self.deserialize_unit(v)
    }

    fn deserialize_newtype_struct<V: Visitor<'de>>(
        self,
        _: &'static str,
        v: V,
    ) -> Result<V::Value> {
        v.visit_newtype_struct(self)
    }

    fn deserialize_seq<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_seq(self)
    }

    fn deserialize_tuple<V: Visitor<'de>>(self, _len: usize, v: V) -> Result<V::Value> {
        self.deserialize_seq(v)
    }

    fn deserialize_tuple_struct<V: Visitor<'de>>(
        self,
        _name: &'static str,
        _len: usize,
        v: V,
    ) -> Result<V::Value> {
        self.deserialize_seq(v)
    }

    fn deserialize_map<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        v.visit_map(self)
    }

    fn deserialize_struct<V: Visitor<'de>>(
        self,
        _name: &'static str,
        _fields: &'static [&'static str],
        v: V,
    ) -> Result<V::Value> {
        self.deserialize_map(v)
    }

    fn deserialize_enum<V: Visitor<'de>>(
        self,
        _name: &'static str,
        _variants: &'static [&'static str],
        v: V,
    ) -> Result<V::Value> {
        v.visit_enum(self)
    }

    fn deserialize_identifier<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.deserialize_string(v)
    }

    fn deserialize_ignored_any<V: Visitor<'de>>(self, v: V) -> Result<V::Value> {
        self.deserialize_any(v)
    }
}

impl<'de, 'a, 'gil> de::SeqAccess<'de> for &'a mut Deserializer<'gil> {
    type Error = Error;

    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>>
    where
        T: de::DeserializeSeed<'de>,
    {
        match self.next()? {
            Some(obj) => {
                let mut deserializer = Deserializer::new(self.py, obj);
                seed.deserialize(&mut deserializer).map(Some)
            }
            None => Ok(None),
        }
    }
}

impl<'de, 'a, 'gil> de::MapAccess<'de> for &'a mut Deserializer<'gil> {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
    where
        K: de::DeserializeSeed<'de>,
    {
        match self.next()? {
            Some(obj) => {
                let (key, value): (PyObject, PyObject) = obj.extract(self.py)?;
                self.pending_value = Some(value);
                let mut deserializer = Deserializer::new(self.py, key);
                seed.deserialize(&mut deserializer).map(Some)
            }
            None => Ok(None),
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
    where
        V: de::DeserializeSeed<'de>,
    {
        match self.pending_value.take() {
            Some(obj) => {
                let mut deserializer = Deserializer::new(self.py, obj);
                seed.deserialize(&mut deserializer)
            }
            None => Err(PyErr::new::<exc::ValueError, _>(
                self.py,
                "no value for MapAccess::next_value_seed to pick up",
            )
            .into()),
        }
    }
}

impl<'de, 'a, 'gil> de::EnumAccess<'de> for &'a mut Deserializer<'gil> {
    type Error = Error;
    type Variant = Deserializer<'gil>;

    fn variant_seed<V>(self, seed: V) -> Result<(V::Value, Self::Variant)>
    where
        V: de::DeserializeSeed<'de>,
    {
        if self.extract::<String>().is_ok() {
            // a string for unit enum variants.
            let name = seed.deserialize(&mut *self)?;
            Ok((name, Deserializer::new(self.py, self.py.None())))
        } else {
            // a dict for complex enum variants.
            let dict: PyDict = self.extract()?;
            let items: Vec<(PyObject, PyObject)> = dict.items(self.py);
            if items.len() != 1 {
                let repr = self.obj.repr(self.py)?;
                let repr = repr.to_string_lossy(self.py);
                let msg = format!("dict for enum should only contain 1 item: {}", repr);
                return Err(PyErr::new::<exc::ValueError, _>(self.py, msg).into());
            }
            let (key, value) = items.into_iter().next().unwrap();
            let name = seed.deserialize(&mut Deserializer::new(self.py, key))?;
            Ok((name, Deserializer::new(self.py, value)))
        }
    }
}

impl<'de, 'gil> de::VariantAccess<'de> for Deserializer<'gil> {
    type Error = Error;

    fn unit_variant(self) -> Result<()> {
        self.extract::<PyNone>()?;
        Ok(())
    }

    fn newtype_variant_seed<T>(mut self, seed: T) -> Result<T::Value>
    where
        T: de::DeserializeSeed<'de>,
    {
        seed.deserialize(&mut self)
    }

    fn tuple_variant<V>(mut self, _len: usize, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_seq(&mut self, visitor)
    }

    fn struct_variant<V>(mut self, _fields: &'static [&'static str], visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_map(&mut self, visitor)
    }
}


/*
 * 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.
 */

//! Integrate cpython with anyhow

use std::collections::BTreeMap;
use std::fmt;

pub use anyhow::Error;
pub use anyhow::Result;
use cpython::exc;
use cpython::FromPyObject;
use cpython::GILGuard;
use cpython::ObjectProtocol;
use cpython::PyClone;
use cpython::PyList;
use cpython::PyModule;
use cpython::PyResult;
use cpython::Python;
use once_cell::sync::Lazy;
use parking_lot::Mutex;

/// Extends the `Result` type to allow conversion to `PyResult` from a native
/// Rust result.
///
/// If the error is created via [`AnyhowResultExt`], the original Python error
/// will be returned.
///
/// # Examples
///
/// ```
/// use anyhow::{format_err, Error};
/// use cpython::{exc, Python, PyResult};
/// use cpython_ext::error::ResultPyErrExt;
///
/// fn fail_if_negative(i: i32) -> Result<i32, Error> {
///    if (i >= 0) {
///       Ok(i)
///    } else {
///       Err(format_err!("{} is negative", i))
///    }
/// }
///
/// fn py_fail_if_negative(py: Python, i: i32) -> PyResult<i32> {
///    fail_if_negative(i).map_pyerr(py)
/// }
/// ```
pub trait ResultPyErrExt<T> {
    fn map_pyerr(self, py: Python<'_>) -> PyResult<T>;
}

/// Extends the `PyResult` type to allow conversion to `Result`.
///
/// # Examples
///
/// ```
/// use anyhow::Result;
/// use cpython_ext::error::{AnyhowResultExt, PyErr};
///
/// fn eval_py() -> Result<i32> {
///     let gil = cpython::Python::acquire_gil();
///     let py = gil.python();
///     let obj = py.eval("1 + 2", None, None).into_anyhow_result()?;
///     obj.extract(py).into_anyhow_result()
/// }
///
/// fn round_trip() -> cpython::PyResult<()> {
///     let gil = cpython::Python::acquire_gil();
///     let py = gil.python();
///     let res = py.eval("1 + 2", None, None).into_anyhow_result();
///     use cpython_ext::error::ResultPyErrExt;
///     res.map(|_| ()).map_pyerr(py)
/// }
/// ```
///
pub trait AnyhowResultExt<T> {
    fn into_anyhow_result(self) -> Result<T>;
}

pub type AnyhowErrorIntoPyErrFunc = fn(Python, &anyhow::Error) -> Option<cpython::PyErr>;

static INTO_PYERR_FUNC_LIST: Lazy<Mutex<BTreeMap<&'static str, AnyhowErrorIntoPyErrFunc>>> =
    Lazy::new(|| Default::default());

/// Register a function to convert [`anyhow::Error`] to [`PyErr`].
/// For multiple functions, those with smaller name are executed first.
/// Registering a function with an existing name will override that function.
///
/// This affects users of `map_pyerr`.
pub fn register(name: &'static str, func: AnyhowErrorIntoPyErrFunc) {
    let mut list = INTO_PYERR_FUNC_LIST.lock();
    list.insert(name, func);
}

impl<T, E: Into<Error>> ResultPyErrExt<T> for Result<T, E> {
    fn map_pyerr(self, py: Python<'_>) -> PyResult<T> {
        self.map_err(|e| {
            let e: anyhow::Error = e.into();

            if let Some(e) = e.downcast_ref::<PyErr>() {
                return e.inner.clone_ref(py);
            }

            for func in INTO_PYERR_FUNC_LIST.lock().values() {
                if let Some(err) = (func)(py, &e) {
                    return err;
                }
            }

            // Nothing matches. Fallback to RuntimeError.
            // Hopefully this is not really used.
            return cpython::PyErr::new::<exc::RuntimeError, _>(py, format!("{:?}", e));
        })
    }
}

pub fn translate_io_error(py: Python, e: &std::io::Error) -> cpython::PyErr {
    let errno = {
        // Attempt to infer the errno from error type.
        //
        // This is useful for Python to test error type for
        // io::Error generated by non-syscalls.  For
        // example, pipe uses userspace pipe implementation
        // and can return BrokenPipe without errno.
        //
        // See rust's decode_error_kind for the reversed
        // direction.
        use std::io::ErrorKind::*;
        match e.kind() {
            NotFound => Some(libc::ENOENT),
            PermissionDenied => Some(libc::EPERM),
            ConnectionRefused => Some(libc::ECONNREFUSED),
            ConnectionReset => Some(libc::ECONNRESET),
            ConnectionAborted => Some(libc::ECONNABORTED),
            NotConnected => Some(libc::ENOTCONN),
            AddrInUse => Some(libc::EADDRINUSE),
            AddrNotAvailable => Some(libc::EADDRNOTAVAIL),
            BrokenPipe => Some(libc::EPIPE),
            AlreadyExists => Some(libc::EEXIST),
            WouldBlock => Some(libc::EWOULDBLOCK),
            InvalidInput => Some(libc::EINVAL),
            TimedOut => Some(libc::ETIMEDOUT),
            Interrupted => Some(libc::EINTR),
            // No clues about the errno based on kind. Expose the raw OS error code.
            _ => e.raw_os_error(),
        }
    };
    cpython::PyErr::new::<exc::IOError, _>(py, (errno, e.to_string()))
}

impl<T> AnyhowResultExt<T> for PyResult<T> {
    fn into_anyhow_result(self) -> Result<T> {
        self.map_err(|e| Error::new(PyErr::from(e)))
    }
}

/// An enhanced version of `PyErr` that implements [`std::error::Error`].
pub struct PyErr {
    inner: cpython::PyErr,
}

impl PyErr {
    pub fn clone(&self, py: Python) -> Self {
        self.inner.clone_ref(py).into()
    }
}

impl From<cpython::PyErr> for PyErr {
    fn from(e: cpython::PyErr) -> PyErr {
        PyErr { inner: e }
    }
}

impl From<PyErr> for cpython::PyErr {
    fn from(e: PyErr) -> cpython::PyErr {
        e.inner
    }
}

impl fmt::Display for PyErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // To read the python error info, we need the gil. It's not safe to blindly acquire the
        // gil, since it may be held by another thread and we could deadlock if other Rust locks
        // are taken and the lock ordering is inconsistent. So let's only produce a meaningful
        // message if we have the gil already.
        // This mostly affects tracing collectors, since they may choose to format the error at any
        // point.
        if GILGuard::check() {
            let gil = Python::acquire_gil();
            let py = gil.python();
            let repr = self
                .inner
                .pvalue
                .as_ref()
                .unwrap_or_else(|| &self.inner.ptype)
                .repr(py)
                .map(|s| s.to_string_lossy(py).to_string())
                .unwrap_or_else(|_| "<error in repr>".into());
            write!(f, "{}", repr)?;
            if std::env::var("RUST_BACKTRACE").is_ok() {
                if let Ok(s) = format_py_error(py, &self.inner) {
                    write!(f, "\n{}", s)?;
                }
            }
        } else {
            write!(f, "<unable to print PyErr - GIL is not held>")?;
        }
        Ok(())
    }
}

impl fmt::Debug for PyErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl std::error::Error for PyErr {}

impl serde::ser::Error for PyErr {
    fn custom<T: std::fmt::Display>(msg: T) -> Self {
        // We should generally not be acquiring the gil here, since we may encounter lock-order
        // races with other Rust locks. But in this case we know serde logic is being invoked from
        // the cpython-ext Serializer and Deserializer classes, which ensure the GIL is held
        // already.
        let gil = Python::acquire_gil();
        let py = gil.python();
        let err = cpython::PyErr::new::<cpython::exc::TypeError, _>(py, msg.to_string());
        Self { inner: err }
    }
}

impl serde::de::Error for PyErr {
    fn custom<T: std::fmt::Display>(msg: T) -> Self {
        // See comment in ser implementation above for info about acquiring the gil here.
        let gil = Python::acquire_gil();
        let py = gil.python();
        let err = cpython::PyErr::new::<cpython::exc::TypeError, _>(py, msg.to_string());
        Self { inner: err }
    }
}

impl PyErr {
    pub fn into_inner(self) -> cpython::PyErr {
        self.inner
    }
}

pub fn format_py_error(py: Python, err: &cpython::PyErr) -> PyResult<String> {
    let traceback = PyModule::import(py, "traceback")?;
    let py_message = traceback.call(
        py,
        "format_exception",
        (&err.ptype, &err.pvalue, &err.ptraceback),
        None,
    )?;

    let py_lines = PyList::extract(py, &py_message)?;

    let lines: Vec<String> = py_lines
        .iter(py)
        .map(|l| l.extract::<String>(py).unwrap_or_default())
        .collect();

    Ok(lines.join(""))
}

#[cfg(test)]
mod tests {
    use anyhow::Context;

    use super::*;

    #[test]
    fn test_dont_lose_anyhow_context() {
        let gil = Python::acquire_gil();
        let py = gil.python();

        let err: Result<()> = Err(anyhow::anyhow!("first")).context("second");

        assert_eq!(
            format!("{}", err.map_pyerr(py).unwrap_err().pvalue.unwrap()),
            "second\n\nCaused by:\n    first"
        );
    }
}


/*
 * 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::Python;

/// Trait for extracting a Rust object from a Python wrapper class.
///
/// A common pattern when writing bindings for Rust objects is to define
/// a Python wrapper class using the `py_class!` macro, with the underlying
/// Rust object stored as a data field within the Python object.
///
/// When Rust code interacts with a Python wrapper, it may want to work
/// with the underlying Rust object directly. This trait provides a means
/// to do so. Note that the `extract_inner` methods takes `&self`, meaning
/// the inner Rust value cannot be moved out of the wrapper. As a result,
/// the inner value will typically be wrapped in something like an `Arc`.
pub trait ExtractInner {
    type Inner;

    fn extract_inner(&self, py: Python) -> Self::Inner;
}

/// Similar to `ExtractInner`, but returns a reference to the wrapped
/// Rust object. Types that implement this trait will automatically
/// implement `ExtractInner` if the inner type implements `Clone`.
pub trait ExtractInnerRef {
    type Inner;

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

impl<T> ExtractInner for T
where
    T: ExtractInnerRef,
    T::Inner: Clone + 'static,
{
    type Inner = <T as ExtractInnerRef>::Inner;

    fn extract_inner(&self, py: Python) -> Self::Inner {
        self.extract_inner_ref(py).clone()
    }
}


/*
 * 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::io;

use cpython::*;

use crate::error::ResultPyErrExt;
use crate::none::PyNone;

/// Wrap a PyObject into a Rust object that implements Rust Read / Write traits.
pub fn wrap_pyio(pyio: PyObject) -> WrappedIO {
    WrappedIO(pyio)
}

pub struct WrappedIO(pub PyObject);

impl io::Read for WrappedIO {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let bytes = self
            .0
            .call_method(py, "read", (buf.len(),), None)
            .map_err(convert_ioerr)?
            .extract::<PyBytes>(py)
            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "read got non-bytes type"))?;
        let bytes = bytes.data(py);
        if bytes.len() > buf.len() {
            let msg = format!("read got {} bytes, expected {}", bytes.len(), buf.len());
            Err(io::Error::new(io::ErrorKind::InvalidData, msg))
        } else {
            buf[..bytes.len()].copy_from_slice(bytes);
            Ok(bytes.len())
        }
    }
}

impl io::Write for WrappedIO {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let len = buf.len();
        let buf = PyBytes::new(py, buf);
        self.0
            .call_method(py, "write", (buf,), None)
            .map_err(convert_ioerr)?;
        Ok(len)
    }

    fn flush(&mut self) -> io::Result<()> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        self.0
            .call_method(py, "flush", NoArgs, None)
            .map_err(convert_ioerr)?;
        Ok(())
    }
}

impl ::io::IsTty for WrappedIO {
    fn is_tty(&self) -> bool {
        (|| -> PyResult<bool> {
            let gil = Python::acquire_gil();
            let py = gil.python();
            let result = self.0.call_method(py, "isatty", NoArgs, None)?;
            result.extract(py)
        })()
        .unwrap_or(false)
    }
}

/// Convert a Python `IOError` to Rust `io::Error`.
fn convert_ioerr(mut pyerr: PyErr) -> io::Error {
    let gil = Python::acquire_gil();
    let py = gil.python();
    let pyerr = pyerr.instance(py);
    if let Ok(errno) = pyerr
        .getattr(py, "errno")
        .and_then(|e| e.extract::<i32>(py))
    {
        io::Error::from_raw_os_error(errno)
    } else {
        let msg = match pyerr.repr(py) {
            Ok(s) => s.to_string_lossy(py).into_owned(),
            Err(_) => "PythonError: unknown".to_string(),
        };
        io::Error::new(io::ErrorKind::Other, msg)
    }
}

py_class!(pub class PyRustWrite |py| {
    data io: RefCell<Box<dyn ::io::Write + Send>>;

    def write(&self, bytes: PyBytes) -> PyResult<usize> {
        let mut io = self.io(py).borrow_mut();
        let bytes = bytes.data(py);
        io.write_all(bytes).map_pyerr(py)?;
        Ok(bytes.len())
    }

    def flush(&self) -> PyResult<PyNone> {
        let mut io = self.io(py).borrow_mut();
        io.flush().map_pyerr(py)?;
        Ok(PyNone)
    }

    def isatty(&self) -> PyResult<bool> {
        let io = self.io(py).borrow();
        Ok(io.is_tty())
    }

    def isstdin(&self) -> PyResult<bool> {
        let io = self.io(py).borrow();
        Ok(io.is_stdin())
    }

    def isstdout(&self) -> PyResult<bool> {
        let io = self.io(py).borrow();
        Ok(io.is_stdout())
    }

    def isstderr(&self) -> PyResult<bool> {
        let io = self.io(py).borrow();
        Ok(io.is_stderr())
    }

    def close(&self) -> PyResult<PyNone> {
        let mut io = self.io(py).borrow_mut();
        io.flush().map_pyerr(py)?;
        *io = Box::new(ClosedIO);
        Ok(PyNone)
    }
});

struct ClosedIO;

impl io::Write for ClosedIO {
    fn write(&mut self, _buf: &[u8]) -> io::Result<usize> {
        Err(io::Error::new(
            io::ErrorKind::NotConnected,
            "stream was closed",
        ))
    }

    fn flush(&mut self) -> io::Result<()> {
        Err(io::Error::new(
            io::ErrorKind::NotConnected,
            "stream was closed",
        ))
    }
}

impl ::io::IsTty for ClosedIO {
    fn is_tty(&self) -> bool {
        false
    }
}

/// Wrap a Rust Write trait object into a Python object.
pub fn wrap_rust_write(py: Python, w: impl ::io::Write + Send + 'static) -> PyResult<PyRustWrite> {
    PyRustWrite::create_instance(py, RefCell::new(Box::new(w)))
}