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

#![allow(non_camel_case_types)]

use std::cell::Cell;
use std::cell::RefCell;
use std::thread::JoinHandle;

use clidispatch::io::IO;
use cpython::*;
use cpython_ext::PyNone;
use cpython_ext::ResultPyErrExt;
use streampager::action::Action;
use streampager::action::ActionSender;
use streampager::bindings::Binding;
use streampager::bindings::Category;
use streampager::bindings::KeyCode;
use streampager::bindings::Keymap;
use streampager::bindings::Modifiers;
use streampager::config::InterfaceMode;
use streampager::control::Change;
use streampager::control::Controller;
use streampager::Pager;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "sptui"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<sptui>(py)?;
    m.add(py, "NONE", Modifiers::NONE.bits())?;
    m.add(py, "SHIFT", Modifiers::SHIFT.bits())?;
    m.add(py, "ALT", Modifiers::ALT.bits())?;
    m.add(py, "CTRL", Modifiers::CTRL.bits())?;
    m.add(py, "SUPER", Modifiers::SUPER.bits())?;
    Ok(m)
}

// From streampager::keymap_file::KeymapFile::parse_keycode
// TODO: export this from streampager
fn parse_keycode(ident: &str) -> Option<KeyCode> {
    use KeyCode::*;
    match ident {
        "Space" => Some(Char(' ')),
        "Cancel" => Some(Cancel),
        "Backspace" => Some(Backspace),
        "Tab" => Some(Tab),
        "Clear" => Some(Clear),
        "Enter" => Some(Enter),
        "Shift" => Some(Shift),
        "Escape" => Some(Escape),
        "Menu" => Some(Menu),
        "LeftMenu" => Some(LeftMenu),
        "RightMenu" => Some(RightMenu),
        "Pause" => Some(Pause),
        "CapsLock" => Some(CapsLock),
        "PageUp" => Some(PageUp),
        "PageDown" => Some(PageDown),
        "End" => Some(End),
        "Home" => Some(Home),
        "LeftArrow" => Some(LeftArrow),
        "RightArrow" => Some(RightArrow),
        "UpArrow" => Some(UpArrow),
        "DownArrow" => Some(DownArrow),
        "Left" => Some(LeftArrow),
        "Right" => Some(RightArrow),
        "Up" => Some(UpArrow),
        "Down" => Some(DownArrow),
        "Select" => Some(Select),
        "Print" => Some(Print),
        "Execute" => Some(Execute),
        "PrintScreen" => Some(PrintScreen),
        "Insert" => Some(Insert),
        "Delete" => Some(Delete),
        "Help" => Some(Help),
        "Applications" => Some(Applications),
        "Sleep" => Some(Sleep),
        "Numpad0" => Some(Numpad0),
        "Numpad1" => Some(Numpad1),
        "Numpad2" => Some(Numpad2),
        "Numpad3" => Some(Numpad3),
        "Numpad4" => Some(Numpad4),
        "Numpad5" => Some(Numpad5),
        "Numpad6" => Some(Numpad6),
        "Numpad7" => Some(Numpad7),
        "Numpad8" => Some(Numpad8),
        "Numpad9" => Some(Numpad9),
        "Multiply" => Some(Multiply),
        "Add" => Some(Add),
        "Separator" => Some(Separator),
        "Subtract" => Some(Subtract),
        "Decimal" => Some(Decimal),
        "Divide" => Some(Divide),
        "NumLock" => Some(NumLock),
        "ScrollLock" => Some(ScrollLock),
        "BrowserBack" => Some(BrowserBack),
        "BrowserForward" => Some(BrowserForward),
        "BrowserRefresh" => Some(BrowserRefresh),
        "BrowserStop" => Some(BrowserStop),
        "BrowserSearch" => Some(BrowserSearch),
        "BrowserFavorites" => Some(BrowserFavorites),
        "BrowserHome" => Some(BrowserHome),
        "VolumeMute" => Some(VolumeMute),
        "VolumeDown" => Some(VolumeDown),
        "VolumeUp" => Some(VolumeUp),
        "MediaNextTrack" => Some(MediaNextTrack),
        "MediaPrevTrack" => Some(MediaPrevTrack),
        "MediaStop" => Some(MediaStop),
        "MediaPlayPause" => Some(MediaPlayPause),
        "ApplicationLeftArrow" => Some(ApplicationLeftArrow),
        "ApplicationRightArrow" => Some(ApplicationRightArrow),
        "ApplicationUpArrow" => Some(ApplicationUpArrow),
        "ApplicationDownArrow" => Some(ApplicationDownArrow),
        other => {
            if other.starts_with('F') && other.chars().skip(1).all(char::is_numeric) {
                let n = other[1..].parse::<u8>().ok()?;
                Some(Function(n))
            } else {
                None
            }
        }
    }
}

fn parse_category(cat: &str) -> Option<Category> {
    match cat {
        "General" => Some(Category::General),
        "Navigation" => Some(Category::Navigation),
        "Presentation" => Some(Category::Presentation),
        "Searching" => Some(Category::Searching),
        _ => None,
    }
}

py_class!(class sptui |py| {
    data controller: RefCell<Controller>;
    data pager_handle: RefCell<Option<JoinHandle<streampager::Result<()>>>>;
    data action_sender: RefCell<ActionSender>;
    data line_count: Cell<usize>;

    // Create a new sptui instance.
    //
    // `bindings` is a list of (`binding`, `keys`) pairs.
    //
    // Each `binding` is one of:
    // * `None`, indicating the keys should be unbound.
    // * A tuple `(category, description, call)` that defines a custom binding.
    // * A string which matches an existing streampager action name.
    // * A tuple `(name, params)` which matches an existing streampager
    //   action name with parameters.
    //
    // Each `keys` is a list of (`mods`, `key`, `hidden`), where:
    // * `mods` is the modifiers (use `NONE`, `SHIFT`, `CTRL`, `ALT` and `SUPER` in
    //   this module for the values, and bitwise-or together for combinations).
    // * `key` is the key (either a single character, or the streampager key
    //   name).
    // * `hidden` is an optional boolean that hides the key from the help screen if
    //   True.
    def __new__(
        _cls,
        title: String,
        bindings: Vec<(PyObject, Vec<PyTuple>)>
    ) -> PyResult<sptui> {
        let controller = Controller::new(title);
        controller.apply_changes(vec![Change::AppendLines {
            contents: vec![]
        }]).unwrap();

        let mut pager = Pager::new_using_stdio().unwrap();

        pager.set_interface_mode(InterfaceMode::FullScreen);
        let file_index = pager.add_controlled_file(&controller).unwrap();

        let mut keymap = Keymap::default();

        for (binding, keys) in bindings {
            let binding = if binding == py.None() {
                None
            } else if let Ok(name) = binding.extract::<String>(py) {
                Some(Binding::parse(name, Vec::new()).map_err(|e| {
                    PyErr::new::<exc::ValueError, _>(py, format!("Invalid binding: {}", e))
                })?)
            } else if let Ok((name, params)) = binding.extract::<(String, Vec<String>)>(py) {
                Some(Binding::parse(name, params).map_err(|e| {
                    PyErr::new::<exc::ValueError, _>(py, format!("Invalid binding: {}", e))
                })?)
            } else if let Ok((cat, desc, callback)) = binding.extract::<(String, String, PyObject)>(py) {
                let cat = parse_category(&cat).ok_or_else(|| {
                    PyErr::new::<exc::ValueError, _>(py, format!("Invalid category '{}'", cat))
                })?;
                Some(Binding::custom(cat, desc, move |f| {
                    if f == file_index {
                        let gil = Python::acquire_gil();
                        callback.call(gil.python(), NoArgs, None).unwrap();
                    }
                }))
            } else {
                return Err(PyErr::new::<exc::ValueError, _>(
                    py, format!("Invalid binding (got {})", binding.get_type(py).name(py))
                ));
            };
            for key in keys {
                let (mods, key, hidden) = {
                    if let Ok((mods, key, hidden)) = key.as_object().extract::<(u8, String, bool)>(py) {
                        (mods, key, hidden)
                    } else if let Ok((mods, key)) = key.as_object().extract::<(u8, String)>(py) {
                        (mods, key, false)
                    } else {
                        return Err(PyErr::new::<exc::ValueError, _>(
                            py, format!("Invalid key (got {})", key.as_object().get_type(py).name(py))
                        ));
                    }
                };
                let mods = Modifiers::from_bits(mods).ok_or_else(|| {
                    PyErr::new::<exc::ValueError, _>(py, format!("Invalid modifiers '0x{:x}'", mods))
                })?;
                let keycode = if key.chars().count() == 1 {
                    KeyCode::Char(key.chars().next().unwrap())
                } else {
                    parse_keycode(&key).ok_or_else(|| {
                        PyErr::new::<exc::ValueError, _>(py, format!("Invalid key '{}'", key))
                    })?
                };
                if hidden {
                    keymap.bind_hidden(mods, keycode, binding.clone());
                } else {
                    keymap.bind(mods, keycode, binding.clone());
                }
            }
        }

        pager.set_keymap(keymap);

        let io = IO::main().unwrap();
        let (prg_read, prg_write) = pipe::pipe();
        pager.set_progress_stream(prg_read);
        io.set_progress_pipe_writer(Some(prg_write)).unwrap();

        let action_sender = pager.action_sender();

        let pager_handle = std::thread::spawn(move || {
            let result = pager.run();
            io.set_progress_pipe_writer(None).unwrap();
            result
        });

        Self::create_instance(
            py,
            RefCell::new(controller),
            RefCell::new(Some(pager_handle)),
            RefCell::new(action_sender),
            Cell::new(0),
        )
    }

    def replace_contents(&self, contents: Vec<Vec<u8>>) -> PyResult<PyNone> {
        let controller = self.controller(py).borrow();
        let lines = self.line_count(py).replace(contents.len());
        let _ = controller.apply_changes(vec![
            Change::ReplaceLines {
                range: 0..lines,
                contents,
            }
        ]);
        Ok(PyNone)
    }

    def wait(&self) -> PyResult<PyNone> {
        if let Some(handle) = self.pager_handle(py).replace(None) {
            py.allow_threads(move || handle.join()).unwrap().map_pyerr(py)?;
        }
        Ok(PyNone)
    }

    def end(&self) -> PyResult<PyNone> {
        let _ = self.action_sender(py).borrow().send(Action::Quit);
        if let Some(handle) = self.pager_handle(py).replace(None) {
            py.allow_threads(move || handle.join()).unwrap().map_pyerr(py)?;
        }
        Ok(PyNone)
    }
});


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

#![allow(non_camel_case_types)]

use ::status::FileStatus;
use ::status::Status;
use ::status::StatusBuilder;
use cpython::*;
use cpython_ext::ExtractInnerRef;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use types::RepoPathBuf;

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

py_class!(pub class status |py| {
    data inner: Status;

    def __new__(
        _cls,
        python_status: PyObject,
    ) -> PyResult<status> {
        let modified = python_status.getattr(py, "modified")?;
        let added = python_status.getattr(py, "added")?;
        let removed = python_status.getattr(py, "removed")?;
        let deleted = python_status.getattr(py, "deleted")?;
        let unknown = python_status.getattr(py, "unknown")?;
        let ignored = python_status.getattr(py, "ignored")?;
        let clean = python_status.getattr(py, "clean")?;

        let builder = StatusBuilder::new()
            .modified(from_python_file_list(py, modified)?)
            .added(from_python_file_list(py, added)?)
            .removed(from_python_file_list(py, removed)?)
            .deleted(from_python_file_list(py, deleted)?)
            .unknown(from_python_file_list(py, unknown)?)
            .ignored(from_python_file_list(py, ignored)?)
            .clean(from_python_file_list(py, clean)?);

        status::create_instance(py, builder.build())
    }

    def __str__(&self) -> PyResult<PyString> {
        Ok(PyString::new(py, &self.inner(py).to_string()))
    }
});

impl ExtractInnerRef for status {
    type Inner = Status;

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

fn from_python_file_list(py: Python, list: PyObject) -> PyResult<Vec<RepoPathBuf>> {
    let list: PyList = list.cast_into(py)?;
    let mut files = vec![];
    for file in list.iter(py) {
        let file: PyString = file.cast_into(py)?;
        let file = file.to_string(py)?;
        let file = RepoPathBuf::from_string(file.to_string()).map_pyerr(py)?;
        files.push(file);
    }
    Ok(files)
}

/// Convert a Rust-native [`Status`] into a Python-native `scmutil.status`.
pub fn to_python_status(py: Python, status: &Status) -> PyResult<PyObject> {
    let modified = PyList::new(py, &[]);
    let added = PyList::new(py, &[]);
    let removed = PyList::new(py, &[]);
    let deleted = PyList::new(py, &[]);
    let unknown = PyList::new(py, &[]);
    let ignored = PyList::new(py, &[]);
    let clean = PyList::new(py, &[]);

    for (file, status) in status.iter() {
        let list = match status {
            FileStatus::Modified => &modified,
            FileStatus::Added => &added,
            FileStatus::Removed => &removed,
            FileStatus::Deleted => &deleted,
            FileStatus::Unknown => &unknown,
            FileStatus::Ignored => &ignored,
            FileStatus::Clean => &clean,
        };
        let pypath: PyPathBuf = file.into();
        list.append(py, pypath.into_py_object(py).into_object());
    }

    let invalid_path: Vec<String> = status
        .invalid_path()
        .iter()
        .map(|p| util::utf8::escape_non_utf8(p))
        .collect();

    // Create the Python-native status object.
    let scmutil_module = py.import("edenscm.scmutil")?;
    let status_class = scmutil_module.get(py, "status")?;
    let lists = (
        modified,
        added,
        removed,
        deleted,
        unknown,
        ignored,
        clean,
        invalid_path.to_py_object(py),
    );
    status_class.call(py, lists, None)
}


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

#![allow(non_camel_case_types)]

use std::cell::Cell;
use std::cell::RefCell;
use std::mem;
use std::sync::Condvar;
use std::sync::Mutex;
use std::thread;
use std::thread::ThreadId;
use std::time::Duration;

use cpython::*;

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

#[derive(Copy, Clone)]
struct Owner {
    thread_id: ThreadId,
    count: usize,
}

// Used to pass MutexGuard into py.allow_threads closure.  "allow_threads"
// requires "Send" to prevent the use of "py". We're not passing "py" and the
// guard is owned by Rust, and the code runs in a same thread. So it is fine.
struct ForceSend<T>(T);
unsafe impl<T> Send for ForceSend<T> {}

fn rust_thread_id() -> ThreadId {
    thread::current().id()
}

impl Owner {
    fn is_none(&self) -> bool {
        self.count == 0
    }

    fn is_some(&self) -> bool {
        !self.is_none()
    }

    fn is_owned(&self) -> bool {
        self.is_some() && self.thread_id == rust_thread_id()
    }

    fn incref(self) -> Self {
        if !self.is_owned() {
            panic!("incref called from wrong thread!");
        }
        Self {
            thread_id: self.thread_id,
            count: self.count + 1,
        }
    }

    fn decref(self) -> Self {
        if !self.is_owned() {
            panic!("decref called from wrong thread!");
        }
        assert!(self.count > 0);
        Self {
            thread_id: self.thread_id,
            count: self.count - 1,
        }
    }

    fn none() -> Self {
        Self {
            thread_id: rust_thread_id(),
            count: 0,
        }
    }

    fn current_thread() -> Self {
        Self {
            thread_id: rust_thread_id(),
            count: 1,
        }
    }
}

// The Condition class can be used as RLock too.
py_class!(class Condition |py| {
    // The order of data fields matters. First declared, first dropped.

    // Wait for "notify"
    data cond_notify: Condvar;

    // Wait for "release" (internal use)
    data cond_release: Condvar;

    // Used to protect the above Condvars
    data mutex_notify: Mutex<()>;
    data mutex_release: Mutex<()>;

    // Thread owner metadata
    data owner: Cell<Owner>;

    def __new__(_cls, lock: Option<PyObject> = None) -> PyResult<PyObject> {
        match lock {
            None => {
                Ok(Condition::create_instance(
                    py,
                    Condvar::new(),
                    Condvar::new(),
                    Mutex::new(()),
                    Mutex::new(()),
                    Cell::new(Owner::none()),
                )?.into_object())
            },
            Some(lock) => {
                // Do not support taking a customized "lock".
                // Fallback to "threading._Condition" in this case.
                let threading = py.import("threading")?;
                threading.call(py, "_Condition", (lock,), None)
            }
        }
    }

    // RLock APIs

    def acquire(&self, blocking: bool = true) -> PyResult<bool> {
        let owner = self.owner(py).get();
        if owner.is_none() {
            self.owner(py).set(Owner::current_thread());
            Ok(true)
        } else if owner.is_owned() {
            let owner = owner.incref();
            self.owner(py).set(owner);
            Ok(true)
        } else {
            if blocking {
                let mutex_release = self.mutex_release(py);
                let cond_release = self.cond_release(py);
                // Blocking. Wait for other threads to "release", or "wait"
                while self.owner(py).get().is_some() {
                    let guard = ForceSend(mutex_release.lock().unwrap());
                    py.allow_threads(|| {
                        let guard = guard; // capture ForceSend into closure
                        let _guard = cond_release.wait(guard.0).unwrap();
                        // Drop _guard to release the lock before acquiring
                        // Python GIL Otherwise this might deadlock with other
                        // threads acquiring mutex_release.
                    });
                    // At this point we don't know whether the lock is free or
                    // not. The above section does not prevent a Python thread
                    // from acquiring the lock again. But we regained GIL so
                    // check it in a loop.
                }
                let old_owner = self.owner(py).replace(Owner::current_thread());
                assert!(old_owner.is_none());
                Ok(true)
            } else {
                Ok(false)
            }
        }
    }

    def release(&self) -> PyResult<Option<bool>> {
        if self._is_owned(py)? {
            let owner = self.owner(py).get().decref();
            self.owner(py).set(owner);
            if owner.is_none() {
                let cond_release = self.cond_release(py);
                let _guard = self.mutex_release(py).lock().unwrap();
                cond_release.notify_one();
            }
            Ok(None)
        } else {
            Err(PyErr::new::<exc::RuntimeError, _>(py, "cannot release un-acquired lock"))
        }
    }

    // Required by RLockTests
    def _is_owned(&self) -> PyResult<bool> {
        let owner = self.owner(py).get();
        Ok(owner.is_owned())
    }

    def __enter__(&self) -> PyResult<bool> {
        self.acquire(py, true)
    }

    def __exit__(&self, _ty: Option<PyType>, _value: PyObject, _traceback: PyObject) -> PyResult<bool> {
        self.release(py).unwrap();

        // Returning False preserves exceptions
        Ok(false)
    }

    // Condition APIs

    def wait(&self, timeout: Option<f64> = None) -> PyResult<Option<bool>> {
        if !self._is_owned(py)? {
            return Err(PyErr::new::<exc::RuntimeError, _>(py, "cannot wait on un-acquired lock"));
        }

        let cond_notify = self.cond_notify(py);
        let cond_release = self.cond_release(py);
        let mutex_notify = self.mutex_notify(py);
        let mutex_release = self.mutex_release(py);

        // Temporarily release the lock so other threads can "acquire" and
        // "notify". Simplified version of "release".
        let owner = self.owner(py).replace(Owner::none());
        {
            let _guard = mutex_release.lock().unwrap();
            cond_release.notify_one();
        }

        {
            let guard = ForceSend(mutex_notify.lock().unwrap());
            py.allow_threads(|| {
                let guard = guard; // capture ForceSend into closure
                // Allow other threads to run "notify", or "acquire". Blocking.
                let _guard = match timeout {
                    None => cond_notify.wait(guard.0).unwrap(),
                    Some(timeout) => {
                        let duration = Duration::from_millis((timeout * 1000.0) as u64);
                        cond_notify.wait_timeout(guard.0, duration).unwrap().0
                    }
                };
            });
        }

        // Need to re-acquire the lock. A simplified version of "acquire".
        while self.owner(py).get().is_some() {
            let guard = ForceSend(mutex_release.lock().unwrap());
            py.allow_threads(|| {
                let guard = guard; // capture ForceSend into closure
                let _guard = cond_release.wait(guard.0).unwrap();
            });
        }

        // Restore owner
        let old_owner = self.owner(py).replace(owner);
        assert!(old_owner.is_none());

        Ok(None)
    }

    def notify(&self, n: usize = 1) -> PyResult<Option<bool>> {
        // Python API requires the lock to be acquired when using "notify",
        // although Rust does not have this restriction.
        if self._is_owned(py)? {
            let cond_notify = self.cond_notify(py);
            // Acquire the mutex. This makes sure "condvar.wait" has released
            // it. Without this, there is a small window between "allow_threads"
            // and "condvar.wait" at which time sending "notify" will be wrong.
            let _guard = self.mutex_notify(py).lock().unwrap();
            for _ in 0..n {
                cond_notify.notify_one();
            }
            Ok(None)
        } else {
            Err(PyErr::new::<exc::RuntimeError, _>(py, "cannot notify on un-acquired lock"))
        }
    }

    def notify_all(&self) -> PyResult<Option<bool>> {
        if self._is_owned(py)? {
            let _guard = self.mutex_notify(py).lock().unwrap();
            self.cond_notify(py).notify_all();
            Ok(None)
        } else {
            Err(PyErr::new::<exc::RuntimeError, _>(py, "cannot notify on un-acquired lock"))
        }
    }

    // Other stuff

    def __repr__(&self) -> PyResult<String> {
        let owner = self.owner(py).get();
        let msg = if owner.is_some() {
            format!("<Condition (owned by {:?}, refcount {})>", owner.thread_id, owner.count)
        } else {
            format!("<Condition (not owned)>")
        };
        Ok(msg)
    }
});

// To workaround Python bug 29988 where "__exit__" can be skipped by Ctrl+C.
// To use this, wrap a [Condition] in: `b = bug29988wrapper(cond)`, and use
// nested `with`: `with b, b, b, b, b: ...` in the Python world.
py_class!(class bug29988wrapper |py| {
    data obj: Condition;
    data entered: Cell<bool>;

    def __new__(_cls, obj: Condition) -> PyResult<bug29988wrapper> {
        bug29988wrapper::create_instance(py, obj, Cell::new(false))
    }

    def __enter__(&self) -> PyResult<PyObject> {
        if !self.entered(py).get() {
            let obj = self.obj(py);
            let _result = obj.acquire(py, true)?;
            self.entered(py).replace(true);
        }
        Ok(py.None())
    }

    def __exit__(&self, _ty: Option<PyType>, _value: PyObject, _traceback: PyObject) -> PyResult<bool> {
        if self.entered(py).get() {
            let obj = self.obj(py);
            obj.release(py)?;
            self.entered(py).replace(false);
        }
        Ok(false)
    }
});

// Reentrant generator. Can be iterated multiple times by using the `iter()` method.
py_class!(class RGenerator |py| {
    // Main generator (non-reentrant).
    //
    // The generator cannot run from multiple threads. It will raise
    // `ValueError: generator already executing` if doing so. One of the goals of
    // `RGenerator` is to use internal locking so `RGenerator` can "run" from
    // multiple threads.
    //
    // Uses `RefCell` instead of `Mutex` so `__traverse__` does not deadlock.
    data iternext: RefCell<Option<PyObject>>;

    // Mutex to ensure there is only one iternext running.
    data iternext_mutex: Mutex<()>;

    // Items produced by iter.
    data iterlist: RefCell<Option<PyList>>;

    // Whether iteration was completed.
    data itercompleted: Cell<bool>;

    def __new__(_cls, gen: PyObject) -> PyResult<Self> {
        Self::init(py, gen)
    }

    /// Obtains an iterator that iterates from the beginning.
    def iter(&self, skip: usize = 0) -> PyResult<RGeneratorIter> {
        RGeneratorIter::create_instance(py, RefCell::new(Some(self.clone_ref(py))), Cell::new(skip))
    }

    /// Iterate to the end of the original generator.
    def itertoend(&self) -> PyResult<usize> {
        if self.itercompleted(py).get() {
            Ok(0)
        } else {
            let iterlist = self.iterlist(py).borrow();
            if let Some(iterlist) = iterlist.as_ref() {
                let iter = self.iter(py, iterlist.len(py))?;
                let iter = ObjectProtocol::iter(iter.as_object(), py)?;
                Ok(iter.count())
            } else {
                Err(unavailable(py))
            }
        }
    }

    def list(&self) -> PyResult<PyList> {
        let iterlist = self.iterlist(py).borrow();
        if let Some(iterlist) = iterlist.as_ref() {
            Ok(iterlist.clone_ref(py))
        } else {
            Err(unavailable(py))
        }
    }

    def completed(&self) -> PyResult<bool> {
        Ok(self.itercompleted(py).get())
    }

    def __traverse__(&self, visit) {
        let iterlist = self.iterlist(py).borrow();
        if let Some(ref obj) = &*iterlist {
            visit.call(obj)?
        }
        drop(iterlist);
        let iternext = self.iternext(py).borrow();
        if let Some(ref obj) = &*iternext {
            visit.call(obj)?
        }
        Ok(())
    }

    def __clear__(&self) {
        let mut list = self.iterlist(py).borrow_mut();
        let obj = mem::replace(&mut *list, None);
        obj.release_ref(py);
        drop(list);
        let mut next = self.iternext(py).borrow_mut();
        let obj = mem::replace(&mut *next, None);
        obj.release_ref(py);
    }
});

impl RGenerator {
    pub(crate) fn init(py: Python, gen: PyObject) -> PyResult<Self> {
        let iter = gen.iter(py)?.into_object();
        let next = match iter.getattr(py, "__next__") {
            Err(_) => iter.getattr(py, "next")?,
            Ok(next) => next,
        };
        Self::create_instance(
            py,
            RefCell::new(Some(next)),
            Mutex::new(()),
            RefCell::new(Some(PyList::new(py, &[]))),
            Cell::new(false),
        )
    }
}

fn unavailable(py: Python) -> PyErr {
    PyErr::new::<exc::ValueError, _>(py, "RGenerator: cannot access fields after garbage collect")
}

py_class!(class RGeneratorIter |py| {
    data rgen: RefCell<Option<RGenerator>>;
    data index: Cell<usize>;

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

    def __next__(&self) -> PyResult<Option<PyObject>> {
        let rgen = self.rgen(py).borrow();
        let rgen = match &*rgen {
            Some(rgen) => rgen,
            None => return Err(unavailable(py)),
        };
        // Ensure that "__next__" is atomic by locking.
        // Cannot rely on Python GIL because iternext.call(py) might release it.
        let mutex = rgen.iternext_mutex(py);
        let locked = mutex.try_lock();
        if let Ok(_) = locked {
            if let Some(next) = &*rgen.iternext(py).borrow() {
                self.next_internal(py, next)
            } else {
                Err(unavailable(py))
            }
        } else {
            // Release Python GIL to give other threads chances to release mutex.
            let _locked = py.allow_threads(|| mutex.lock().unwrap());
            if let Some(next) = &*rgen.iternext(py).borrow() {
                self.next_internal(py, &next)
            } else {
                Err(unavailable(py))
            }
        }
    }

    def __traverse__(&self, visit) {
        let rgen = self.rgen(py).borrow();
        if let Some(ref obj) = &*rgen {
            visit.call(obj)?
        }
        Ok(())
    }

    def __clear__(&self) {
        let mut rgen = self.rgen(py).borrow_mut();
        let obj = mem::replace(&mut *rgen, None);
        obj.release_ref(py);
    }
});

impl RGeneratorIter {
    // The caller should use locking to ensure `iternext` is not being called
    // from another thread.
    fn next_internal(&self, py: Python, iternext: &PyObject) -> PyResult<Option<PyObject>> {
        let rgen = self.rgen(py).borrow();
        let rgen = match &*rgen {
            Some(rgen) => rgen,
            None => return Err(unavailable(py)),
        };
        let index = self.index(py).get();
        let iterlist = rgen.iterlist(py).borrow();
        if let Some(iterlist) = &*iterlist {
            while iterlist.len(py) <= index && !rgen.itercompleted(py).get() {
                match iternext.call(py, NoArgs, None) {
                    Ok(item) => {
                        iterlist.append(py, item);
                    }
                    Err(err) => {
                        // Could be StopIteration.
                        rgen.itercompleted(py).set(true);
                        return Err(err);
                    }
                };
            }

            let result = if iterlist.len(py) > index {
                Some(iterlist.get_item(py, index))
            } else {
                None
            };
            self.index(py).set(index + 1);
            Ok(result)
        } else {
            Err(unavailable(py))
        }
    }
}

#[cfg(test)]
#[cfg(not(fbcode_build))]
mod tests {
    use std::thread;

    use super::*;

    #[test]
    fn test_rgenerator_iter_multi_threads() {
        let rgen = with_py(|py| {
            let gen = py.eval("(i for i in range(1000))", None, None).unwrap();
            RGenerator::init(py, gen).unwrap()
        });
        let mut rgen_list = Vec::new();
        let n = 40;
        with_py(|py| {
            for _ in 0..n {
                rgen_list.push(rgen.clone_ref(py));
            }
        });

        let threads = rgen_list
            .into_iter()
            .map(move |rgen| {
                thread::spawn(move || {
                    let iter: RGeneratorIter =
                        with_py(|py| RGenerator::iter(&rgen, py, 0).unwrap());
                    let mut count = 0;
                    while let Ok(Some(_)) = with_py(|py| iter.__next__(py)) {
                        count += 1;
                    }
                    assert_eq!(count, 1000);
                    let v: Vec<u32> =
                        with_py(|py| rgen.list(py).unwrap().into_object().extract(py).unwrap());
                    assert_eq!(v, (0..1000).collect::<Vec<u32>>());
                    assert!(with_py(|py| rgen.completed(py).unwrap()));
                })
            })
            .collect::<Vec<_>>();
        for t in threads {
            t.join().unwrap();
        }
    }

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

#![allow(non_camel_case_types)]

use std::cell::Cell;
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::ops::DerefMut;
use std::sync::Arc;

use cpython::*;
use cpython_ext::Bytes;
use cpython_ext::PyNone;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;
use lazy_static::lazy_static;
use parking_lot::Mutex;
use python3_sys as ffi;
use tracing::Level;
use tracing_collector::model::Action;
use tracing_collector::model::EspanId;
use tracing_collector::model::TreeSpans;
use tracing_collector::TracingData;
use tracing_runtime_callsite::CallsiteInfo;
use tracing_runtime_callsite::CallsiteKey;
use tracing_runtime_callsite::EventKindType;
use tracing_runtime_callsite::KindType;
use tracing_runtime_callsite::RuntimeCallsite;
use tracing_runtime_callsite::SpanKindType;

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

    m.add_class::<tracingdata>(py)?;
    m.add_class::<meta>(py)?;
    m.add_class::<wrapfunc>(py)?;
    m.add_class::<wrapiter>(py)?;
    m.add(py, "isheaptype", py_fn!(py, is_heap_type(obj: PyType)))?;

    impl_getsetattr::<wrapfunc>(py);
    impl_getsetattr::<wrapiter>(py);

    let singleton = tracingdata::create_instance(py, DATA.clone())?;
    m.add(py, "singleton", singleton)?;

    m.add(py, "LEVEL_TRACE", LEVEL_TRACE)?;
    m.add(py, "LEVEL_DEBUG", LEVEL_DEBUG)?;
    m.add(py, "LEVEL_INFO", LEVEL_INFO)?;
    m.add(py, "LEVEL_WARN", LEVEL_WARN)?;
    m.add(py, "LEVEL_ERROR", LEVEL_ERROR)?;
    m.add_class::<EventCallsite>(py)?;
    m.add_class::<SpanCallsite>(py)?;
    m.add_class::<instrument>(py)?;
    impl_getsetattr::<InstrumentFunction>(py);

    m.add(
        py,
        "updateenvfilter",
        py_fn!(py, updateenvfilter(dirs: &str)),
    )?;

    Ok(m)
}

lazy_static! {
    // This is public so other libraries can replace it.
    pub static ref DATA: Arc<Mutex<TracingData>> = Arc::new(Mutex::new(TracingData::new()));
}

py_class!(class tracingdata |py| {
    data data: Arc<Mutex<TracingData>>;

    def __new__(_cls) -> PyResult<tracingdata> {
        Self::create_instance(py, Arc::new(Mutex::new(tracing_collector::TracingData::new())))
    }

    @staticmethod
    def deserialize(bytes: PyBytes, binary: bool = true) -> PyResult<tracingdata> {
        let bytes = bytes.data(py);
        let data: TracingData = if binary {
            mincode::deserialize(bytes).map_pyerr(py)?
        } else {
            serde_json::from_slice(bytes).map_pyerr(py)?
        };
        Self::create_instance(py, Arc::new(Mutex::new(data)))
    }

    /// Serialize to bytes.
    def serialize(&self, binary: bool = true) -> PyResult<Bytes> {
        let data = self.data(py).lock();
        if binary {
            let bytes = mincode::serialize(data.deref()).map_pyerr(py)?;
            Ok(Bytes::from(bytes))
        } else {
            let json = serde_json::to_string(data.deref()).map_pyerr(py)?;
            Ok(Bytes::from(json))
        }
    }

    /// Add a span. Return SpanId.
    /// The Span is not entered automatically.
    def span(&self, metadata: Vec<(String, String)>) -> PyResult<u64> {
        let metadata: Vec<(&str, &str)> = metadata.iter().map(|(k, v)| (k.as_ref(), v.as_ref())).collect();
        let mut data = self.data(py).lock();
        Ok(data.add_espan(&metadata, None).0)
    }

    /// Edit fields of a previously added span.
    def edit(&self, id: u64,  metadata: Vec<(String, String)>) -> PyResult<PyNone> {
        let mut data = self.data(py).lock();
        data.edit_espan(EspanId(id), metadata);
        Ok(PyNone)
    }

    /// Enter a span.
    def enter(&self, id: u64) -> PyResult<PyNone> {
        let mut data = self.data(py).lock();
        data.add_action(EspanId(id), Action::EnterSpan);
        Ok(PyNone)
    }

    /// Exit a span.
    def exit(&self, id: u64) -> PyResult<PyNone> {
        let mut data = self.data(py).lock();
        data.add_action(EspanId(id), Action::ExitSpan);
        Ok(PyNone)
    }

    /// Add an event.
    def event(&self, metadata: Vec<(String, String)>) -> PyResult<PyNone> {
        let metadata: Vec<(&str, &str)> = metadata.iter().map(|(k, v)| (k.as_ref(), v.as_ref())).collect();
        let mut data = self.data(py).lock();
        let id = data.add_espan(&metadata, None);
        data.add_action(id, Action::Event);
        Ok(PyNone)
    }

    /// Export as Trace Event.
    def traceevent(&self) -> PyResult<PyObject> {
        let data = self.data(py).lock();
        let trace_event = data.trace_event(Default::default());
        cpython_ext::ser::to_object(py, &trace_event)
    }

    /// Export as Trace Event.
    def traceeventjson(&self) -> PyResult<Bytes> {
        let data = self.data(py).lock();
        let mut buf = Vec::new();
        data.write_trace_event_json(&mut buf, Default::default()).map_pyerr(py)?;
        Ok(Bytes::from(buf))
    }

    /// Export as ASCII.
    ///
    /// `minduration` specifies the minimal duration threshold in micro seconds.
    /// The default value is 10000 (10 milliseconds).
    def ascii(&self, minduration: u64 = 10000) -> PyResult<Str> {
        let mut opts = tracing_collector::model::AsciiOptions::default();
        opts.min_duration_micros_to_hide = minduration;
        Ok(self.data(py).lock().ascii(&opts).into())
    }

    /// Get the TreeSpans for all threads.
    ///
    /// Return {(pid, tid): treespans}.
    def treespans(&self) -> PyResult<HashMap<(u64, u64), treespans>> {
        let data = self.data(py).lock();
        let spans = data.tree_spans().into_iter().map(|(k, v)|{
            (k, treespans::create_instance(py, v).expect("create_instance should succeed"))
        }).collect();
        Ok(spans)
    }

    /// Swap with the global singleton.
    def __enter__(&self) -> PyResult<PyNone> {
        self.swap_with_singleton(py);
        Ok(PyNone)
    }

    /// Swap (back) with the global singleton.
    def __exit__(&self, _ty: Option<PyType>, _value: PyObject, _traceback: PyObject) -> PyResult<bool> {
        self.swap_with_singleton(py);
        Ok(false) // Do not suppress exception
    }
});

impl tracingdata {
    /// Swap self with the global singleton.
    fn swap_with_singleton(&self, py: Python) {
        let mut self_data = self.data(py).lock();
        if let Some(mut global_data) = DATA.try_lock() {
            // in this case, self_data == global_data, no need to swap
            std::mem::swap(self_data.deref_mut(), global_data.deref_mut())
        }
    }
}

fn getattr(py: Python, obj: &PyObject, name: &str) -> PyObject {
    obj.getattr(py, name).unwrap_or_else(|_| py.None())
}

fn getattr_opt(py: Python, obj: &PyObject, name: &str) -> Option<PyObject> {
    obj.getattr(py, name).ok()
}

fn tostr(py: Python, obj: PyObject) -> String {
    obj.str(py)
        .map(|s| s.to_string_lossy(py).to_string())
        .unwrap_or_else(|_| "<missing>".to_string())
}

fn tostr_opt(py: Python, obj: PyObject) -> Option<String> {
    obj.str(py)
        .map(|s| Some(s.to_string_lossy(py).to_string()))
        .unwrap_or_else(|_| None)
}

// Decorator to set "meta" attribute.
py_class!(class meta |py| {
    data meta_obj: PyObject;

    def __new__(_cls, *args, **kwargs) -> PyResult<meta> {
        if (args.len(py) == 1 && kwargs.is_none()) || (args.len(py) == 0 && kwargs.is_some()) {
            let meta = match kwargs {
                Some(kwargs) => kwargs.items_list(py).into_object(),
                None => args.get_item(py, 0),
            };
            Self::create_instance(py, meta)
        } else {
            Err(PyErr::new::<exc::TypeError, _>(py, "invalid meta arguments"))
        }
    }

    def __call__(&self, obj: PyObject) -> PyResult<PyObject> {
        obj.setattr(py, "meta", self.meta_obj(py))?;
        Ok(obj)
    }
});

py_class!(class wrapfunc |py| {
    data inner: PyObject;
    data name: String;
    data module: String;
    data lineno: String;
    data meta: Option<PyObject>;
    data is_generator: bool;
    data last_espan_id: Cell<EspanId>;
    data push_callback: Option<PyObject>;
    data pop_callback: Option<PyObject>;

    def __new__(_cls,
        obj: PyObject,
        meta: Option<PyObject> = None,
        classname: Option<String> = None,
        push_callback: Option<PyObject> = None,
        pop_callback: Option<PyObject> = None
    ) -> PyResult<PyObject> {
        Self::new(py, obj, meta, classname, push_callback, pop_callback)
    }

    def __call__(&self, *args, **kwargs) -> PyResult<PyObject> {
        // Attention: make sure DATA.lock() does not overlap with Python
        // operations. "Simple" Python code like `getattr(a, b)` can
        // potentially call DATA.lock() and cause deadlock.

        // Prepare extra (dynamic) metadata.
        // This calls into Python and cannot take DATA.lock().
        let mut extra_meta: Option<Vec<(String, String)>> = None;
        if let Some(meta) = self.meta(py) {
            let meta = if meta.is_callable(py) {
                // meta: (*args, **kwargs) -> [(str, str)]
                meta.call(py, args, kwargs)?
            } else {
                // meta: [(str, str)]
                meta.clone_ref(py)
            };
            if meta == py.None() {
                // Special case: bypass logging.
                return self.inner(py).call(py, args, kwargs);
            }
            let meta = meta.extract::<Vec<(String, String)>>(py)?;
            extra_meta = Some(meta);
        }

        // Enter Span.
        let espan_id = {
            let last_id = self.last_espan_id(py).get();
            let name = self.name(py);
            let module = self.module(py);
            let line = self.lineno(py);
            let basic_meta: [(&str, &str); 3] = [("name", &name), ("module_path", &module), ("line", &line)];

            // Okay to lock - pure Rust code.
            let mut data = DATA.lock();
            let espan_id = match extra_meta {
                None => {
                    // Static metadata. Avoid dynamic allocations. Try reuse Espans.
                    data.add_espan(&basic_meta[..], Some(last_id))
                },
                Some(extra) => {
                    // Dynamic metadata.
                    let meta: Vec<(&str, &str)> = basic_meta
                        .iter()
                        .cloned()
                        .chain(extra.iter().map(|(k, v)| (k.as_ref(), v.as_ref())))
                        .collect();
                    data.add_espan(&meta, Some(last_id))
                }
            };
            data.add_action(espan_id, Action::EnterSpan);
            espan_id
        };
        self.last_espan_id(py).set(espan_id);

        if let Some(callback) = &self.push_callback(py) {
            callback.call(py, (espan_id.0,), None)?;
        }

        // This calls into Python and cannot take DATA.lock().
        let mut result = self.inner(py).call(py, args, kwargs);

        // Wrap generator automatically.
        if *self.is_generator(py) {
            if let Ok(ref obj) = result {
                if let Ok(obj) = wrapiter::new(py, obj.clone_ref(py)) {
                    result = Ok(obj.into_object());
                }
            }
        }

        if let Some(callback) = &self.pop_callback(py) {
            callback.call(py, NoArgs, None)?;
        }

        // Exit Span.
        {
            // Okay to lock - pure Rust code.
            let mut data = DATA.lock();
            data.add_action(espan_id, Action::ExitSpan);
        }
        result
    }

    def spanid(&self) -> PyResult<u64> {
        Ok(self.last_espan_id(py).get().0)
    }
});

impl wrapfunc {
    fn new(
        py: Python,
        obj: PyObject,
        meta: Option<PyObject>,
        class_name: Option<String>,
        push_callback: Option<PyObject>,
        pop_callback: Option<PyObject>,
    ) -> PyResult<PyObject> {
        if let Ok(wrapped) = obj.extract::<wrapfunc>(py) {
            // No need to wrap again.
            return Ok(wrapped.into_object());
        }

        // Static metadata for a function - name, module and line number.
        // To reduce cost of __call__, cache them in Rust native form.
        let code = getattr(py, &obj, "__code__");
        let mut name = tostr(py, getattr(py, &obj, "__name__"));
        let module = tostr(py, getattr(py, &obj, "__module__"));
        let lineno = tostr(py, getattr(py, &code, "co_firstlineno"));

        // If the callsite provides a class name, use it.
        if let Some(class_name) = class_name {
            name = format!("{}.{}", class_name, name);
        }

        // Function wrapping is used a lot in hg extensions (via mercurial.
        // extensions.wrapfunction). Add the module name to make it easier to
        // check what the function really is.
        // For example, `dispatch.*runcommand` might be wrapped by the undo,
        // journal, copytrace, clienttelemetry, sparse extensions.  By showing
        // `journal.runcommand`, `copytrace._runcommand` instead of
        // `runcommand`, it's easier to tell what's going on.
        if let Some(module_last_name) = module.rsplit(".").nth(0) {
            // Only keep the last part of module name. There is limited space
            // in a span, and common prefix like `edenscm` is not
            // very interesting.
            if module_last_name != "<missing>" {
                name = format!("{}.{}", module_last_name, name);
            }
        }

        // `meta` is `[(str, str)]` or `(*args, **kwargs) -> [(str, str)]`
        // to provide dynamic metadata. It's sometimes inconvenient to
        // pass "meta" through `__new__`. So we also check the `meta`
        // attribute set by the `meta` decorator. This allows the
        // following syntax:
        //
        //    @tracing.wrapfunc
        //    @tracing.meta(...)
        //    def f(...):
        //       ...
        let meta = meta.or_else(|| obj.getattr(py, "meta").ok());

        // Pre-calculate whether this function is a generator.
        // See `inspect.isgeneratorfunction` from Python stdlib.
        let mut is_generator = false;
        if let Ok(flags) = getattr(py, &code, "co_flags").extract::<u64>(py) {
            const CO_GENERATOR: u64 = 32;
            if flags & CO_GENERATOR != 0 {
                is_generator = true;
            }
        }

        let wrapped = Self::create_instance(
            py,
            obj,
            name,
            module,
            lineno,
            meta,
            is_generator,
            Cell::new(EspanId(0)),
            push_callback,
            pop_callback,
        )?;
        Ok(wrapped.into_object())
    }
}

py_class!(class wrapiter |py| {
    data inner: PyObject;
    data name: String;
    data last_espan_id: Cell<EspanId>;

    def __new__(_cls, obj: PyObject) -> PyResult<wrapiter> {
        Self::new(py, obj)
    }

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

    def __next__(&self) -> PyResult<Option<PyObject>> {
        // PERF: Could be a bit faster using unsafe `ffi::PyIter_Next` directly.
        let mut iter = PyIterator::from_object(py, self.inner(py).clone_ref(py))?;

        // Enter Span.
        let espan_id = {
            let last_id = self.last_espan_id(py).get();
            let name = self.name(py);
            let meta: [(&str, &str); 2] = [("name", &name), ("cat", "generator")];
            let mut data = DATA.lock();
            let espan_id = data.add_espan(&meta[..], Some(last_id));
            data.add_action(espan_id, Action::EnterSpan);
            espan_id
        };

        // next(generator)
        let result = match iter.next() {
            None => Ok(None),
            Some(Err(err)) => Err(err),
            Some(Ok(obj)) => Ok(Some(obj)),
        };

        // Exit Span.
        {
            let mut data = DATA.lock();
            data.add_action(espan_id, Action::ExitSpan);
        }
        self.last_espan_id(py).set(espan_id);
        result
    }
});

impl wrapiter {
    fn new(py: Python, obj: PyObject) -> PyResult<Self> {
        if let Ok(wrapped) = obj.extract::<Self>(py) {
            return Ok(wrapped);
        }
        let name = format!("{}.next", tostr(py, getattr(py, &obj, "__name__")));
        Self::create_instance(py, obj, name, Cell::new(EspanId(0)))
    }
}

/// Test whether a `type` (aka. `class`) object is a heap type or not.
/// This is not exposed in Python stdlib. But can be useful to check
/// whether `setattr` is supported or not, which decides whether that
/// class can be traced or not.
///
/// Practically, a heap type is usually defined in Python land using
/// the `class` keyword. A non-heap type is usually defined using native
/// languages. A heap type usually has a `__dict__` slot so it can
/// store attributes and support `setattr`.
fn is_heap_type(_py: Python, typeobj: PyType) -> PyResult<bool> {
    let type_ptr: *mut ffi::PyTypeObject = typeobj.as_type_ptr();
    let result = (unsafe { *type_ptr }.tp_flags & ffi::Py_TPFLAGS_HEAPTYPE) != 0;
    Ok(result)
}

/// Add T.__get__ and __set__ so it can proxy attributes
/// like __doc__ to the original object.
fn impl_getsetattr<T: PythonTypeWithInner>(py: Python) {
    // rust-cpython does not provide a safe way to define __get__.
    // So we have to use some unsafe ffi.
    let type_object: PyType = T::type_object(py);
    let type_ptr: *mut ffi::PyTypeObject = type_object.as_type_ptr();

    extern "C" fn getattr<T: PythonTypeWithInner>(
        this: *mut ffi::PyObject,
        name: *mut ffi::PyObject,
    ) -> *mut ffi::PyObject {
        // This function is called by the Python interpreter.
        // So GIL is already held.
        let py = unsafe { Python::assume_gil_acquired() };

        // Convert raw ffi pointer to friendly rust-cpython objects.
        let this = unsafe { PyObject::from_borrowed_ptr(py, this) };
        let this = unsafe { T::unchecked_downcast_from(this) };
        let name = unsafe { PyObject::from_borrowed_ptr(py, name) };

        match this.inner_obj(py).getattr(py, name) {
            Err(err) => {
                err.restore(py);
                std::ptr::null_mut()
            }
            Ok(obj) => obj.steal_ptr(),
        }
    }

    extern "C" fn setattr<T: PythonTypeWithInner>(
        this: *mut ffi::PyObject,
        name: *mut ffi::PyObject,
        value: *mut ffi::PyObject,
    ) -> std::os::raw::c_int {
        let py = unsafe { Python::assume_gil_acquired() };

        // Convert raw ffi pointer to friendly rust-cpython objects.
        let this = unsafe { PyObject::from_borrowed_ptr(py, this) };
        let this = unsafe { T::unchecked_downcast_from(this) };
        let name = unsafe { PyObject::from_borrowed_ptr(py, name) };
        let value = unsafe { PyObject::from_borrowed_ptr(py, value) };

        match this.inner_obj(py).setattr(py, name, value) {
            Err(err) => {
                err.restore(py);
                -1
            }
            Ok(_) => 0,
        }
    }

    // `__get__`, useful to bind `self` to a free function.
    // See https://docs.python.org/2/howto/descriptor.html.
    extern "C" fn descr_get<T: PythonTypeWithInner>(
        this: *mut ffi::PyObject,
        obj: *mut ffi::PyObject,
        typeobj: *mut ffi::PyObject,
    ) -> *mut ffi::PyObject {
        let py = unsafe { Python::assume_gil_acquired() };

        let this = unsafe { PyObject::from_borrowed_ptr(py, this) };
        let this = unsafe { T::unchecked_downcast_from(this) };
        let inner = this.inner_obj(py);

        // Need to call inner->ob_type->tp_descr_get.
        // There does not seem to have a Python API for calling this.
        let inner_type = inner.get_type(py);
        let inner_type_ptr: *mut ffi::PyTypeObject = inner_type.as_type_ptr();
        let inner_descr_get = unsafe { *inner_type_ptr }.tp_descr_get;
        if let Some(descr_get) = inner_descr_get {
            // Delegate to the original descr_get implementation.
            // Most interestingly, `func_descr_get` in `funcobject.c`.
            let result: *mut ffi::PyObject =
                unsafe { descr_get(this.into_object().steal_ptr(), obj, typeobj) };
            result
        } else {
            this.into_object().steal_ptr()
        }
    }

    // Modify the type object to make __get__ working.
    unsafe {
        (*type_ptr).tp_getattro = Some(getattr::<T>);
        (*type_ptr).tp_setattro = Some(setattr::<T>);
        (*type_ptr).tp_descr_get = Some(descr_get::<T>);
    }
}

trait PythonTypeWithInner: PythonObjectWithTypeObject + PythonObject {
    fn inner_obj<'a>(&'a self, _py: Python<'a>) -> &'a PyObject;
}

impl PythonTypeWithInner for wrapfunc {
    fn inner_obj<'a>(&'a self, py: Python<'a>) -> &'a PyObject {
        self.inner(py)
    }
}

impl PythonTypeWithInner for wrapiter {
    fn inner_obj<'a>(&'a self, py: Python<'a>) -> &'a PyObject {
        self.inner(py)
    }
}

py_class!(pub class treespans |py| {
    data spans: TreeSpans<String>;

    /// Convert into plain Python objects:
    ///
    /// [{"start": micros, "duration": micros | null, "children": [index], **meta}]
    def flatten(&self) -> PyResult<PyObject> {
        cpython_ext::ser::to_object(py, self.spans(py))
    }

    /// Get a "flat" list of spans by name.
    def byname(&self, name: String) -> PyResult<PyObject> {
        let spans = self.spans(py).walk().filter(move |walker, span| {
            if span.meta.get("name") == Some(&name) {
                walker.step_out();
                true
            } else {
                false
            }
        }).collect::<Vec<_>>();
        cpython_ext::ser::to_object(py, &spans)
    }
});

py_class!(pub class TracingSpan |py| {
    // The fields of a struct are dropped in declaration order.
    // "entered" is borrowed from "span". Declare it first.
    data entered: RefCell<Option<UnsafePyLeaked<tracing::span::Entered<'static>>>>;
    @shared data span: tracing::Span;

    def __enter__(&self) -> PyResult<Self> {
        let mut entered = self.entered(py).borrow_mut();
        if entered.is_some() {
            Err(PyErr::new::<exc::ValueError, _>(py, "tracing span was already entered"))
        } else {
            // safety: "entered" lifetime won't exceed its parent (Span).
            *entered = Some(unsafe {
                self.span(py).leak_immutable().map(py, |s| { s.enter() } )
            });
            Ok(self.clone_ref(py))
        }
    }

    def __exit__(&self, _ty: Option<PyType>, _value: PyObject, _traceback: PyObject) -> PyResult<bool> {
        let mut entered = self.entered(py).borrow_mut();
        *entered = None;
        Ok(false) // Do not suppress exception
    }

    /// Record a value. The field name must be predefined with the callsite.
    def record(&self, name: &str, value: FieldValue) -> PyResult<PyNone> {
        if let Some(value) = value.to_opt_tracing_value() {
            self.span(py).borrow().record(name, &value.as_ref());
        }
        Ok(PyNone)
    }

    /// Returns true if this span was disabled by the subscriber.
    def is_disabled(&self) -> PyResult<bool> {
        Ok(self.span(py).borrow().is_disabled())
    }

    /// Returns this span's Id, if it is enabled.
    def id(&self) -> PyResult<Option<u64>> {
        Ok(self.span(py).borrow().id().map(|i| i.into_u64()))
    }
});

py_class!(pub class SpanCallsite |py| {
    data inner: &'static RuntimeCallsite<SpanKindType>;

    // Create a `Callsite` for span in Rust tracing eco-system.
    def __new__(
        _cls,
        obj: PyObject, /* func, or str */
        target: Option<String> = None,
        name: Option<String> = None,
        level: usize = LEVEL_INFO,
        file: Option<String> = None,
        line: Option<usize> = None,
        module: Option<String> = None,
        fieldnames: Option<Vec<String>> = None,
    ) -> PyResult<Self> {
        let callsite = new_callsite(py, obj, target, name, level, file, line, module, fieldnames)?;
        Self::create_instance(py, callsite)
    }

    /// Create a span with the given field values.
    def span(&self, values: Option<Vec<FieldValue>> = None) -> PyResult<TracingSpan> {
        // Convert values to Option<Box<dyn tracing::Value>>.
        let values: Vec<Option<Box<dyn tracing::Value>>> = match values.as_ref() {
            None => Vec::new(),
            Some(list) => list.iter().map(|v| v.to_opt_tracing_value()).collect(),
        };
        let span = self.inner(py).create_span(&values);
        TracingSpan::create_instance(py, Default::default(), span)
    }

    /// Check if this callsite is enabled for logging.
    def isenabled(&self) -> PyResult<bool> {
        Ok(self.inner(py).is_enabled())
    }
});

py_class!(pub class EventCallsite |py| {
    data inner: &'static RuntimeCallsite<EventKindType>;

    // Create a `Callsite` for event in Rust tracing eco-system.
    def __new__(
        _cls,
        obj: PyObject, /* id(obj) for identity */
        target: Option<String> = None,
        name: Option<String> = None,
        level: usize = LEVEL_INFO,
        file: Option<String> = None,
        line: Option<usize> = None,
        module: Option<String> = None,
        fieldnames: Option<Vec<String>> = None,
    ) -> PyResult<Self> {
        let callsite = new_callsite(py, obj, target, name, level, file, line, module, fieldnames)?;
        Self::create_instance(py, callsite)
    }

    /// Create a event with the given field values.
    def event(&self, values: Option<Vec<FieldValue>> = None) -> PyResult<PyNone> {
        // Convert values to Option<Box<dyn tracing::Value>>.
        let values: Vec<Option<Box<dyn tracing::Value>>> = match values.as_ref() {
            None => Vec::new(),
            Some(list) => list.iter().map(|v| v.to_opt_tracing_value()).collect(),
        };
        self.inner(py).create_event(&values);
        Ok(PyNone)
    }

    /// Check if this callsite is enabled for logging.
    def isenabled(&self) -> PyResult<bool> {
        Ok(self.inner(py).is_enabled())
    }
});

/// Create a runtime callsite.
///
/// `obj` should be a Python function or frame so we can figure out the right
/// "callsite" identity and fill up information like module, file, line, etc.
fn new_callsite<K: KindType>(
    py: Python,
    obj: PyObject, /* func, or str */
    target: Option<String>,
    name: Option<String>,
    level: usize,
    file: Option<String>,
    line: Option<usize>,
    module: Option<String>,
    fieldnames: Option<Vec<String>>,
) -> PyResult<&'static RuntimeCallsite<K>> {
    enum ObjType {
        Frame,
        Func,
    }

    // Extract the "code" object and "line". Support frame and function object.
    let (code, line, obj_type) = if let Some(code) = getattr_opt(py, &obj, "f_code") {
        // obj is a frame object (f_code, f_lineno, f_globals, etc.)
        let line = line.or_else(|| getattr(py, &obj, "f_lineno").extract::<usize>(py).ok());
        (code, line, ObjType::Frame)
    } else if let Some(code) = getattr_opt(py, &obj, "__code__") {
        // obj is a function object (__code__, __name__, __module__, etc.)
        let line = line.or_else(|| {
            getattr(py, &obj, "co_firstlineno")
                .extract::<usize>(py)
                .ok()
        });
        (code, line, ObjType::Func)
    } else {
        return Err(PyErr::new::<exc::TypeError, _>(
            py,
            "callsite: expected frame or function object",
        ));
    };
    let id: CallsiteKey = (code.as_ptr() as usize, line.unwrap_or_default());

    let callsite = tracing_runtime_callsite::create_callsite::<K, _>(id, || {
        // Populate other fields: name, module, file.
        let (mut name, module) = match obj_type {
            ObjType::Frame => {
                let name = name.or_else(|| tostr_opt(py, getattr(py, &code, "co_name")));
                let module = module.or_else(|| {
                    getattr_opt(py, &obj, "f_globals").and_then(|g| {
                        g.get_item(py, "__name__")
                            .ok()
                            .and_then(|n| n.extract::<String>(py).ok())
                    })
                });
                (name, module)
            }
            ObjType::Func => {
                let name = name.or_else(|| tostr_opt(py, getattr(py, &obj, "__name__")));
                let module = module.or_else(|| tostr_opt(py, getattr(py, &obj, "__module__")));
                (name, module)
            }
        };

        // Discard the "<lambda>" name. It's pointless.
        if name.as_deref() == Some("<lambda>") {
            name = None;
        };

        // code object provides file name.
        let file = file.or_else(|| tostr_opt(py, getattr(py, &code, "co_filename")));

        CallsiteInfo {
            name: name.unwrap_or_default(),
            // Rewrite Python module "foo.bar" to Rust form "foo::bar". The `.` is not supported
            // by env filter syntax.
            target: target
                .unwrap_or_else(|| module.as_deref().unwrap_or_default().replace('.', "::")),
            level: usize_to_level(level),
            file,
            line: line.map(|l| l as u32),
            module_path: module,
            field_names: fieldnames.unwrap_or_default(),
        }
    });
    Ok(callsite)
}

// Instrument decorator. Construct a callable that accepts a function
// and returns InstrumentFunction
py_class!(class instrument |py| {
    data name: Option<String>;
    data target: Option<String>;
    data level: usize;
    data skip: Option<Vec<String>>;
    data meta: Option<Vec<(String, FieldValue)>>;

    def __new__(_cls, *args, **kwargs) -> PyResult<PyObject> {
        let decorator = {
            let get = |name| -> Option<PyObject> { kwargs?.get_item(py, name) };
            let get_str = |name| -> Option<String> { get(name)?.extract(py).ok() };
            let name: Option<String> = get_str("name");
            let target: Option<String> = get_str("target");
            let level: usize = get("level").and_then(|l| l.extract(py).ok()).unwrap_or(LEVEL_INFO);
            let skip: Option<Vec<String>> = get("skip").and_then(|l| l.extract(py).ok());
            let meta: Option<Vec<(String, FieldValue)>> = kwargs.map(|kwargs| {
                // The rest of the kwargs.
                kwargs.items(py).into_iter().filter_map(|(name, value)| -> Option<_> {
                    let name = name.extract::<String>(py).ok()?;
                    if name == "name" || name == "target" || name == "level" || name == "skip" {
                        None
                    } else {
                        let value = value.extract::<FieldValue>(py).ok()?;
                        Some((name, value))
                    }
                }).collect()
            });
            Self::create_instance(py, name, target, level, skip, meta)?
        };

        match args.len(py) {
            0 => Ok(decorator.into_object()),
            1 => {
                // Single function: decoratoring function
                let func = args.get_item(py, 0);
                let decorated = decorator.__call__(py, func)?;
                Ok(decorated.into_object())
            }
            _ => {
                Err(PyErr::new::<exc::TypeError, _>(py, "instrument expect 0 or 1 args"))
            }
        }
    }

    // Decorate a function. Analyze its arguments.
    def __call__(&self, func: PyObject) -> PyResult<InstrumentFunction> {
        // See inspect.getargs for how to extract arguments.
        let code = func.getattr(py, "__code__")?;

        // Predefined fields (name, value). Includes instrumented args and meta.
        let fields: Vec<(String, InstrumentFieldValue)> = {
            // "static" normal parameters (no *args, **kwds) that are also not skipped.
            // "names" include args and kwargs.
            let names: Vec<String> = code.getattr(py, "co_varnames")?.extract(py)?;
            // "count" only includes normal parameters (no args, or kwargs).
            let count: usize = code.getattr(py, "co_argcount")?.extract(py)?;
            let is_skipped: Box<dyn Fn(&String) -> bool> = match self.skip(py) {
                None => Box::new(|_n| false),
                Some(skip) => Box::new(move |n| skip.contains(n)),
            };
            let mut fields: Vec<(String, InstrumentFieldValue)> =
                names.into_iter().enumerate().take(count).filter_map(|(i, n)| {
                    if is_skipped(&n) {
                        None
                    } else {
                        Some((n, InstrumentFieldValue::FunctionArg(i)))
                    }
                }).collect();
            // Also include names from the metadata.
            if let Some(meta) = self.meta(py) {
                for (key, value) in meta.clone() {
                    fields.push((key, InstrumentFieldValue::Value(value)));
                }
            }
            fields
        };

        let field_names: Vec<String> = fields.iter().map(|t| t.0.clone()).collect();
        let callsite = new_callsite::<SpanKindType>(
            py,
            func.clone_ref(py),
            self.target(py).clone(),
            self.name(py).clone(),
            *self.level(py),
            None /* file */,
            None /* line */,
            None /* module */,
            Some(field_names),
        )?;

        let field_values: Vec<InstrumentFieldValue> = fields.into_iter().map(|t| t.1).collect();
        InstrumentFunction::create_instance(py, func, callsite, field_values)
    }
});

// Instrumented (Python) function.
//
// This is similar to wrapfunc with the differences:
// - it writes to tracing eco-system instead of TracingData (more overhead since
//   spans are recreated every time (?), but level filtering works)
// - it supports logging arguments.
// - it does not wrap iterable or returned iterables.
py_class!(class InstrumentFunction |py| {
    data inner: PyObject;
    data callsite: &'static RuntimeCallsite<SpanKindType>;
    data field_values: Vec<InstrumentFieldValue>;

    def __call__(&self, *args, **kwargs) -> PyResult<PyObject> {
        // Prepare a span. To do so we need field values.
        let owned_values: Vec<FieldValue> = self.field_values(py).iter().map(|v| {
             match v {
                InstrumentFieldValue::FunctionArg(i) => {
                    let i: usize = *i;
                    if args.len(py) <= i {
                        FieldValue::None
                    } else {
                        let arg = args.get_item(py, i);
                        arg.extract::<FieldValue>(py).unwrap_or(FieldValue::None)
                    }
                },
                InstrumentFieldValue::Value(v) => v.clone(),
            }
        }).collect();
        let values: Vec<Option<Box<dyn tracing::Value>>> =
            owned_values.iter().map(|v| v.to_opt_tracing_value()).collect();

        let span = self.callsite(py).create_span(&values);
        let _entered = span.enter();

        self.inner(py).call(py, args, kwargs)
    }
});

impl PythonTypeWithInner for InstrumentFunction {
    fn inner_obj<'a>(&'a self, py: Python<'a>) -> &'a PyObject {
        self.inner(py)
    }
}

/// FieldValue for instrumented function.
/// Could be a function parameter, or a predefined FieldValue.
#[derive(Debug)]
enum InstrumentFieldValue {
    FunctionArg(usize),
    Value(FieldValue),
}

/// Represent the value type of a field without references.
/// Conceptually the owned value of `tracing::Value`.
#[derive(Clone, Debug)]
pub enum FieldValue {
    Str(String),
    Int(i64),
    None,
}

impl FieldValue {
    fn to_opt_tracing_value<'a>(&'a self) -> Option<Box<dyn tracing::Value + 'a>> {
        match self {
            FieldValue::Str(s) => Some(Box::new(s.as_str())),
            FieldValue::Int(i) => Some(Box::new(i)),
            FieldValue::None => None,
        }
    }
}

impl<'a> FromPyObject<'a> for FieldValue {
    fn extract(py: Python, obj: &PyObject) -> PyResult<FieldValue> {
        if let Ok(i) = obj.extract::<i64>(py) {
            Ok(FieldValue::Int(i))
        } else if obj == &py.None() {
            Ok(FieldValue::None)
        } else {
            let s = obj.extract::<String>(py)?;
            Ok(FieldValue::Str(s))
        }
    }
}

fn usize_to_level(level: usize) -> Level {
    match level {
        LEVEL_TRACE => Level::TRACE,
        LEVEL_DEBUG => Level::DEBUG,
        LEVEL_INFO => Level::INFO,
        LEVEL_WARN => Level::WARN,
        LEVEL_ERROR => Level::ERROR,
        _ => Level::ERROR,
    }
}

const LEVEL_TRACE: usize = 0;
const LEVEL_DEBUG: usize = 1;
const LEVEL_INFO: usize = 2;
const LEVEL_WARN: usize = 3;
const LEVEL_ERROR: usize = 4;

fn updateenvfilter(py: Python, dirs: &str) -> PyResult<PyNone> {
    tracing_reload::update_env_filter_directives(dirs).map_pyerr(py)?;
    Ok(PyNone)
}


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

#![allow(non_camel_case_types)]

//! Python bindings for treestate.
//!
//! This is a Rust implementation of the dirstate concept for Mercurial, using a tree structure
//! in an append-only storage back-end.
//!
//! The directory state stores information for all files in a working copy that are of interest
//! to Mercurial.  In particular, for each file in the working copy it stores the mode flags,
//! size, and modification time of the file.  These can be compared with current values to
//! determine if the file has changed.
//!
//! The directory state also stores files that are in the working copy parent manifest but have
//! been marked as removed.

use std::sync::Arc;

use ::treestate::filestate::FileStateV2;
use ::treestate::filestate::StateFlags;
use ::treestate::store::BlockId;
use ::treestate::tree::VisitorResult;
use ::treestate::treestate::TreeState;
use anyhow::Error;
use cpython::*;
use cpython_ext::AnyhowResultExt;
use cpython_ext::PyPath;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use parking_lot::Mutex;
use pathmatcher::DirectoryMatch;
use pathmatcher::Matcher;
use pypathmatcher::PythonMatcher;
use types::RepoPathBuf;
use vfs::VFS;

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

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "treestate"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<treestate>(py)?;
    m.add(py, "EXIST_P1", StateFlags::EXIST_P1.to_bits())?;
    m.add(py, "EXIST_P2", StateFlags::EXIST_P2.to_bits())?;
    m.add(py, "EXIST_NEXT", StateFlags::EXIST_NEXT.to_bits())?;
    m.add(py, "IGNORED", StateFlags::IGNORED.to_bits())?;
    m.add(py, "NEED_CHECK", StateFlags::NEED_CHECK.to_bits())?;
    m.add(py, "COPIED", StateFlags::COPIED.to_bits())?;
    m.add(py, "tohgstate", py_fn!(py, flags_to_hg_state(flags: u16)))?;
    Ok(m)
}

impl treestate {
    pub fn get_state(&self, py: Python) -> Arc<Mutex<TreeState>> {
        self.state(py).clone()
    }
}

py_class!(pub class treestate |py| {
    data state: Arc<Mutex<TreeState>>;

    // This should only be used for hg doctor.
    @staticmethod
    def openraw(path: PyPathBuf, root_id: u64) -> PyResult<treestate> {
        let treestate = TreeState::open(
            path.as_path(),
            BlockId(root_id),
            VFS::new(path.as_path().to_path_buf()).map_pyerr(py)?.case_sensitive(),
        ).map_pyerr(py)?;
        Self::create_instance(py, Arc::new(Mutex::new(treestate)))
    }

    // This should only be used for tests.
    @staticmethod
    def new(directory: PyPathBuf) -> PyResult<treestate> {
        let treestate = TreeState::new(
            directory.as_path(),
            VFS::new(directory.as_path().to_path_buf()).map_pyerr(py)?.case_sensitive(),
        ).map_pyerr(py)?.0;
        Self::create_instance(py, Arc::new(Mutex::new(treestate)))
    }

    def flush(&self) -> PyResult<u64> {
        // Save changes to the existing file.
        let mut state = self.state(py).lock();
        let root_id = convert_result(py, state.flush())?;
        Ok(root_id.0)
    }

    def reset(&self, directory: PyPathBuf) -> PyResult<u64> {
        let mut treestate = self.state(py).lock();
        let (new_treestate, root_id) = convert_result(py, TreeState::new(
            directory.as_path(),
            VFS::new(directory.as_path().to_path_buf()).map_pyerr(py)?.case_sensitive(),
        ))?;
        *treestate = new_treestate;
        Ok(root_id.0)
    }

    def filename(&self) -> PyResult<String> {
        convert_result(py, self.state(py).lock().file_name())
    }

    def saveas(&self, directory: &PyPath) -> PyResult<u64> {
        // Save as a new file. Return `BlockId` that can be used in constructor.
        let mut state = self.state(py).lock();
        let root_id = convert_result(py, state.write_new(directory))?;
        Ok(root_id.0)
    }

    def __len__(&self) -> PyResult<usize> {
        Ok(self.state(py).lock().len())
    }

    def __contains__(&self, path: PyPathBuf) -> PyResult<bool> {
        let mut state = self.state(py).lock();
        let file = convert_result(py, state.get(path.as_utf8_bytes()))?;
        // A lot of places require "__contains__(path)" to be "False" if "path" is "?" state
        let visible_flags = StateFlags::EXIST_P1 | StateFlags::EXIST_P2 | StateFlags::EXIST_NEXT;
        Ok(match file {
            Some(file) => file.state.intersects(visible_flags),
            None => false,
        })
    }

    def get(&self, path: &PyPath, default: Option<(u16, u32, i32, i32, Option<PyPathBuf>)>) -> PyResult<Option<(u16, u32, i32, i32, Option<PyPathBuf>)>> {
        let mut state = self.state(py).lock();
        let path = path.as_utf8_bytes();

        assert!(!path.ends_with(b"/"));

        let file = convert_result(py, state.get(path))?;
        Ok(file.map_or(default, |file|
                    Some((file.state.to_bits(),
                     file.mode,
                     file.size,
                     file.mtime,
                     file.copied.as_ref().map(|path| PyPathBuf::from_utf8_bytes(path.to_vec()).unwrap())))))
    }

    def insert(
        &self, path: &PyPath, bits: u16, mode: u32, size: i32, mtime: i32, copied: Option<PyPathBuf>
    ) -> PyResult<PyObject> {
        let mut flags = StateFlags::from_bits_truncate(bits);
        // For special mtime or size, mark them as "NEED_CHECK" automatically.
        if mtime < 0 || size < 0 {
            flags |= StateFlags::NEED_CHECK;
        }

        // Also fix-up COPIED bit so they stay consistent.
        if copied.as_ref().is_some() {
            flags |= StateFlags::COPIED;
        } else {
            flags -= StateFlags::COPIED;
        };

        let file = FileStateV2 { mode, size, mtime, copied: copied.map(|copied| copied.as_utf8_bytes().to_vec().into_boxed_slice()), state: flags };
        let path = path.as_utf8_bytes();
        let mut state = self.state(py).lock();
        convert_result(py, state.insert(path, &file))?;
        Ok(py.None())
    }

    def remove(&self, path: &PyPath) -> PyResult<bool> {
        let mut state = self.state(py).lock();
        convert_result(py, state.remove(path.as_utf8_bytes()))
    }

    def getdir(&self, path: &PyPath) -> PyResult<Option<(u16, u16)>> {
        let mut state = self.state(py).lock();
        let path = path.as_utf8_bytes();

        let dir = convert_result(py, state.get_dir(path))?;
        Ok(dir.map(|state| (state.union.to_bits(), state.intersection.to_bits())))
    }

    def hasdir(&self, path: &PyPath) -> PyResult<bool> {
        let mut state = self.state(py).lock();
        let path = path.as_utf8_bytes();
        Ok(convert_result(py, state.has_dir(path))?)
    }

    def walk(
        &self,
        setbits: u16,
        unsetbits: u16,
        dirfilter: Option<PyObject> = None
    ) -> PyResult<Vec<PyPathBuf>> {
        // Get all file paths with `setbits` all set, `unsetbits` all unset.
        // If dirfilter is provided. It is a callable that takes a directory path, and returns True
        // if the path should be skipped.
        assert_eq!(setbits & unsetbits, 0, "setbits cannot overlap with unsetbits");
        let setbits = StateFlags::from_bits_truncate(setbits);
        let unsetbits = StateFlags::from_bits_truncate(unsetbits);
        let mask = setbits | unsetbits;
        let mut state = self.state(py).lock();
        let mut result = Vec::new();
        convert_result(py, state.visit(
            &mut |components, _state| {
                let path = PyPathBuf::from_utf8_bytes(components.concat()).expect("path should be utf-8");
                result.push(path);
                Ok(VisitorResult::NotChanged)
            },
            &|components, dir| {
                if let Some(state) = dir.get_aggregated_state() {
                    if !state.union.contains(setbits) || state.intersection.intersects(unsetbits) {
                        return false;
                    }
                }
                if let Some(ref dirfilter) = dirfilter {
                    let path = PyPathBuf::from_utf8_bytes(components.concat()).expect("path should be utf-8");
                    if let Ok(result) = dirfilter.call(py, (path,), None) {
                        if let Ok(result) = result.is_true(py) {
                            if result {  // should skip
                                return false;  // do not visit
                            }
                        }
                    }
                }
                true  // do visit
            },
            &|_, file| file.state & mask == setbits,
        ))?;
        Ok(result)
    }

    /// Tracked files filtered by the matcher.
    def matches(&self, matcher: PyObject) -> PyResult<Vec<PyPathBuf>> {
        let matcher = PythonMatcher::new(py, matcher);

        let mut state = self.state(py).lock();
        let mut result = Vec::new();
        let mask = StateFlags::EXIST_P1 | StateFlags::EXIST_P2 | StateFlags::EXIST_NEXT;
        convert_result(py, state.visit(
            &mut |components, _state| {
                let path = RepoPathBuf::from_utf8(components.concat()).expect("path should be utf-8");
                if  matcher.matches_file(&path)? {
                    let path = PyPathBuf::from_string(path.into_string());
                    result.push(path);
                }
                Ok(VisitorResult::NotChanged)
            },
            &|components, dir| {
                if let Some(state) = dir.get_aggregated_state() {
                    if !state.union.intersects(mask) {
                        return false;
                    }
                }
                let mut binary_path = components.concat();
                // Remove the trailing slash.
                assert_eq!(binary_path.pop().unwrap_or(b'/'), b'/');
                let path = RepoPathBuf::from_utf8(binary_path).expect("path should be utf-8");
                match matcher.matches_directory(&path) {
                    Ok(DirectoryMatch::Nothing) => false,  // do not visit
                    _ => true,  // do visit
                }
            },
            &|_, file| file.state.intersects(mask),
        ))?;
        Ok(result)
    }

    def tracked(&self, prefix: &PyPath) -> PyResult<Vec<PyPathBuf>> {
        // prefix limits the result to given prefix (ex. ["dir1/", "dir2/"]). To get all tracked
        // files, set prefix to an empty list.
        // Not ideal as a special case. But the returned list is large and it needs to be fast.
        // It's basically walk(EXIST_P1, 0) + walk(EXIST_P2, 0) + walk(EXIST_NEXT).
        let mut state = self.state(py).lock();
        let mut result = Vec::new();
        let mask = StateFlags::EXIST_P1 | StateFlags::EXIST_P2 | StateFlags::EXIST_NEXT;
        let prefix = split_path(prefix.as_utf8_bytes());
        convert_result(py, state.visit(
            &mut |components, _state| {
                let path = PyPathBuf::from_utf8_bytes(components.concat()).expect("path should be utf-8");
                result.push(path);
                Ok(VisitorResult::NotChanged)
            },
            &|path, dir| {
                if path.iter().zip(prefix.iter()).any(|(a, b)| a != b) {
                    // common prefix does not match
                    false
                } else {
                    match dir.get_aggregated_state() {
                        None => true,
                        Some(state) => state.union.intersects(mask),
                    }
                }
            },
            &|path, file| {
                if path.len() == prefix.len() {
                    // must be an exact match
                    *path == prefix
                } else if path.len() < prefix.len() {
                    // file outside given prefix
                    false
                } else {
                    file.state.intersects(mask)
                }
            }
        ))?;
        Ok(result)
    }

    def getfiltered(
        &self, path: &PyPath, filter: PyObject, filterid: u64
    ) -> PyResult<Vec<PyPathBuf>> {
        let mut state = self.state(py).lock();

        let result = convert_result(py, state.get_filtered_key(
            path.as_utf8_bytes(),
            &mut |path| {
                let path = PyPathBuf::from_utf8_bytes(path.to_vec()).expect("path should be utf-8");
                let filtered = filter
                    .call(py, (&path,), None).into_anyhow_result()?
                    .extract::<PyPathBuf>(py).into_anyhow_result()?;
                Ok(filtered.into_utf8_bytes().into_boxed_slice())
            },
            filterid,
        ))?;

        Ok(result.into_iter().map(|o| PyPathBuf::from_utf8_bytes(o.into_vec()).expect("path should be utf-8")).collect())
    }

    def pathcomplete(
        &self, prefix: &PyPath, setbits: u16, unsetbits: u16, matchcallback: PyObject,
        fullpaths: bool
    ) -> PyResult<PyObject> {
        let setbits = StateFlags::from_bits_truncate(setbits);
        let unsetbits = StateFlags::from_bits_truncate(unsetbits);
        let mask = setbits | unsetbits;
        let mut state = self.state(py).lock();
        let prefix = prefix.as_utf8_bytes();

        convert_result(py, state.path_complete(
            prefix,
            fullpaths,
            &|file| file.state & mask == setbits,
            &mut |components| {
                let path = PyPathBuf::from_utf8_bytes(components.concat()).expect("path should be utf-8");
                matchcallback.call(py, (path,), None).into_anyhow_result()?;
                Ok(())
            },
        ))?;

        Ok(py.None())
    }

    // Import another map of dirstate tuples into this treestate. Note: copymap is not imported.
    def importmap(&self, old_map: PyObject) -> PyResult<Option<PyObject>> {
        let mut tree = self.state(py).lock();
        let items = old_map.call_method(py, "items", NoArgs, None)?;
        let iter = PyIterator::from_object(
            py, items.call_method(py, "__iter__", NoArgs, None)?)?;

        for item in iter {
            let item_tuple = item?.extract::<PyTuple>(py)?;
            let path = item_tuple.get_item(py, 0).extract::<PyPathBuf>(py)?;
            let data = item_tuple.get_item(py, 1).extract::<PySequence>(py)?;
            let state = data.get_item(py, 0)?.extract::<PyString>(py)?.data(py).to_string(py)?.bytes().next().unwrap();
            let mode = data.get_item(py, 1)?.extract::<u32>(py)?;
            let size = data.get_item(py, 2)?.extract::<i32>(py)?;
            let mtime = data.get_item(py, 3)?.extract::<i32>(py)?;
            // Mercurial uses special "size"s to represent "otherparent" if state is "n".
            // See "size = -2" in mercurial/dirstate.py
            let flags = match size {
                -2 => StateFlags::EXIST_P2,
                _ => StateFlags::EXIST_P1,
            };
            let flags = match state {
                b'n' => flags | StateFlags::EXIST_NEXT,
                b'm' => StateFlags::EXIST_P1 | StateFlags::EXIST_P2 | StateFlags::EXIST_NEXT,
                b'r' => flags,
                b'a' => StateFlags::EXIST_NEXT,
                _ => StateFlags::empty(),
            };
            if !flags.is_empty() {
                let file = FileStateV2 { mode, size, mtime, copied: None, state: flags };
                convert_result(py, tree.insert(path.as_utf8_bytes(), &file))?;
            }
        }
        Ok(None)
    }

    def invalidatemtime(&self, fsnow: i32) -> PyResult<PyObject> {
        let mut state = self.state(py).lock();
        convert_result(py, state.invalidate_mtime(fsnow))?;
        Ok(py.None())
    }

    def getmetadata(&self) -> PyResult<PyBytes> {
        let state = self.state(py).lock();
        let metadata = PyBytes::new(py, state.metadata_bytes());
        Ok(metadata)
    }

    def setmetadata(&self, metadata: PyBytes) -> PyResult<PyObject> {
        let mut state = self.state(py).lock();
        let metadata = metadata.data(py);
        state.set_metadata_bytes(metadata);
        Ok(py.None())
    }
});

/// Convert StateFlags to Mercurial dirstate state
fn flags_to_hg_state(_py: Python, flags: u16) -> PyResult<&'static str> {
    let flags = StateFlags::from_bits_truncate(flags);
    Ok(
        match (
            flags.intersects(StateFlags::EXIST_P1 | StateFlags::EXIST_P2),
            flags.contains(StateFlags::EXIST_P1 | StateFlags::EXIST_P2),
            flags.contains(StateFlags::EXIST_NEXT),
        ) {
            (true, true, true) => "m",  // merge
            (true, false, true) => "n", // normal
            (true, _, false) => "r",    // remove
            (false, _, true) => "a",    // add
            (false, _, false) => "?",   // untracked
        },
    )
}

/// Convert a Result to PyResult
fn convert_result<T>(py: Python, result: Result<T>) -> PyResult<T> {
    result.map_pyerr(py)
}

/// Convert "dir1/dir2/file1" to ["dir1/", "dir2/", "file1"]
fn split_path(path: &[u8]) -> Vec<&[u8]> {
    // convert prefix to a vec like ["dir/", "dir2/", "file"]
    if path == b"/" {
        return Vec::new();
    }
    let mut components = Vec::new();
    let mut next_index = 0;
    for (index, byte) in path.iter().enumerate() {
        if *byte == b'/' {
            components.push(&path[next_index..index + 1]);
            next_index = index + 1;
        }
    }
    if next_index < path.len() {
        components.push(&path[next_index..]);
    }
    components
}


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

use cpython::*;
use cpython_ext::ResultPyErrExt;
use vlqencoding::VLQDecode;
use vlqencoding::VLQDecodeAt;
use vlqencoding::VLQEncode;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "vlq"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(
        py,
        "decodeat",
        py_fn!(py, decodeat(data: PyBytes, offset: usize)),
    )?;
    m.add(py, "decode", py_fn!(py, decode(data: PyBytes)))?;
    m.add(py, "encode", py_fn!(py, encode(value: u64)))?;
    m.add(py, "read", py_fn!(py, read(io: PyObject)))?;
    Ok(m)
}

fn decodeat(py: Python, data: PyBytes, offset: usize) -> PyResult<(u64, usize)> {
    let (value, len) = data.data(py).read_vlq_at(offset).map_pyerr(py)?;
    Ok((value, len))
}

fn decode(py: Python, data: PyBytes) -> PyResult<u64> {
    let value = data.data(py).read_vlq().map_pyerr(py)?;
    Ok(value)
}

fn encode(py: Python, value: u64) -> PyResult<PyBytes> {
    let mut buf = Vec::new();
    buf.write_vlq(value).unwrap();
    Ok(PyBytes::new(py, &buf))
}

fn read(py: Python, io: PyObject) -> PyResult<u64> {
    let mut io = cpython_ext::wrap_pyio(io);
    let value = io.read_vlq().map_pyerr(py)?;
    Ok(value)
}