/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fs;
use std::io;
use std::io::BufRead;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io::Write;
use std::path::Path;

use crate::errors::IoResultExt;
use crate::errors::ResultExt;
use crate::lock::ScopedDirLock;
use crate::log::GenericPath;
use crate::log::Log;
use crate::log::LogMetadata;
use crate::log::OpenOptions;
use crate::log::META_FILE;
use crate::log::PRIMARY_FILE;
use crate::log::PRIMARY_HEADER;
use crate::log::PRIMARY_START_OFFSET;
use crate::repair::OpenOptionsOutput;
use crate::repair::OpenOptionsRepair;
use crate::repair::RepairMessage;
use crate::utils;
use crate::utils::mmap_path;

// Repair
impl OpenOptions {
    /// Attempt to repair a broken [`Log`] at the given directory.
    ///
    /// This is done by truncating entries in the primary log, and rebuilding
    /// corrupted indexes.
    ///
    /// Backup files are written for further investigation.
    ///
    /// Return message useful for human consumption.
    pub fn repair(&self, dir: impl Into<GenericPath>) -> crate::Result<String> {
        let dir = dir.into();
        let dir = match dir.as_opt_path() {
            Some(dir) => dir,
            None => return Ok(format!("{:?} is not on disk. Nothing to repair.\n", &dir)),
        };

        let result: crate::Result<_> = (|| {
            if !dir.exists() {
                return Ok(format!("{:?} does not exist. Nothing to repair.\n", dir));
            }

            let lock = ScopedDirLock::new(dir)?;
            let mut message = RepairMessage::new(dir);
            message += &format!("Processing IndexedLog: {:?}\n", dir);

            let primary_path = dir.join(PRIMARY_FILE);
            let meta_path = dir.join(META_FILE);

            // Make sure the header of the primary log file is okay.
            (|| -> crate::Result<()> {
                #[allow(clippy::never_loop)]
                let header_corrupted = loop {
                    if let Err(e) = primary_path.metadata() {
                        if e.kind() == io::ErrorKind::NotFound {
                            break true;
                        }
                    }
                    let mut file = fs::OpenOptions::new()
                        .read(true)
                        .open(&primary_path)
                        .context(&primary_path, "cannot open for read")?;
                    let mut buf = [0; PRIMARY_START_OFFSET as usize];
                    break match file.read_exact(&mut buf) {
                        Ok(_) => buf != PRIMARY_HEADER,
                        Err(_) => true,
                    };
                };
                if header_corrupted {
                    let mut file = fs::OpenOptions::new()
                        .write(true)
                        .create(true)
                        .open(&primary_path)
                        .context(&primary_path, "cannot open for write")?;
                    file.write_all(PRIMARY_HEADER)
                        .context(&primary_path, "cannot re-write header")?;
                    let _ = utils::fix_perm_file(&file, false);
                    message += "Fixed header in log\n";
                }
                Ok(())
            })()
            .context("while making sure log has the right header")?;

            // Make sure the "primary_len" is large enough.
            (|| -> crate::Result<()> {
                let primary_len = primary_path
                    .metadata()
                    .context(&primary_path, "cannot read fs metadata")?
                    .len();
                match LogMetadata::read_file(&meta_path)
                    .context("repair cannot fix metadata corruption")
                {
                    Ok(meta) => {
                        // If metadata can be read, trust it.
                        if meta.primary_len > primary_len {
                            use fs2::FileExt;
                            // Log was truncated for some reason...
                            // (This should be relatively rare)
                            // Fill Log with 0s.
                            let file = fs::OpenOptions::new()
                                .write(true)
                                .open(&primary_path)
                                .context(&primary_path, "cannot open for write")?;
                            file.allocate(meta.primary_len)
                                .context(&primary_path, "cannot fallocate")?;
                            message += &format!(
                                "Extended log to {:?} bytes required by meta\n",
                                meta.primary_len
                            );
                        }
                    }
                    Err(meta_err) => {
                        // Attempt to rebuild metadata.
                        let meta = LogMetadata::new_with_primary_len(primary_len);
                        meta.write_file(&meta_path, self.fsync)
                            .context("while recreating meta")
                            .source(meta_err)?;
                        message += "Rebuilt metadata\n";
                    }
                }
                Ok(())
            })()
            .context("while making sure log.length >= meta.log_length")?;

            // Reload the latest log without indexes.
            //
            // At this time log is likely open-able.
            //
            // Try to open it with indexes so we might reuse them. If that
            // fails, retry with all indexes disabled.
            let mut log = self
                .open_with_lock(&dir.into(), &lock)
                .or_else(|_| {
                    self.clone()
                        .index_defs(Vec::new())
                        .open(GenericPath::from(dir))
                })
                .context("cannot open log for repair")?;

            let mut iter = log.iter();

            // Read entries until hitting a checksum error.
            let mut entry_count = 0;
            while let Some(Ok(_)) = iter.next() {
                entry_count += 1;
            }

            let valid_len = iter.next_offset;
            assert!(valid_len >= PRIMARY_START_OFFSET);
            assert!(valid_len <= log.meta.primary_len);

            if valid_len == log.meta.primary_len {
                message += &format!(
                    "Verified {} entries, {} bytes in log\n",
                    entry_count, valid_len
                );
            } else {
                message += &format!(
                    "Verified first {} entries, {} of {} bytes in log\n",
                    entry_count, valid_len, log.meta.primary_len
                );

                // Backup the part to be truncated.
                (|| -> crate::Result<()> {
                    let mut primary_file = fs::OpenOptions::new()
                        .read(true)
                        .open(&primary_path)
                        .context(&primary_path, "cannot open for read")?;
                    let backup_path = dir.join(format!(
                        "log.bak.epoch{}.offset{}",
                        log.meta.epoch, valid_len
                    ));
                    let mut backup_file = fs::OpenOptions::new()
                        .create_new(true)
                        .write(true)
                        .open(&backup_path)
                        .context(&backup_path, "cannot open")?;

                    primary_file
                        .seek(SeekFrom::Start(valid_len))
                        .context(&primary_path, "cannot seek")?;

                    let mut reader = io::BufReader::new(primary_file);
                    loop {
                        let len = {
                            let buf = reader.fill_buf().context(&primary_path, "cannot read")?;
                            if buf.is_empty() {
                                break;
                            }
                            backup_file
                                .write_all(buf)
                                .context(&backup_path, "cannot write")?;
                            buf.len()
                        };
                        reader.consume(len);
                    }
                    message += &format!("Backed up corrupted log to {:?}\n", backup_path);
                    Ok(())
                })()
                .context("while trying to backup corrupted log")?;

                // Update metadata. Invalidate indexes.
                // Bump epoch since this is a non-append-only change.
                // Reload disk buffer.
                log.meta.primary_len = valid_len;
                log.meta.indexes.clear();
                log.meta.epoch = log.meta.epoch.wrapping_add(1);
                log.disk_buf = mmap_path(&primary_path, valid_len)?;

                log.meta
                    .write_file(&meta_path, log.open_options.fsync)
                    .context("while trying to update metadata with verified log length")?;
                message += &format!("Reset log size to {}\n", valid_len);
            }

            // Also rebuild corrupted indexes.
            // Without this, indexes are empty until the next `sync`, which
            // can lead to bad performance.
            log.open_options.index_defs = self.index_defs.clone();
            message += &log
                .rebuild_indexes_with_lock(false, &lock)
                .context("while trying to update indexes with reapired log")?;

            Ok(message.into_string())
        })();

        result.context(|| format!("in log::OpenOptions::repair({:?})", dir))
    }
}

impl OpenOptionsRepair for OpenOptions {
    fn open_options_repair(&self, dir: impl AsRef<Path>) -> crate::Result<String> {
        OpenOptions::repair(self, dir.as_ref())
    }
}

impl OpenOptionsOutput for OpenOptions {
    type Output = Log;

    fn open_path(&self, path: &Path) -> crate::Result<Self::Output> {
        self.open(path)
    }
}

impl OpenOptions {
    /// Attempt to change a [`Log`] at the given directory so it becomes
    /// empty and hopefully recovers from some corrupted state.
    ///
    /// Warning: This deletes data, and there is no backup!
    pub fn delete_content(&self, dir: impl Into<GenericPath>) -> crate::Result<()> {
        let dir = dir.into();
        let dir = match dir.as_opt_path() {
            Some(dir) => dir,
            None => return Ok(()),
        };
        let result: crate::Result<()> = (|| {
            // Ensure the directory exist.
            utils::mkdir_p(dir)?;

            // Prevent other writers.
            let lock = ScopedDirLock::new(dir)?;

            // Replace the metadata to an empty state.
            let meta = LogMetadata::new_with_primary_len(PRIMARY_START_OFFSET);
            let meta_path = dir.join(META_FILE);
            meta.write_file(&meta_path, self.fsync)?;

            // Replace the primary log.
            let primary_path = dir.join(PRIMARY_FILE);
            utils::atomic_write_plain(&primary_path, PRIMARY_HEADER, self.fsync)?;

            // Replace indexes so they become empty.
            let log = self
                .clone()
                .create(true)
                .open_with_lock(&dir.into(), &lock)
                .context("cannot open")?;
            log.rebuild_indexes_with_lock(true, &lock)?;

            Ok(())
        })();

        result.context(|| format!("in log::OpenOptions::delete_content({:?})", dir))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// Test RotateLog behavior when fd is limited.

#[cfg(unix)]
mod unix_tests {
    use indexedlog::log::IndexDef;
    use indexedlog::log::IndexOutput;
    use indexedlog::rotate::OpenOptions;
    use tempfile::tempdir;

    const MAX_NOFILE: libc::rlim_t = 256;

    #[test]
    fn test_low_fileno_limit() {
        let verbose = std::env::var("VERBOSE").is_ok();
        for i in 10..30 {
            if verbose {
                eprintln!("Testing RLIMIT_NOFILE = {}", i);
            }
            set_rlimit_nofile(i);
            test_multithread_sync()
        }
    }

    // Test writing using multi-threads. Verify that although some threads error
    // out, the resulting on-disk state is still consistent - data can be opened and
    // read if fileno limit is lifted.
    fn test_multithread_sync() {
        let verbose = std::env::var("VERBOSE").is_ok();
        let dir = tempdir().unwrap();

        // Release mode runs much faster.
        const THREAD_COUNT: u8 = if cfg!(debug_assertions) { 10 } else { 30 };
        const WRITE_COUNT_PER_THREAD: u8 = if cfg!(debug_assertions) { 10 } else { 50 };

        // Some indexes. They have different lag_threshold.
        fn index_ref(data: &[u8]) -> Vec<IndexOutput> {
            vec![IndexOutput::Reference(0..data.len() as u64)]
        }
        let indexes = vec![IndexDef::new("key1", index_ref).lag_threshold(1)];
        let index_len = indexes.len();
        let open_opts = OpenOptions::new()
            .create(true)
            .max_log_count(20)
            .max_bytes_per_log(50)
            .index_defs(indexes);

        use std::sync::Arc;
        use std::sync::Barrier;
        let barrier = Arc::new(Barrier::new(THREAD_COUNT as usize));
        let threads: Vec<_> = (0..THREAD_COUNT)
            .map(|i| {
                let barrier = barrier.clone();
                let open_opts = open_opts.clone();
                let path = dir.path().join("rotatelog");
                std::thread::spawn(move || {
                    barrier.wait();
                    let run = || -> indexedlog::Result<()> {
                        // This might fail with fileno limit
                        let mut log = open_opts.clone().open(&path)?;
                        for j in 1..=WRITE_COUNT_PER_THREAD {
                            let buf = [i, j];
                            log.append(&buf).expect("append should not fail");
                            if j % (i + 1) == 0 || j == WRITE_COUNT_PER_THREAD {
                                // This might fail with fileno limit
                                log.sync()?;
                            }
                            if j % (i + 2) == 0 {
                                // Reopen log. This might fail with fileno limit.
                                log = open_opts.clone().open(&path)?;
                            }
                        }
                        Ok(())
                    };
                    match run() {
                        Ok(_) => {}
                        Err(err) => {
                            if verbose {
                                eprintln!(
                                    " thread {}: {}",
                                    i,
                                    format!("{:?}", err)
                                        .replace("\n\n", "\n")
                                        .replace("\n", "\n  ")
                                )
                            }
                        }
                    }
                })
            })
            .collect();

        // Wait for them. Some of the threads might fail.
        for thread in threads {
            thread.join().expect("joined");
        }

        // Check that if rlimit is restored, then the log can still be opened, and the indexes are
        // functional.
        set_rlimit_nofile(256);
        let log = open_opts.open(dir.path()).unwrap();
        for entry in log.iter().map(|d| d.unwrap()) {
            for index_id in 0..index_len {
                for index_value in log.lookup(index_id, entry.to_vec()).unwrap() {
                    assert_eq!(index_value.unwrap(), entry);
                }
            }
        }
    }

    fn set_rlimit_nofile(n: libc::rlim_t) {
        unsafe {
            let limit = libc::rlimit {
                rlim_cur: n,
                rlim_max: MAX_NOFILE,
            };
            libc::setrlimit(libc::RLIMIT_NOFILE, &limit as *const libc::rlimit);
        }
    }
}


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

// for macros
pub use insta;

pub fn setup() {
    const WORKSPACE: &str = "INSTA_WORKSPACE_ROOT";
    const UPDATE: &str = "INSTA_UPDATE";
    if std::env::var(WORKSPACE).is_err() {
        let mut root = std::path::PathBuf::from(file!());
        root.pop();
        root.pop();
        std::env::set_var(WORKSPACE, root);
    }
    if std::env::var(UPDATE).is_err() {
        std::env::set_var(UPDATE, "no");
    }
}

pub fn run(
    test_name: &str,
    is_cargo: bool,
    snapshot: &str,
    file: &str,
    function_name: &str,
    module: &str,
    line: u32,
    expr: &str,
) {
    let command = if is_cargo {
        "INSTA_UPDATE=1 cargo test ..."
    } else {
        "buck test ... -- --env INSTA_UPDATE=1"
    };
    println!(
        "{:=^80}\n",
        format!(" Run `{}` to update snapshots ", command)
    );
    let file_name = std::path::Path::new(file)
        .file_name()
        .and_then(|p| p.to_str())
        .unwrap();
    insta::_macro_support::assert_snapshot(
        test_name.into(),
        snapshot,
        "unused",
        function_name,
        // buck builds have a _unittest module suffix which cargo doesn't
        // this makes the snapshot location consistent on both
        &module.replacen("_unittest", "", 1),
        file_name,
        line,
        expr,
    )
    .unwrap();
}

/// Assert that the serde json representation of given expression matches the snapshot
/// stored on disk.
#[macro_export]
macro_rules! assert_json {
    ($value: expr, $test_name: ident) => {{
        $crate::setup();

        $crate::run(
            stringify!($test_name),
            option_env!("CARGO_MANIFEST_DIR").is_some(),
            &serde_json::to_string(&$value).unwrap(),
            file!(),
            $crate::insta::_function_name!(),
            module_path!(),
            line!(),
            stringify!($value),
        );
    }};
}


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

use std::io::stderr;
use std::io::stdin;
use std::io::stdout;
use std::io::Cursor;
use std::io::IsTerminal;
use std::sync::Weak;

use crate::IsTty;

impl IsTty for std::io::Stdin {
    fn is_tty(&self) -> bool {
        stdin().is_terminal()
    }
    fn is_stdin(&self) -> bool {
        true
    }
}

impl IsTty for std::io::Stdout {
    fn is_tty(&self) -> bool {
        stdout().is_terminal()
    }
    fn is_stdout(&self) -> bool {
        true
    }
}

impl IsTty for std::io::Stderr {
    fn is_tty(&self) -> bool {
        stderr().is_terminal()
    }
    fn is_stderr(&self) -> bool {
        true
    }
}

impl IsTty for Vec<u8> {
    fn is_tty(&self) -> bool {
        false
    }
}

impl<'a> IsTty for &'a [u8] {
    fn is_tty(&self) -> bool {
        false
    }
}

impl<T> IsTty for Cursor<T> {
    fn is_tty(&self) -> bool {
        false
    }
}

impl IsTty for crate::IOInput {
    fn is_tty(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.input.is_tty()
    }
    fn is_stdin(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.input.is_stdin()
    }
    fn pager_active(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.input.pager_active()
    }
}

impl IsTty for crate::IOOutput {
    fn is_tty(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.output.is_tty()
    }
    fn is_stdout(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.output.is_stdout()
    }
    fn pager_active(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        inner.output.pager_active()
    }
}

impl IsTty for crate::IOError {
    fn is_tty(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        if let Some(error) = inner.error.as_ref() {
            error.is_tty()
        } else {
            false
        }
    }
    fn is_stderr(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        if let Some(error) = inner.error.as_ref() {
            error.is_stderr()
        } else {
            false
        }
    }
    fn pager_active(&self) -> bool {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return false,
        };
        let inner = inner.io_state.lock();
        if let Some(error) = inner.error.as_ref() {
            error.pager_active()
        } else {
            false
        }
    }
}

pub(crate) struct PipeWriterWithTty {
    inner: pipe::PipeWriter,
    pretend_tty: bool,
    pub(crate) pretend_stdout: bool,
}

impl std::io::Write for PipeWriterWithTty {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.inner.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

impl IsTty for PipeWriterWithTty {
    fn is_tty(&self) -> bool {
        self.pretend_tty
    }
    fn is_stdout(&self) -> bool {
        self.pretend_stdout
    }
    fn pager_active(&self) -> bool {
        true
    }
}

impl PipeWriterWithTty {
    pub fn new(inner: pipe::PipeWriter, pretend_tty: bool) -> Self {
        Self {
            inner,
            pretend_tty,
            pretend_stdout: false,
        }
    }
}


/*
 * 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::io;
use std::mem;
use std::sync::Arc;
use std::sync::Weak;
use std::thread::spawn;
use std::time::Duration;
use std::time::Instant;

use configmodel::Config;
use configmodel::ConfigExt;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use parking_lot::RwLock;
use pipe::pipe;
use pipe::PipeWriter;
use streampager::action::Action;
use streampager::config::InterfaceMode;
use streampager::config::WrappingMode;
use streampager::Pager;
use term::make_real_term;
use term::DumbTerm;
use term::DumbTty;
use term::Term;
use term::DEFAULT_TERM_HEIGHT;
use term::DEFAULT_TERM_WIDTH;
use termwiz::surface::change::ChangeSequence;
use termwiz::surface::Change;

mod impls;
mod term;

use crate::impls::PipeWriterWithTty;

// IO is Clone, but care must be taken to drop the IO object normally
// to ensure things are cleaned up before the process exits.
#[derive(Clone)]
pub struct IO {
    inner: Arc<Inner>,
}

/// Implements `io::Read` on the input stream.
#[derive(Clone)]
pub struct IOInput(Weak<Inner>);

/// Implements `io::Write` on the output stream.
#[derive(Clone)]
pub struct IOOutput(Weak<Inner>);

/// Implements `io::Write` on the error stream.
#[derive(Clone)]
pub struct IOError(Weak<Inner>);

/// Provides a way to set progress, without requiring the `&IO` reference.
#[derive(Clone)]
pub struct IOProgress(Weak<Inner>);

struct Inner {
    io_state: Mutex<IOState>,
    // Use a separate Mutex for quitting the pager without blocking.
    //
    // Note: wait_pager is in io_state so quit_pager during wait_pager
    // won't block.
    pager_quit_func: Mutex<Option<Box<dyn FnOnce() + Send>>>,
}

struct IOState {
    input: Box<dyn Read>,
    output: Box<dyn Write>,
    error: Option<Box<dyn Write>>,
    pager_progress: Option<Box<dyn Term + Send + Sync>>,

    term: Option<Box<dyn Term + Send + Sync>>,

    // Used to decide whether to render progress bars.
    output_on_new_line: bool,
    error_on_new_line: bool,
    // Whether progress is non-empty.
    progress_has_content: bool,
    // Whether progress (stderr) and stdout (is likely) sharing output.
    progress_conflict_with_output: bool,
    // Whether to redirect stdout writes to stderr. More useful for pager use-case.
    redirect_err_to_out: bool,

    // How many (nested) blocks want progress output disabled.
    progress_disabled: usize,

    // Function to wait for the pager to cleanup (restore terminal state).
    // Might block, unless `pager_quit_func` is called right before.
    pager_wait_func: Option<Box<dyn FnOnce() + Send>>,
}

/// The "main" IO used by the process.
///
/// This global state makes it easier for Python bindings
/// (ex. "pyio") to obtain the IO state without needing
/// to pass the state across layers. This is similar to
/// `std::io::stdout` etc being globally accessible.
///
/// Use `IO::set_main()` to set the main IO, and `IO::main()`
/// to obtain the "main" `IO`.
static MAIN_IO_REF: Lazy<RwLock<Option<Weak<Inner>>>> = Lazy::new(Default::default);

fn colors_disabled_via_env() -> bool {
    hgplain::is_plain(Some("color")) || std::env::var("TERM").ok().as_deref() == Some("dumb")
}

pub trait IsTty {
    fn is_tty(&self) -> bool;

    fn is_stdin(&self) -> bool {
        false
    }
    fn is_stdout(&self) -> bool {
        false
    }
    fn is_stderr(&self) -> bool {
        false
    }

    /// Whether this connection is capable of colors, ignoring whether
    /// the user has enabled colors.
    fn can_color(&self) -> bool {
        self.is_tty() && !colors_disabled_via_env()
    }

    fn pager_active(&self) -> bool {
        false
    }
}

pub trait Read: io::Read + IsTty + Any + Send + Sync {
    fn as_any(&self) -> &dyn Any;

    fn type_name(&self) -> &'static str {
        std::any::type_name::<Self>()
    }
}

pub trait Write: io::Write + IsTty + Any + Send + Sync {
    fn as_any(&self) -> &dyn Any;

    fn type_name(&self) -> &'static str {
        std::any::type_name::<Self>()
    }
}

impl<T: io::Read + IsTty + Any + Send + Sync> Read for T {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

impl<T: io::Write + IsTty + Any + Send + Sync> Write for T {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

// Read from input.
impl io::Read for IOInput {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(0),
        };
        let mut inner = inner.io_state.lock();
        inner.input.read(buf)
    }
}

// Write to error.
impl io::Write for IOError {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(buf.len()),
        };
        let mut inner = inner.io_state.lock();
        if inner.redirect_err_to_out {
            inner.clear_progress_for_output()?;
            inner.output_on_new_line = buf.ends_with(b"\n");
            return inner.output.write(buf);
        }
        inner.clear_progress_for_error()?;
        inner.error_on_new_line = buf.ends_with(b"\n");
        if let Some(error) = inner.error.as_mut() {
            error.write(buf)
        } else {
            Ok(buf.len())
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(()),
        };
        let mut inner = inner.io_state.lock();
        if let Some(error) = inner.error.as_mut() {
            error.flush()?;
        }
        Ok(())
    }
}

// Write to output.
impl io::Write for IOOutput {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(buf.len()),
        };
        let mut inner = inner.io_state.lock();
        inner.clear_progress_for_output()?;
        inner.output_on_new_line = buf.ends_with(b"\n");
        inner.output.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(()),
        };
        let mut inner = inner.io_state.lock();
        inner.output.flush()
    }
}

impl IOProgress {
    /// Set progress to the given text.
    pub fn set(&self, changes: &[Change]) -> io::Result<()> {
        let inner = match Weak::upgrade(&self.0) {
            Some(inner) => inner,
            None => return Ok(()),
        };
        let mut inner = inner.io_state.lock();
        inner.set_progress(changes)
    }

    pub fn term_size(&self) -> (usize, usize) {
        if let Some(inner) = Weak::upgrade(&self.0) {
            inner.io_state.lock().term_size()
        } else {
            (DEFAULT_TERM_WIDTH, DEFAULT_TERM_HEIGHT)
        }
    }
}

impl IO {
    pub fn with_input<R>(&self, f: impl FnOnce(&mut dyn Read) -> R) -> R {
        f(self.inner.io_state.lock().input.as_mut())
    }

    pub fn with_output<R>(&self, f: impl FnOnce(&dyn Write) -> R) -> R {
        f(self.inner.io_state.lock().output.as_ref())
    }

    pub fn with_error<R>(&self, f: impl FnOnce(Option<&dyn Write>) -> R) -> R {
        f(self.inner.io_state.lock().error.as_deref())
    }

    /// Returns a clonable value that impls [`io::Write`] to `error` stream.
    /// The output is associated with the `IO` so if the `IO` starts a pager,
    /// the error stream will be properly redirected to the pager.
    ///
    /// If this IO is dropped, the IOError stream will be redirected to null.
    pub fn error(&self) -> IOError {
        IOError(Arc::downgrade(&self.inner))
    }

    /// Returns a clonable value that impls [`io::Read`] from `input` stream.
    /// The output is associated with the `IO` so if the `IO` starts a pager,
    /// the error stream will be properly redirected to the pager.
    ///
    /// If this IO is dropped, the IOInput stream will be redirected to null.
    pub fn input(&self) -> IOInput {
        IOInput(Arc::downgrade(&self.inner))
    }

    /// Returns a clonable value that impls [`io::Write`] to `output` stream.
    /// The output is associated with the `IO` so if the `IO` starts a pager,
    /// the error stream will be properly redirected to the pager.
    ///
    /// If this IO is dropped, the IOError stream will be redirected to null.
    pub fn output(&self) -> IOOutput {
        IOOutput(Arc::downgrade(&self.inner))
    }

    /// Returns a clonable value that provides a way to set progress text.
    pub fn progress(&self) -> IOProgress {
        IOProgress(Arc::downgrade(&self.inner))
    }

    pub fn new<IS, OS, ES>(input: IS, output: OS, error: Option<ES>) -> Self
    where
        IS: Read + 'static,
        OS: Write + 'static,
        ES: Write + 'static,
    {
        let progress_conflict_with_output = match &error {
            None => false, // No progress bar.
            Some(e) => e.is_tty() && output.is_tty(),
        };

        let inner = Inner {
            io_state: Mutex::new(IOState {
                input: Box::new(input),
                output: Box::new(output),
                error: error.map(|e| Box::new(e) as Box<dyn Write>),
                pager_progress: None,
                term: None,
                progress_conflict_with_output,
                output_on_new_line: true,
                error_on_new_line: true,
                progress_has_content: false,
                progress_disabled: 0,
                redirect_err_to_out: false,
                pager_wait_func: None,
            }),
            pager_quit_func: Default::default(),
        };

        Self {
            inner: Arc::new(inner),
        }
    }

    /// Wait for the pager to exit, and restore outputs to stdio.
    /// Might block if the pager is waiting for the user to exit.
    pub fn wait_pager(&self) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        inner.flush()?;

        // Drop the piped streams (to the pager).
        // XXX: Stdio is hard-coded for wait_pager.
        inner.input = Box::new(io::stdin());
        inner.output = Box::new(io::stdout());
        inner.error = Some(Box::new(io::stderr()));
        inner.redirect_err_to_out = false;
        inner.pager_progress = None;

        // This might block but shouldn't block quit_pager.
        inner.wait_pager();

        // pager_quit_func is no longer needed.
        let _ = self.inner.pager_quit_func.lock().take();

        Ok(())
    }

    /// Quit the pager now and wait for it to complete cleanup.
    /// Does not restore `input`, `output`, `error`, should only
    /// be used before exiting.
    pub fn quit_pager(&self) {
        let mut lock = self.inner.pager_quit_func.lock();
        let mut func = None;
        mem::swap(&mut func, &mut lock);
        if let Some(func) = func {
            drop(lock);
            func();
        }
    }

    pub fn write(&self, data: impl AsRef<[u8]>) -> io::Result<()> {
        let data = data.as_ref();
        let mut inner = self.inner.io_state.lock();
        inner.clear_progress_for_output()?;
        inner.output_on_new_line = data.ends_with(b"\n");
        inner.output.write_all(data)?;
        Ok(())
    }

    pub fn write_err(&self, data: impl AsRef<[u8]>) -> io::Result<()> {
        let data = data.as_ref();
        let mut inner = self.inner.io_state.lock();
        if inner.redirect_err_to_out {
            inner.clear_progress_for_output()?;
            inner.output_on_new_line = data.ends_with(b"\n");
            inner.output.write_all(data)?;
            return Ok(());
        }
        inner.clear_progress_for_error()?;
        inner.error_on_new_line = data.ends_with(b"\n");
        if let Some(ref mut error) = inner.error {
            error.write_all(data)?;
        }
        Ok(())
    }

    pub fn set_progress(&self, changes: &[Change]) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        inner.set_progress(changes)
    }

    pub fn set_progress_str(&self, data: &str) -> io::Result<()> {
        self.set_progress(&[data.into()])
    }

    pub fn flush(&self) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        inner.flush()
    }

    pub fn stdio() -> Self {
        let progress_conflict_with_output = io::stderr().is_tty() && io::stdout().is_tty();
        let inner = Inner {
            io_state: Mutex::new(IOState {
                input: Box::new(io::stdin()),
                output: Box::new(io::stdout()),
                error: Some(Box::new(io::stderr())),
                pager_progress: None,
                term: None,
                progress_conflict_with_output,
                progress_has_content: false,
                progress_disabled: 0,
                output_on_new_line: true,
                error_on_new_line: true,
                redirect_err_to_out: false,
                pager_wait_func: None,
            }),
            pager_quit_func: Default::default(),
        };
        Self {
            inner: Arc::new(inner),
        }
    }

    pub fn setup_term(&mut self) -> Result<(), termwiz::Error> {
        let mut inner = self.inner.io_state.lock();

        if std::env::var_os("TESTTMP").is_some() {
            // Use dumb terminal with static width/height for tests.
            inner.term = Some(Box::new(DumbTerm::new(DumbTty::new(Box::new(
                io::stderr(),
            )))?));
        } else {
            inner.term = Some(make_real_term()?);
        }

        Ok(())
    }

    /// Obtain the main IO.
    ///
    /// The main IO must be set via `set_main` and is still alive.
    /// Otherwise, this function will return an error.
    pub fn main() -> io::Result<Self> {
        let opt_main_io = MAIN_IO_REF.read();
        let main_io = match opt_main_io.as_ref() {
            Some(io) => io,
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::NotConnected,
                    "IO::main() is not available (call set_main first)",
                ));
            }
        };

        if let Some(inner) = Weak::upgrade(&*main_io) {
            Ok(Self { inner })
        } else {
            Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "IO::main() is not available (dropped)",
            ))
        }
    }

    /// Set the current IO as the main IO.
    ///
    /// Note: If the current IO gets dropped then the main IO will be dropped
    /// too and [`IO::main`] will return an error.
    pub fn set_main(&self) {
        let mut main_io_ref = MAIN_IO_REF.write();
        *main_io_ref = Some(Arc::downgrade(&self.inner));
    }

    /// Check if the pager is active.
    pub fn is_pager_active(&self) -> bool {
        let state = self.inner.io_state.lock();
        state.is_pager_active()
    }

    /// Starts a pager.
    ///
    /// It is recommended to run [`IO::flush`] and [`IO::wait_pager`] before exiting.
    pub fn start_pager(&self, config: &dyn Config) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        if inner.is_pager_active() {
            return Ok(());
        }

        inner.set_progress(&[])?;

        let mut pager = Pager::new_using_system_terminal()
            .or_else(|_| Pager::new_using_stdio())
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

        // Configure the pager.
        // The Hybrid mode is similar to "-FX" from "less".
        let mut interface_mode = InterfaceMode::Hybrid;
        // Similar to "less" behavior - lines are wrapped and copy-paste preserves long ines.
        let mut wrapping_mode = WrappingMode::GraphemeBoundary;
        if let Some(mode_str) = config.get("pager", "interface") {
            let mode = InterfaceMode::from(mode_str.as_ref());
            interface_mode = mode;
        }
        if let Ok(Some(wrapping_mode_str)) = config.get_opt::<String>("pager", "wrapping-mode") {
            match wrapping_mode_str.to_lowercase().as_str() {
                "word" => wrapping_mode = WrappingMode::WordBoundary,
                "unwrapped" => wrapping_mode = WrappingMode::Unwrapped,
                _ => {}
            }
        }
        pager.set_wrapping_mode(wrapping_mode);
        // Similar to "less" default.
        pager.set_scroll_past_eof(config.must_get("pager", "scroll-past-eof").unwrap_or(false));
        pager.set_interface_mode(interface_mode);
        // Setting this to "false" will preserve pending terminal input in
        // "direct" mode at the cost of disabling streampager user input
        // handling.
        pager.set_startup_poll_input(
            config
                .must_get("pager", "startup-poll-input")
                .unwrap_or(false),
        );

        let (out_read, out_write) = pipe();
        let (err_read, err_write) = pipe();
        let (prg_read, prg_write) = pipe();

        let out_is_tty = inner.output.is_tty();
        let out_is_stdout = inner.output.is_stdout();
        let err_is_tty = inner
            .error
            .as_ref()
            .map(|e| e.is_tty())
            .unwrap_or_else(|| out_is_tty);

        inner.flush()?;
        inner.output = {
            let mut pipe = PipeWriterWithTty::new(out_write, out_is_tty);
            pipe.pretend_stdout = out_is_stdout;
            Box::new(pipe)
        };
        pager
            .add_stream(out_read, "")
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;

        // Only use the pager for error stream if error stream is a tty.
        // This makes `hg 2>foo` works as expected.
        if err_is_tty {
            inner.error = Some(Box::new(PipeWriterWithTty::new(err_write, err_is_tty)));
            let separate =
                config.get_opt::<bool>("pager", "separate-stderr").ok() == Some(Some(true));
            inner.redirect_err_to_out = !separate;
            if separate {
                pager
                    .add_error_stream(err_read, "stderr")
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
            }
        }

        let mut pager_term = DumbTerm::new(DumbTty::new(Box::new(prg_write)))
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
        pager_term.set_separator(0x0C);
        inner.pager_progress = Some(Box::new(pager_term));
        pager.set_progress_stream(prg_read);

        let pager_action_sender = pager.action_sender();
        let pager_thread_handler = Arc::new(Mutex::new(Some(spawn(|| {
            let _ = pager.run();
        }))));

        let pager_thread_handler_wait = pager_thread_handler.clone();
        inner.pager_wait_func = Some(Box::new(move || {
            if let Some(handler) = pager_thread_handler_wait.lock().take() {
                let _ = handler.join();
            }
        }));

        self.inner.pager_quit_func.lock().replace(Box::new(move || {
            let _ = pager_action_sender.send(Action::Quit);

            if let Some(handler) = pager_thread_handler.lock().take() {
                let start = Instant::now();

                // Wait up to 10ms for streampager to shut down.
                while !handler.is_finished() && start.elapsed() < Duration::from_millis(10) {
                    std::thread::sleep(Duration::from_millis(1));
                }
            }
        }));

        Ok(())
    }

    /// Disable progress rendering.
    /// - `disable_progress(true)` disables progress rendering. It can be nested.
    /// - `disable_progress(false)` cancels out a `disable_progress(true)`.
    ///   If all `disable_progress(true)` are canceled out, restore the progress
    ///   rendering.
    pub fn disable_progress(&self, disabled: bool) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        if disabled {
            inner.progress_disabled += 1;
            inner.set_progress(&[])?;
        } else {
            if inner.progress_disabled == 0 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "disable_progress(false) called without matching disable_progress(true)",
                ));
            }
            inner.progress_disabled -= 1;
        }
        Ok(())
    }

    pub fn set_progress_pipe_writer(&self, progress: Option<PipeWriter>) -> io::Result<()> {
        let mut inner = self.inner.io_state.lock();
        inner.pager_progress = match progress {
            Some(progress) => {
                let mut pager_term = DumbTerm::new(DumbTty::new(Box::new(progress)))
                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
                pager_term.set_separator(0x0C);
                Some(Box::new(pager_term))
            }
            None => None,
        };
        Ok(())
    }
}

impl IOState {
    pub(crate) fn flush(&mut self) -> io::Result<()> {
        self.output.flush()?;
        if let Some(ref mut error) = self.error {
            error.flush()?;
        }
        Ok(())
    }

    /// Clear the progress (temporarily) for other output.
    fn clear_progress_for_error(&mut self) -> io::Result<()> {
        if self.progress_has_content && self.pager_progress.is_none() {
            self.progress_has_content = false;
            if let Some(ref mut term) = self.term {
                write_term_progress(term, &[])
                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
            }
        }
        Ok(())
    }

    /// Clear the progress (temporarily) if it ("stderr") conflicts with "stdout" output.
    fn clear_progress_for_output(&mut self) -> io::Result<()> {
        // If self.progress is set. Progress writes to streampager, and does not need clear.
        if self.progress_conflict_with_output && self.pager_progress.is_none() {
            self.clear_progress_for_error()?;
        }
        Ok(())
    }

    fn term_size(&mut self) -> (usize, usize) {
        if let Some(ref mut term) = self.term {
            if let Ok((cols, rows)) = term.size() {
                return (cols, rows);
            }
        }

        (DEFAULT_TERM_WIDTH, DEFAULT_TERM_HEIGHT)
    }

    fn set_progress(&mut self, mut changes: &[Change]) -> io::Result<()> {
        let inner = self;
        if inner.progress_disabled > 0 {
            changes = &[];
        }

        if let Some(ref mut progress) = inner.pager_progress {
            write_term_progress(progress, changes)
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
        } else {
            if !inner.output_on_new_line || !inner.error_on_new_line {
                // There is a line that hasn't ended.
                // Not suitable to render progress bars.
                changes = &[];
            }

            // Fast path: empty progress, unchanged.
            if changes.is_empty() && !inner.progress_has_content {
                return Ok(());
            }

            // Flush pending output if it might conflict with progress.
            if inner.progress_conflict_with_output {
                inner.output.flush()?;
            }

            if let Some(ref mut term) = inner.term {
                inner.progress_has_content = !changes.is_empty();
                write_term_progress(term, changes)
                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
            }
        }
        Ok(())
    }

    /// Wait for the pager to exit and cleanup (restore terminal).
    fn wait_pager(&mut self) {
        let mut func = None;
        mem::swap(&mut func, &mut self.pager_wait_func);
        if let Some(func) = func {
            func();
        }
    }

    fn is_pager_active(&self) -> bool {
        self.pager_wait_func.is_some()
    }
}

/// Write data to the progress area by clearing everything after
/// cursor, writing data, then moving cursor back.
fn write_term_progress(
    term: &mut Box<dyn Term + Send + Sync>,
    changes: &[Change],
) -> Result<(), termwiz::Error> {
    let (cols, rows) = term.size()?;
    let mut change_seq = ChangeSequence::new(rows, cols);
    change_seq.add(Change::ClearToEndOfScreen(Default::default()));
    change_seq.add_changes(changes.to_vec());
    change_seq.move_to((0, 0));
    term.render(&change_seq.consume())?;

    Ok(())
}

impl Drop for IOState {
    fn drop(&mut self) {
        let _ = self.set_progress(&[]);
        let _ = self.flush();
        // Drop the output and error. This sends EOF to pager.
        self.output = Box::new(Vec::new());
        self.error = None;
        self.pager_progress = None;
        self.wait_pager();
    }
}


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

use std::io;

use termwiz::caps::Capabilities;
use termwiz::render::terminfo::TerminfoRenderer;
use termwiz::render::RenderTty;
use termwiz::surface::Change;
use termwiz::terminal::Terminal;
use termwiz::Result;

use crate::IsTty;

pub(crate) const DEFAULT_TERM_WIDTH: usize = 80;
pub(crate) const DEFAULT_TERM_HEIGHT: usize = 25;

#[cfg(windows)]
mod windows_term;

#[cfg(unix)]
mod unix_term;

/// Term is a minimally skinny abstraction over termwiz::Terminal.
/// It makes it easy to swap in other things for testing.
pub(crate) trait Term {
    fn render(&mut self, changes: &[Change]) -> Result<()>;
    fn size(&mut self) -> Result<(usize, usize)>;
}

impl<T: Terminal> Term for T {
    fn render(&mut self, changes: &[Change]) -> Result<()> {
        Terminal::render(self, changes)?;
        self.flush()?;
        Ok(())
    }

    fn size(&mut self) -> Result<(usize, usize)> {
        let size = self.get_screen_size()?;
        Ok((size.cols, size.rows))
    }
}

/// DumbTerm allows writing termwiz Changes to an arbitrary writer,
/// ignoring lack of ttyness and using a default terminal size.
pub(crate) struct DumbTerm<W: RenderTty + io::Write> {
    tty: W,
    renderer: TerminfoRenderer,
    separator: Option<u8>,
}

impl<W: RenderTty + io::Write> DumbTerm<W> {
    pub fn new(tty: W) -> Result<Self> {
        Ok(Self {
            tty,
            renderer: TerminfoRenderer::new(caps()?),
            separator: None,
        })
    }

    pub fn set_separator(&mut self, sep: u8) {
        self.separator = Some(sep);
    }
}

impl<W: RenderTty + io::Write> Term for DumbTerm<W> {
    fn render(&mut self, changes: &[Change]) -> Result<()> {
        self.renderer.render_to(changes, &mut self.tty)?;
        if let Some(sep) = self.separator {
            self.tty.write_all(&[sep])?;
            self.tty.flush()?;
        }
        Ok(())
    }

    fn size(&mut self) -> Result<(usize, usize)> {
        self.tty.get_size_in_cells()
    }
}

pub(crate) struct DumbTty {
    write: Box<dyn io::Write + Send + Sync>,
}

impl DumbTty {
    pub fn new(write: Box<dyn io::Write + Send + Sync>) -> Self {
        Self { write }
    }
}

impl RenderTty for DumbTty {
    fn get_size_in_cells(&mut self) -> Result<(usize, usize)> {
        Ok((DEFAULT_TERM_WIDTH, DEFAULT_TERM_HEIGHT))
    }
}

impl io::Write for DumbTty {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.write.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.write.flush()
    }
}

fn caps() -> Result<Capabilities> {
    let hints = termwiz::caps::ProbeHints::new_from_env().mouse_reporting(Some(false));
    termwiz::caps::Capabilities::new_with_hints(hints)
}

pub(crate) fn make_real_term() -> Result<Box<dyn Term + Send + Sync>> {
    // Don't use the real termwiz terminal yet because:
    //   1. On Windows, it disables automatic \n -> \r\n conversion.
    //   2. On Mac, we were detecting /dev/tty as usable, but ended up blocking when dropping the UnixTerminal object (when invoked via buck).
    //   3. Termwiz sets up a SIGWINCH handler which causes crash in Python crecord stuff.

    #[cfg(windows)]
    {
        let stderr = io::stderr();
        if stderr.is_tty() {
            let tty = windows_term::WindowsTty::new(Box::new(stderr));
            return Ok(Box::new(DumbTerm::new(tty)?));
        }
    }

    #[cfg(unix)]
    {
        let stderr = io::stderr();
        if stderr.is_tty() {
            let tty = unix_term::UnixTty::new(Box::new(stderr));
            return Ok(Box::new(DumbTerm::new(tty)?));
        }
    }

    termwiz::bail!("no suitable term output file");
}


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

use std::io;
use std::os::unix::prelude::AsRawFd;

use termwiz::render::RenderTty;

pub(crate) trait RawWrite: io::Write + AsRawFd + Send + Sync {}

impl RawWrite for io::Stderr {}

/// UnixTty is a shim tty to avoid using the real termwiz
/// UnixTerminal for now.
pub(crate) struct UnixTty {
    write: Box<dyn RawWrite>,
}

impl UnixTty {
    pub fn new(write: Box<dyn RawWrite>) -> Self {
        Self { write }
    }
}

impl RenderTty for UnixTty {
    fn get_size_in_cells(&mut self) -> termwiz::Result<(usize, usize)> {
        match terminal_size::terminal_size_using_fd(self.write.as_raw_fd()) {
            Some((width, height)) => Ok((width.0 as _, height.0 as _)),
            // Fallback size, just in case.
            None => Ok((super::DEFAULT_TERM_WIDTH, super::DEFAULT_TERM_HEIGHT)),
        }
    }
}

impl io::Write for UnixTty {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.write.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.write.flush()
    }
}


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

use std::io;
use std::os::windows::io::AsRawHandle;

use termwiz::render::RenderTty;

pub(crate) trait RawWrite: io::Write + AsRawHandle + Send + Sync {}

impl RawWrite for io::Stderr {}

/// WindowsTty is a shim tty to avoid using the real termwiz
/// WindowsTerminal for now. The WindowsTerminal disables automatic
/// "\n" -> "\r\n" conversion, so it messes up our normal newline
/// delimited output that goes to stderr/stdout.
pub(crate) struct WindowsTty {
    write: Box<dyn RawWrite>,
}

impl WindowsTty {
    pub fn new(write: Box<dyn RawWrite>) -> Self {
        Self { write }
    }
}

impl RenderTty for WindowsTty {
    fn get_size_in_cells(&mut self) -> termwiz::Result<(usize, usize)> {
        match terminal_size::terminal_size_using_handle(self.write.as_raw_handle()) {
            Some((width, height)) => Ok((width.0 as _, height.0 as _)),
            // Fallback size, just in case.
            None => Ok((super::DEFAULT_TERM_WIDTH, super::DEFAULT_TERM_HEIGHT)),
        }
    }
}

impl io::Write for WindowsTty {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.write.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.write.flush()
    }
}


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

use std::io::Write;

use io::IO;
use lazystr::LazyStr;

/// TermLogger mixes the IO object with knowledge of output verbosity.
pub struct TermLogger {
    output: io::IOOutput,
    error: io::IOError,
    quiet: bool,
    verbose: bool,
}

impl TermLogger {
    pub fn new(io: &IO) -> Self {
        TermLogger {
            output: io.output(),
            error: io.error(),
            quiet: false,
            verbose: false,
        }
    }

    pub fn with_quiet(mut self, quiet: bool) -> Self {
        self.quiet = quiet;
        self
    }

    pub fn with_verbose(mut self, verbose: bool) -> Self {
        self.verbose = verbose;
        self
    }

    /// Write to stdout if not --quiet.
    pub fn info(&mut self, msg: impl LazyStr) {
        if !self.quiet {
            Self::write(&mut self.output, msg.to_str())
        }
    }

    /// Write to stderr.
    pub fn warn(&mut self, msg: impl AsRef<str>) {
        Self::write(&mut self.error, msg)
    }

    /// Write to stdout if --verbose.
    pub fn verbose(&mut self, msg: impl LazyStr) {
        if self.verbose {
            Self::write(&mut self.output, msg.to_str())
        }
    }

    /// Short client program name.
    pub fn cli_name(&self) -> &'static str {
        identity::cli_name()
    }

    pub fn flush(&mut self) {
        let _ = self.output.flush();
        let _ = self.error.flush();
    }

    fn write(mut w: impl Write, msg: impl AsRef<str>) {
        let msg = identity::default().punch(msg.as_ref());

        if let Err(err) = || -> std::io::Result<()> {
            w.write_all(msg.as_bytes())?;
            if !msg.ends_with('\n') {
                w.write_all(b"\n")?;
            }
            Ok(())
        }() {
            tracing::warn!(?msg, ?err, "error writing command output");
        }
    }
}

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

    #[test]
    fn test_quiet() {
        let io = IO::new("".as_bytes(), Vec::new(), Some(Vec::new()));
        let mut logger = TermLogger::new(&io).with_quiet(true);
        logger.info("hello");
        logger.warn("error");
        assert_eq!(get_stdout(&io), "");
        assert_eq!(get_stderr(&io), "error\n");
    }

    #[test]
    fn test_default() {
        let io = IO::new("".as_bytes(), Vec::new(), Some(Vec::new()));
        let mut logger = TermLogger::new(&io);
        logger.info("status");
        logger.verbose("verbose");
        logger.warn("warn");
        assert_eq!(get_stdout(&io), "status\n");
        assert_eq!(get_stderr(&io), "warn\n");
    }

    #[test]
    fn test_verbose() {
        let io = IO::new("".as_bytes(), Vec::new(), Some(Vec::new()));
        let mut logger = TermLogger::new(&io);
        logger.verbose(|| -> String {
            panic!("don't call me!");
        });
        assert_eq!(get_stdout(&io), "");

        let mut logger = logger.with_verbose(true);
        logger.verbose(|| "okay".to_string());
        assert_eq!(get_stdout(&io), "okay\n");
    }

    fn get_stdout(io: &IO) -> String {
        let stdout = io.with_output(|o| o.as_any().downcast_ref::<Vec<u8>>().unwrap().clone());
        String::from_utf8(stdout).unwrap()
    }

    fn get_stderr(io: &IO) -> String {
        let stderr = io.with_error(|o| {
            o.unwrap()
                .as_any()
                .downcast_ref::<Vec<u8>>()
                .unwrap()
                .clone()
        });
        String::from_utf8(stderr).unwrap()
    }
}


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

use std::io;
use std::io::Write;

use termwiz::caps::Capabilities;
pub use termwiz::caps::ColorLevel;
use termwiz::caps::ProbeHints;
use termwiz::cell::CellAttributes;
use termwiz::cell::Intensity;
use termwiz::cell::Underline;
use termwiz::color::AnsiColor;
use termwiz::color::ColorSpec;
use termwiz::color::RgbColor;
use termwiz::render::terminfo::TerminfoRenderer;
use termwiz::render::RenderTty;
use termwiz::surface::Change;

/// Evaluate style specs given supported color level, yielding a
/// CellAttributes object with corresponding fields filled in.
///
/// Style spec format:
///    - effect: a concrete text modifier (e.g. "green" or "bold").
///              Various color formats are supported depending on color level:
///                 4-bit: green, red_background, etc.
///                 8-bit: color123, color100_background, etc
///                 24-bit: #FFF, #A1B2C3, DarkOrange2 (and more - see termwiz)
///    - style: effect(+effect)*
///             A list of effects separated by "+". Effects are only applied
///             if all effects in the list are valid.
///    - spec: style(:style)*
///            Priority order list of styles. First valid style wins.
///    - specs: spec( spec)*
///             Space separated list of specs. All specs are applied in order.
///
/// Examples:
///
/// Pick one from "color214" and "yellow", then combined with "bold":
///
/// ```text
/// color214:yellow bold
/// ```
///
/// Pick one from "color214" and "yellow+bold":
///
/// ```text
/// color214:yellow+bold
/// ```
pub fn eval_style(level: ColorLevel, style_specs: &str) -> CellAttributes {
    let mut attrs = CellAttributes::blank();

    'specs: for spec in style_specs.split_whitespace() {
        for style in spec.split(':') {
            if eval_effects(level, style.split('+'), &mut attrs) {
                continue 'specs;
            }
        }

        tracing::warn!(spec, "couldn't apply style spec");
    }

    attrs
}

struct DumbTty<'a> {
    w: &'a mut dyn Write,
}

impl RenderTty for DumbTty<'_> {
    fn get_size_in_cells(&mut self) -> termwiz::Result<(usize, usize)> {
        Ok((80, 26))
    }
}

impl io::Write for DumbTty<'_> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.w.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.w.flush()
    }
}

// The main purpose of this object is to cache the Capabilities.
pub struct Styler {
    level: ColorLevel,
    renderer: TerminfoRenderer,
}

fn enable_compat_mode() -> bool {
    std::env::var("SL_DISABLE_TERMINFO_COMPAT").is_err()
}

impl Styler {
    /// Initialize a styler, letting termwiz sniff the supported color level.
    pub fn new() -> termwiz::Result<Styler> {
        let mut hints = ProbeHints::new_from_env()
            .force_terminfo_render_to_use_ansi_sgr(Some(enable_compat_mode()));
        if cfg!(test) || std::env::var("TESTTMP").is_ok() {
            // Use a consistent (and non-existent) terminal to avoid differences in tests.
            hints = hints.term(Some("fake-term".to_string()));
        }
        let caps = Capabilities::new_with_hints(hints)?;
        Ok(Styler {
            level: caps.color_level(),
            renderer: TerminfoRenderer::new(caps),
        })
    }

    /// Initialize a styler using specified color level. Python uses
    /// this for now to keep closer compatibility with the Python
    /// color support detection.
    pub fn from_level(level: ColorLevel) -> termwiz::Result<Styler> {
        let mut hints = ProbeHints::default()
            .color_level(Some(level))
            .force_terminfo_render_to_use_ansi_sgr(Some(enable_compat_mode()));
        if cfg!(test) || std::env::var("TESTTMP").is_ok() {
            // Use a consistent (and non-existent) terminal to avoid differences in tests.
            hints = hints.term(Some("fake-term".to_string()));
        }
        let caps = Capabilities::new_with_hints(hints)?;
        let renderer = TerminfoRenderer::new(caps);
        Ok(Styler { level, renderer })
    }

    pub fn render_bytes(&mut self, style_specs: &str, text: &str) -> termwiz::Result<Vec<u8>> {
        let mut buf: Vec<u8> = Vec::new();
        self.render(&mut buf, style_specs, text)?;
        Ok(buf)
    }

    pub fn render(
        &mut self,
        w: &mut dyn Write,
        style_specs: &str,
        text: &str,
    ) -> termwiz::Result<()> {
        if style_specs.is_empty() {
            w.write_all(text.as_bytes())?;
            return Ok(());
        }

        let mut tty = DumbTty { w };

        // Line breaks within escape sequences don't look right, so
        // process each line's contents separately.
        for (idx, line) in text.split('\n').enumerate() {
            if idx > 0 {
                tty.write_all(b"\n")?;
            }

            if line.is_empty() {
                continue;
            }

            self.renderer.render_to(
                &[
                    Change::AllAttributes(eval_style(self.level, style_specs)),
                    Change::Text(line.to_string()),
                    Change::AllAttributes(CellAttributes::blank()),
                ],
                &mut tty,
            )?;
        }

        Ok(())
    }
}

/// Apply given effects to attrs iff all effects are valid.
fn eval_effects<'a>(
    level: ColorLevel,
    effects: impl IntoIterator<Item = &'a str>,
    attrs: &mut CellAttributes,
) -> bool {
    let mut tentative_attrs = attrs.clone();

    for mut effect in effects {
        if eval_non_color(effect, &mut tentative_attrs) {
            continue;
        }

        let mut is_bg = false;
        if let Some(bg_name) = effect.strip_suffix("_background") {
            effect = bg_name;
            is_bg = true;
        }

        if let Some(ansi) = ansi_color(effect) {
            set_color(is_bg, &mut tentative_attrs, ansi);
            continue;
        }

        if let Some(color_256) = effect.strip_prefix("color") {
            if level == ColorLevel::Sixteen {
                return false;
            }

            if let Ok(idx) = color_256.parse::<u8>() {
                set_color(is_bg, &mut tentative_attrs, ColorSpec::PaletteIndex(idx));
                continue;
            }
        }

        // Supports various 24 bit color formats including our
        // standard #FFF or #FFFFFF.
        let rgb = if effect.starts_with('#') && effect.len() == 4 {
            // termwiz converts #FFF to #F0F0F0, but we want #FFFFFF (a la CSS).
            RgbColor::from_named_or_rgb_string(&format!(
                "#{0}{0}{1}{1}{2}{2}",
                &effect[1..2],
                &effect[2..3],
                &effect[3..4]
            ))
        } else {
            RgbColor::from_named_or_rgb_string(effect)
        };
        if let Some(rgb) = rgb {
            if level != ColorLevel::TrueColor {
                return false;
            }

            set_color(is_bg, &mut tentative_attrs, rgb);
            continue;
        }

        tracing::warn!(effect, "unknown style effect");
        return false;
    }

    // Only updated attrs if all effects were valid.
    *attrs = tentative_attrs;
    true
}

fn eval_non_color(effect: &str, attrs: &mut CellAttributes) -> bool {
    match effect {
        "none" => {
            *attrs = CellAttributes::blank();
        }
        "bold" => {
            attrs.set_intensity(Intensity::Bold);
        }
        "italic" => {
            attrs.set_italic(true);
        }
        "underline" => {
            attrs.set_underline(Underline::Single);
        }
        "inverse" => {
            attrs.set_reverse(true);
        }
        "dim" => {
            attrs.set_intensity(Intensity::Half);
        }
        _ => {
            return false;
        }
    }

    true
}

fn set_color(is_bg: bool, attrs: &mut CellAttributes, color: impl Into<ColorSpec>) {
    if is_bg {
        attrs.set_background(color.into());
    } else {
        attrs.set_foreground(color.into());
    }
}

fn ansi_color(name: &str) -> Option<AnsiColor> {
    Some(match name {
        "black" => AnsiColor::Black,
        "red" => AnsiColor::Maroon,
        "green" => AnsiColor::Green,
        "yellow" => AnsiColor::Olive,
        "blue" => AnsiColor::Navy,
        "magenta" => AnsiColor::Purple,
        "cyan" => AnsiColor::Teal,
        "white" => AnsiColor::Silver,
        "brightblack" => AnsiColor::Grey,
        "brightred" => AnsiColor::Red,
        "brightgreen" => AnsiColor::Lime,
        "brightyellow" => AnsiColor::Yellow,
        "brightblue" => AnsiColor::Blue,
        "brightmagenta" => AnsiColor::Fuchsia,
        "brightcyan" => AnsiColor::Aqua,
        "brightwhite" => AnsiColor::White,
        _ => None?,
    })
}

#[cfg(test)]
mod test {
    use termwiz::caps::Capabilities;
    use termwiz::caps::ColorLevel::*;
    use termwiz::caps::ProbeHints;
    use termwiz::render::terminfo::TerminfoRenderer;
    use termwiz::surface::Change;

    use super::*;

    macro_rules! assert_spec {
        ($level:ident, $specs:tt, $want:expr $(,)?) => {
            let bytes = specs_to_bytes($level, $specs);
            assert_eq!(bytes, $want, $specs);
        };
    }

    #[test]
    fn test_eval_style() {
        // refer to https://en.wikipedia.org/wiki/ANSI_escape_code

        assert_spec!(Sixteen, "green", b"\x1B[32m");
        assert_spec!(Sixteen, "green_background", b"\x1B[42m");
        assert_spec!(
            Sixteen,
            "red blue_background italic underline inverse dim",
            b"\x1B[0m\x1B[2m\x1B[4m\x1B[7m\x1B[3m\x1B[31m\x1B[44m",
        );

        assert_spec!(Sixteen, "color231", b"");
        assert_spec!(Sixteen, "color231:red", b"\x1B[31m");
        assert_spec!(Sixteen, "color231:red:green", b"\x1B[31m");
        assert_spec!(TwoFiftySix, "color231:red:green", b"\x1B[38:5:231m");
        assert_spec!(
            TwoFiftySix,
            "color231_background:red:green",
            b"\x1B[48:5:231m"
        );

        assert_spec!(Sixteen, "bold+#FFF:green", b"\x1B[32m");
        assert_spec!(TwoFiftySix, "bold+#FFF:green", b"\x1B[32m");
        assert_spec!(
            TrueColor,
            "bold+#FFF:green",
            b"\x1B[0m\x1B[1m\x1B[38:2::255:255:255m",
        );
    }

    #[test]
    fn test_render_bytes() {
        let mut styler = Styler::from_level(TwoFiftySix).unwrap();

        assert_eq!(
            styler.render_bytes("red", "hello\nthere\n").unwrap(),
            b"\x1B[31mhello\x1B[39m\n\x1B[31mthere\x1B[39m\n"
        );
    }

    fn specs_to_bytes(level: ColorLevel, specs: &str) -> Vec<u8> {
        let cell_attrs = eval_style(level, specs);
        let mut buf: Vec<u8> = Vec::new();
        let hints = ProbeHints::default()
            .color_level(Some(level))
            // Use a consistent (and non-existent) terminal to avoid differences in tests.
            .term(Some("fake-term".to_string()));
        let caps = Capabilities::new_with_hints(hints).unwrap();
        let mut renderer = TerminfoRenderer::new(caps);
        let mut tty = DumbTty { w: &mut buf };
        renderer
            .render_to(&[Change::AllAttributes(cell_attrs)], &mut tty)
            .unwrap();

        buf
    }
}


/*
 * 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 configmodel::convert::FromConfigValue;
use configmodel::Config;

mod effects;

pub use effects::eval_style;
pub use effects::ColorLevel;
pub use effects::Styler;

enum ColorMode {
    Off,
    Debug,
    Auto,
    Always,
}

impl ColorMode {
    pub fn from_config(val: &str, is_cli_flag: bool) -> Self {
        match val {
            "always" => Self::Always,
            "debug" => Self::Debug,
            "auto" => Self::Auto,

            // Otherwise, truthy CLI flag means "always", truthy config value means "auto".
            _ => match (bool::try_from_str(val).unwrap_or_default(), is_cli_flag) {
                (true, true) => Self::Always,
                (true, false) => Self::Auto,
                (false, _) => Self::Off,
            },
        }
    }
}

pub fn should_color(config: &dyn Config, file: &dyn io::Write) -> bool {
    if hgplain::is_plain(Some("color")) {
        return false;
    }

    let mode = ColorMode::from_config(
        &config.get("ui", "color").unwrap_or_default(),
        config
            .get_sources("ui", "color")
            .last()
            .map_or(false, |s| s.source().as_ref() == "--color"),
    );

    match mode {
        ColorMode::Off | ColorMode::Debug => false,
        ColorMode::Always => {
            // Call can_color() even if mode==always because can_color()
            // has side effects of configuring the terminal to better
            // support colors.
            let _ = file.can_color();
            true
        }
        ColorMode::Auto => file.can_color(),
    }
}


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

pub trait LazyStr {
    fn to_str<'a>(self) -> Cow<'a, str>;
}

impl<F: FnOnce() -> String> LazyStr for F {
    fn to_str<'a>(self) -> Cow<'a, str> {
        self().into()
    }
}

impl LazyStr for &'static str {
    fn to_str<'a>(self) -> Cow<'a, str> {
        self.into()
    }
}

impl LazyStr for String {
    fn to_str<'a>(self) -> Cow<'a, str> {
        self.into()
    }
}

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

    fn use_lazy<'a>(l: impl LazyStr) -> Cow<'a, str> {
        l.to_str()
    }

    fn ignore_lazy(l: impl LazyStr) {
        let _ = l;
    }

    #[test]
    fn test_lazy_str() {
        assert_eq!(use_lazy("foo"), "foo");
        assert_eq!(use_lazy("bar".to_string()), "bar");
        assert_eq!(use_lazy(|| "baz".to_string()), "baz");

        // sanity check laziness works
        ignore_lazy(|| panic!("oops"));
    }
}


/*
 * 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 lz4_pyframe::compress;
use lz4_pyframe::decompress;
use minibench::bench;
use minibench::elapsed;
use rand_core::RngCore;
use rand_core::SeedableRng;

fn main() {
    let mut rng = rand_chacha::ChaChaRng::seed_from_u64(0);
    let mut buf = vec![0u8; 100_000000];
    rng.fill_bytes(&mut buf);
    let compressed = compress(&buf).unwrap();

    bench("compress (100M)", || {
        elapsed(|| {
            compress(&buf).unwrap();
        })
    });

    bench("decompress (~100M)", || {
        elapsed(|| {
            decompress(&compressed).unwrap();
        })
    });
}


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

mod lz4;

pub use lz4::LZ4Error as Error;

pub use crate::lz4::compress;
pub use crate::lz4::compresshc;
pub use crate::lz4::decompress;
pub use crate::lz4::decompress_into;
pub use crate::lz4::decompress_size;
pub type Result<T> = std::result::Result<T, Error>;


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

use std::io::Cursor;

use byteorder::LittleEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use libc::c_int;
use libc::c_void;
use lz4_sys::LZ4StreamDecode;
use lz4_sys::LZ4StreamEncode;
use lz4_sys::LZ4_compressBound;
use lz4_sys::LZ4_compress_continue;
use lz4_sys::LZ4_createStream;
use lz4_sys::LZ4_createStreamDecode;
use lz4_sys::LZ4_decompress_safe_continue;
use lz4_sys::LZ4_freeStream;
use lz4_sys::LZ4_freeStreamDecode;
use thiserror::Error;

use crate::Result;

const HEADER_LEN: usize = 4;

// These function should be exported by lz4-sys. For now, we just declare them.
// See https://github.com/lz4/lz4/blob/dev/lib/lz4hc.h
//
// int LZ4_compress_HC_continue (LZ4_streamHC_t* streamHCPtr, const char* src, char* dst, int
//     srcSize, int maxDstSize);
// LZ4_streamHC_t* LZ4_createStreamHC(void);
// int LZ4_freeStreamHC (LZ4_streamHC_t* streamHCPtr);
#[repr(C)]
struct LZ4_streamHC_t(c_void);
extern "C" {
    fn LZ4_compress_HC_continue(
        LZ4_stream: *mut LZ4_streamHC_t,
        src: *const u8,
        dst: *mut u8,
        srcSize: c_int,
        maxDstSize: c_int,
    ) -> c_int;

    fn LZ4_createStreamHC() -> *mut LZ4_streamHC_t;
    fn LZ4_freeStreamHC(streamHCPtr: *mut LZ4_streamHC_t) -> c_int;
}

#[derive(Debug, Error)]
pub enum LZ4Error {
    #[error("{message:?}")]
    Generic { message: String },

    #[error("{source:?}")]
    Decompression {
        #[from]
        source: LZ4DecompressionError,
    },

    #[error("{source:?}")]
    Io {
        #[from]
        source: std::io::Error,
    },
}

#[derive(Error, Debug)]
#[error(
    "lz4 decompressed data does not match expected length. \
     Expected '{expected:?}' vs Actual '{actual:?}'"
)]
pub struct LZ4DecompressionError {
    expected: usize,
    actual: usize,
}

struct StreamDecoder(pub *mut LZ4StreamDecode);
impl Drop for StreamDecoder {
    fn drop(&mut self) {
        if !self.0.is_null() {
            let error = unsafe { LZ4_freeStreamDecode(self.0) };
            if error != 0 {
                panic!("unable to free stream decoder");
            }
        }
    }
}

struct StreamEncoder(pub *mut LZ4StreamEncode);
impl Drop for StreamEncoder {
    fn drop(&mut self) {
        if !self.0.is_null() {
            let error = unsafe { LZ4_freeStream(self.0) };
            if error != 0 {
                panic!("unable to free stream encoder");
            }
        }
    }
}

/// Read decompressed size from a u32 header.
pub fn decompress_size(data: &[u8]) -> Result<usize> {
    if data.is_empty() {
        Ok(0)
    } else {
        let mut cur = Cursor::new(data);
        let max_decompressed_size = cur.read_u32::<LittleEndian>()? as usize;
        Ok(max_decompressed_size)
    }
}

struct StreamEncoderHC(pub *mut LZ4_streamHC_t);
impl StreamEncoderHC {
    fn new() -> Self {
        StreamEncoderHC(unsafe { LZ4_createStreamHC() })
    }
}
impl Drop for StreamEncoderHC {
    fn drop(&mut self) {
        if !self.0.is_null() {
            let error = unsafe { LZ4_freeStreamHC(self.0) };
            if error != 0 {
                panic!("unable to free stream encoder");
            }
        }
    }
}

/// Decompress into a preallocated buffer. The size of `dest` must
/// match what [decompress_size] returns.
pub fn decompress_into(data: &[u8], dest: &mut [u8]) -> Result<()> {
    let stream = StreamDecoder(unsafe { LZ4_createStreamDecode() });
    if stream.0.is_null() {
        return Err(LZ4Error::Generic {
            message: "Unable to construct lz4 stream decoder".to_string(),
        }
        .into());
    }
    if !dest.is_empty() {
        let data = &data[HEADER_LEN..];
        let source = data.as_ptr();
        let read: i32 = check_error(unsafe {
            LZ4_decompress_safe_continue(
                stream.0,
                source,
                dest.as_mut_ptr() as *mut u8,
                data.len() as i32,
                dest.len() as i32,
            )
        })?;
        if read != dest.len() as i32 {
            return Err(LZ4DecompressionError {
                expected: dest.len(),
                actual: read as usize,
            }
            .into());
        }
    }
    Ok(())
}

pub fn decompress(data: &[u8]) -> Result<Vec<u8>> {
    let max_decompressed_size = decompress_size(data)?;
    if max_decompressed_size == 0 {
        return Ok(Vec::new());
    }
    let mut dest = Vec::<u8>::with_capacity(max_decompressed_size);
    unsafe { dest.set_len(max_decompressed_size) };
    decompress_into(data, &mut dest)?;
    Ok(dest)
}

pub fn compress(data: &[u8]) -> Result<Vec<u8>> {
    let max_compressed_size = (check_error(unsafe { LZ4_compressBound(data.len() as i32) })?
        + HEADER_LEN as i32) as usize;

    let stream = StreamEncoder(unsafe { LZ4_createStream() });
    if stream.0.is_null() {
        return Err(LZ4Error::Generic {
            message: "unable to construct LZ4 stream encoder".to_string(),
        }
        .into());
    }

    let source = data.as_ptr();
    let mut dest = Vec::<u8>::with_capacity(max_compressed_size);
    dest.write_u32::<LittleEndian>(data.len() as u32)?;

    if !data.is_empty() {
        unsafe { dest.set_len(max_compressed_size) };
        let written: i32 = check_error(unsafe {
            LZ4_compress_continue(
                stream.0,
                source,
                dest.as_mut_ptr().add(HEADER_LEN),
                data.len() as i32,
            )
        })?;
        if written < dest.len() as i32 {
            dest.truncate(written as usize + HEADER_LEN);
        }
    }
    Ok(dest)
}

pub fn compresshc(data: &[u8]) -> Result<Vec<u8>> {
    let max_compressed_size = (check_error(unsafe { LZ4_compressBound(data.len() as i32) })?
        + HEADER_LEN as i32) as usize;

    let stream = StreamEncoderHC::new();
    if stream.0.is_null() {
        return Err(LZ4Error::Generic {
            message: "unable to construct LZ4 stream encoder".to_string(),
        }
        .into());
    }

    let source = data.as_ptr();
    let mut dest = Vec::<u8>::with_capacity(max_compressed_size);
    dest.write_u32::<LittleEndian>(data.len() as u32)?;

    if !data.is_empty() {
        unsafe { dest.set_len(max_compressed_size) };
        let written: i32 = check_error(unsafe {
            LZ4_compress_HC_continue(
                stream.0,
                source,
                dest.as_mut_ptr().add(HEADER_LEN),
                data.len() as c_int,
                (max_compressed_size - HEADER_LEN) as c_int,
            )
        })?;
        if written < dest.len() as i32 {
            dest.truncate(written as usize + HEADER_LEN);
        }
    }
    Ok(dest)
}

fn check_error(result: i32) -> Result<i32> {
    if result < 0 {
        return Err(LZ4Error::Generic {
            message: format!("lz4 failed with error '{:?}'", result),
        }
        .into());
    }

    Ok(result)
}

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

    use super::*;

    fn check_roundtrip<T: AsRef<[u8]>>(data: T) -> (Vec<u8>, bool) {
        let data = data.as_ref();
        let compressed = compress(data).unwrap();
        let compressedhc = compresshc(data).unwrap();
        let decompressed = decompress(&compressed).unwrap();
        let decompressedhc = decompress(&compressedhc).unwrap();
        (
            compressed,
            data[..] == decompressed[..] && data[..] == decompressedhc[..],
        )
    }

    #[test]
    fn test_roundtrip() {
        let data = &b"\x00\x01\x02hello world long string easy easy easy easy compress\xF0\xFA"[..];
        let (compressed, roundtrips) = check_roundtrip(&data);
        assert!(compressed.len() < data.len());
        assert!(roundtrips);
    }

    #[test]
    fn test_empty() {
        let data = &b"";
        let (compressed, roundtrips) = check_roundtrip(&data);
        assert_eq!(compressed, &[0u8, 0, 0, 0][..]);
        assert!(roundtrips);
    }

    #[test]
    fn test_short() {
        let data = &b"0"[..];
        let compressed = compress(data).unwrap();

        // Short strings should compress to be longer than they used to be.
        assert!(compressed.len() > data.len());

        // But decompress to their original form.
        let decompressed = decompress(&compressed).unwrap();
        assert_eq!(data, &*decompressed);
    }

    quickcheck! {
        fn test_quickcheck_roundtrip(data: Vec<u8>) -> bool {
            check_roundtrip(&data).1
        }
    }
}


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

//! manifest - The contents of the repository at a specific commit.
//!
//! The history of the repository is recorded in the commit graph. Each commit has a manifest
//! associated with it. The manifest specifies the revision for all the files available in the
//! repository. The file path and file revision are then used to retrieve the contents of the
//! file thus achieving the reconstruction of the entire repository state.

#[cfg(any(test, feature = "for-tests"))]
pub mod testutil;

use anyhow::Result;
use pathmatcher::Matcher;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck_arbitrary_derive::Arbitrary;
use types::HgId;
use types::PathComponentBuf;
use types::RepoPath;
use types::RepoPathBuf;

/// Manifest describes a mapping between file path ([`String`]) and file metadata ([`FileMetadata`]).
/// Fundamentally it is just a Map<file_path, file_metadata>.
///
/// It can be assumed that Manifest interacts with an underlying store for persistance. These
/// interactions may fail due to a variety of reasons. Such failures will be propagated up as Error
/// return statuses.
///
/// Another common failure is passing in a path that the manifest has labeled as a directory. File
/// paths composed of directory names and file names. Querying for paths that the Manifest has
/// determined previously to be directories will result in Errors.
// TODO: Add method for batch modification, takes iterator of added, removed, changed, or
// maybe (path, Option<FileMetadata>) where None signals removal.
// TODO: A batch API allows us to move to having all nodes have a computed hash without losing
// performance. It also allows us to get rid of the flush method.
pub trait Manifest {
    /// Inspects the manifest for the given path. Returns available metadata.
    /// If the path is pointing to an file then Some(FsNodeMetadata::File) is returned with then
    /// file_metadata associated with the file. If the path is poitning to a directory then
    /// Some(FsNodeMetadata::Directory) is returned. If the path is not found then None is
    /// returned.
    // TODO: add default implementation
    fn get(&self, path: &RepoPath) -> Result<Option<FsNodeMetadata>>;

    fn get_ignore_case(&self, path: &RepoPath) -> Result<Option<FsNodeMetadata>>;

    /// Lists the immediate contents of directory in a manifest (non-recursive).
    /// Given a path, the manifest will return:
    /// * List::NotFound when the path is not present in the manifest
    /// * List::File when the path points to a file
    /// * List::Directory when the path points to a directory
    ///    wraps the names of the files and directories in this directory
    // TODO: add default implementation
    fn list(&self, path: &RepoPath) -> Result<List>;

    /// Associates a file path with specific file metadata.
    /// A call with a file path that already exists results in an override or the old metadata.
    fn insert(&mut self, file_path: RepoPathBuf, file_metadata: FileMetadata) -> Result<()>;

    /// Removes a file from the manifest (equivalent to removing it from the repository).
    /// A call with a file path that does not exist in the manifest is a no-op.
    fn remove(&mut self, file_path: &RepoPath) -> Result<Option<FileMetadata>>;

    /// Persists the manifest so that it can be retrieved at a later time. Returns a note
    /// representing the identifier for saved manifest.
    fn flush(&mut self) -> Result<HgId>;

    /// Retrieve the FileMetadata associated with a path.
    /// Paths that were not set will return None.
    fn get_file(&self, file_path: &RepoPath) -> Result<Option<FileMetadata>> {
        let result = self.get(file_path)?.and_then(|fs_hgid| match fs_hgid {
            FsNodeMetadata::File(file_metadata) => Some(file_metadata),
            FsNodeMetadata::Directory(_) => None,
        });
        Ok(result)
    }

    /// Returns an iterator over all the files in the Manifest that satisfy the given Matcher.
    fn files<'a, M: 'static + Matcher + Sync + Send>(
        &'a self,
        matcher: M,
    ) -> Box<dyn Iterator<Item = Result<File>> + 'a>;

    /// Returns an iterator over all directories found in the paths of the files in the Manifest
    /// that satisfy the given Matcher.
    // TODO: add default implementation
    fn dirs<'a, M: 'static + Matcher + Sync + Send>(
        &'a self,
        matcher: M,
    ) -> Box<dyn Iterator<Item = Result<Directory>> + 'a>;

    /// Retuns an iterator of all the differences in files between two Manifest instances of the
    /// same type.
    // TODO: add default implementation
    fn diff<'a, M: Matcher>(
        &'a self,
        other: &'a Self,
        matcher: &'a M,
    ) -> Result<Box<dyn Iterator<Item = Result<DiffEntry>> + 'a>>;

    /// Calculates modified directories between two Manifests.
    ///
    /// A directory is considered "modified" if:
    /// - It only exists in one Manifest.
    /// - It exists in both Manifests. A direct sub-item (file or dir)
    ///   is added, removed, or renamed.
    ///
    /// Similar to when the OS needs to update directory mtime.
    /// Modifying content of a file alone does not count as modifying its
    /// parent directory.
    fn modified_dirs<'a, M: Matcher>(
        &'a self,
        other: &'a Self,
        matcher: &'a M,
    ) -> Result<Box<dyn Iterator<Item = Result<DirDiffEntry>> + 'a>>;
}

/// The result of a list operation. Given a path, the manifest will return:
/// * List::NotFound when the path is not present in the manifest
/// * List::File when the path points to a file
/// * List::Directory when the path points to a directory
// TODO: add FileMetadata to File and "DirectoryMetadata" to Directory
#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum List {
    NotFound,
    File,
    Directory(Vec<(PathComponentBuf, FsNodeMetadata)>),
}

/// FsNodeMetadata short for file system node.
/// The manifest tracks a list of files. However file systems are hierarchical structures
/// composed of directories and files at the end. For different operations it is useful to have
/// a representation for file or directory. A good example is listing a directory. This structure
/// helps us represent that notion.
#[derive(Clone, Copy, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum FsNodeMetadata {
    File(FileMetadata),
    Directory(Option<HgId>),
}

/// A directory entry in a manifest.
///
/// Consists of the full path to the directory. Directories may or may not be assigned
/// identifiers. When an identifier is available it is also returned.
// TODO: Move hgid to a new DirectoryMetadata struct.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Directory {
    pub path: RepoPathBuf,
    pub hgid: Option<HgId>,
}

impl Directory {
    pub fn new(path: RepoPathBuf, hgid: Option<HgId>) -> Self {
        Self { path, hgid }
    }
}

/// A file entry in a manifest.
///
/// Consists of the full path to the file along with the associated file metadata.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct File {
    pub path: RepoPathBuf,
    pub meta: FileMetadata,
}

impl File {
    pub fn new(path: RepoPathBuf, meta: FileMetadata) -> Self {
        Self { path, meta }
    }
}

/// The contents of the Manifest for a file.
/// * hgid: used to determine the revision of the file in the repository.
/// * file_type: determines the type of the file.
#[derive(Clone, Copy, Debug, Default, Ord, PartialOrd, Eq, PartialEq, Hash)]
#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
pub struct FileMetadata {
    pub hgid: HgId,
    pub file_type: FileType,
}

/// The types of files (leaf nodes in a tree).
///
/// The type needs to round-trip tree serialization.
#[derive(Clone, Copy, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum FileType {
    /// Regular files.
    Regular,
    /// Executable files. Like Regular files but with the executable flag set.
    Executable,
    /// Symlinks. Their targets are not limited to repository paths. They can point anywhere.
    Symlink,
    /// Git submodule. It's up to the higher layer to decide what to do with them.
    GitSubmodule,
}

impl Default for FileType {
    fn default() -> Self {
        FileType::Regular
    }
}

impl FileMetadata {
    pub fn new(hgid: HgId, file_type: FileType) -> Self {
        Self { hgid, file_type }
    }

    /// Creates `FileMetadata` with file_type set to `FileType::Regular`.
    pub fn regular(hgid: HgId) -> Self {
        Self::new(hgid, FileType::Regular)
    }

    /// Creates `FileMetadata` with file_type set to `FileType::Executable`.
    pub fn executable(hgid: HgId) -> Self {
        Self::new(hgid, FileType::Executable)
    }

    /// Creates `FileMetadata` with file_type set to `FileType::Symlink`.
    pub fn symlink(hgid: HgId) -> Self {
        Self::new(hgid, FileType::Symlink)
    }
}

/// Represents a file that is different between two tree manifests.
#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct DiffEntry {
    pub path: RepoPathBuf,
    pub diff_type: DiffType,
}

#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct DirDiffEntry {
    /// Path of the directory.
    pub path: RepoPathBuf,
    /// Exist on the left side.
    pub left: bool,
    /// Exist on the right side.
    pub right: bool,
}

impl DiffEntry {
    pub fn new(path: RepoPathBuf, diff_type: DiffType) -> Self {
        DiffEntry { path, diff_type }
    }

    pub fn left(file: File) -> Self {
        Self::new(file.path, DiffType::LeftOnly(file.meta))
    }

    pub fn right(file: File) -> Self {
        Self::new(file.path, DiffType::RightOnly(file.meta))
    }

    pub fn changed(left: File, right: File) -> Self {
        debug_assert!(left.path == right.path);
        Self::new(left.path, DiffType::Changed(left.meta, right.meta))
    }
}

#[derive(Clone, Copy, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum DiffType {
    LeftOnly(FileMetadata),
    RightOnly(FileMetadata),
    Changed(FileMetadata, FileMetadata),
}

impl DiffType {
    /// Returns the metadata of the file in the left manifest when it exists.
    pub fn left(&self) -> Option<FileMetadata> {
        match self {
            DiffType::LeftOnly(left_metadata) => Some(*left_metadata),
            DiffType::RightOnly(_) => None,
            DiffType::Changed(left_metadata, _) => Some(*left_metadata),
        }
    }

    /// Returns the metadata of the file in the right manifest when it exists.
    pub fn right(&self) -> Option<FileMetadata> {
        match self {
            DiffType::LeftOnly(_) => None,
            DiffType::RightOnly(right_metadata) => Some(*right_metadata),
            DiffType::Changed(_, right_metadata) => Some(*right_metadata),
        }
    }
}

#[cfg(any(test, feature = "for-tests"))]
impl quickcheck::Arbitrary for FileType {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        g.choose(&[FileType::Regular, FileType::Executable, FileType::Symlink])
            .unwrap()
            .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.
 */

use types::testutil::*;

use crate::File;
use crate::FileMetadata;

pub fn make_meta(hex: &str) -> FileMetadata {
    FileMetadata::regular(hgid(hex))
}

pub fn make_file(path: &str, hex: &str) -> File {
    File {
        path: repo_path_buf(path),
        meta: make_meta(hex),
    }
}


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

use std::sync::Arc;

use manifest::FileMetadata;
use manifest::Manifest;
use manifest_tree::testutil::*;
use manifest_tree::TreeManifest;
use manifest_tree::TreeStore;
use minibench::bench;
use minibench::elapsed;
use pathmatcher::AlwaysMatcher;
use quickcheck::Arbitrary;
use quickcheck::StdGen;
use rand::SeedableRng;
use rand_chacha::ChaChaRng;
use types::testutil::generate_repo_paths;
use types::HgId;
use types::RepoPathBuf;

const INIT_SET_COUNT: usize = 4_000_000;
const OP_COUNT: usize = 1_000_000;

// See https://github.com/rust-lang/rust/issues/64102
pub fn black_box<T>(dummy: T) -> T {
    unsafe {
        let ret = std::ptr::read_volatile(&dummy);
        std::mem::forget(dummy);
        ret
    }
}

pub fn generate_entries<G: quickcheck::Gen>(
    entry_count: usize,
    qc_gen: &mut G,
) -> Vec<(RepoPathBuf, FileMetadata)> {
    let repo_paths = generate_repo_paths(entry_count, qc_gen);
    let mut result = Vec::with_capacity(entry_count);
    for path in repo_paths.into_iter() {
        let hgid = HgId::arbitrary(qc_gen);
        result.push((path, FileMetadata::regular(hgid)));
    }
    result
}

pub fn finalize(
    store: &TestStore,
    manifest: &mut TreeManifest,
    parent_manifests: Vec<&TreeManifest>,
) -> HgId {
    let mut manifest_id = Default::default();
    for (path, hgid, raw, _, _) in manifest.finalize(parent_manifests).unwrap() {
        store.insert(&path, hgid, raw).unwrap();
        if path.is_empty() {
            manifest_id = hgid;
        }
    }
    manifest_id
}

// Run with: cargo bench --features for-tests
// In this benchmark we focus on the `finalize` model of writing to storage.
// This model required the caller to write the data to a store that they have.
fn main() {
    let rng = ChaChaRng::from_seed([0u8; 32]);
    let mut qc_gen = StdGen::new(rng, 10);
    let store = Arc::new(TestStore::new());
    let entries = generate_entries(INIT_SET_COUNT + OP_COUNT, &mut qc_gen);
    let initial_entries = &entries[..INIT_SET_COUNT];
    let op_entries = &entries[INIT_SET_COUNT..];
    let mut initial_manifest = TreeManifest::ephemeral(store.clone());
    for (path, file_metadata) in initial_entries.iter() {
        initial_manifest
            .insert(path.to_owned(), *file_metadata)
            .unwrap();
    }
    println!("initial_entry_cnt = {}", initial_entries.len());
    // Iterate through the tree before it gets committed to storage
    bench("iterate_files_ephemeral", || {
        // In this block we rely on the fact that we have access to the initial manifest before
        // it is finalized. This is an optimization to redo file insertion.
        elapsed(|| {
            for file in initial_manifest.files(&AlwaysMatcher::new()) {
                black_box(file).unwrap();
            }
        })
    });
    let initial_manifest_id = finalize(&store, &mut initial_manifest, vec![]);

    // Iterate through the durable entries that are all loaded in memory
    bench("iterate_files_durable_in_memory", || {
        // In this block we rely on the fact that we did not take the initial manifest out of
        // memory. This is just an optimization to avoid reloading a manifest.
        let manifest = initial_manifest.clone();
        elapsed(|| {
            for file in manifest.files(&AlwaysMatcher::new()) {
                black_box(file).unwrap();
            }
        })
    });

    // Iterate through the durable entries while loading them from storage
    bench("iterate_files_durable_load", || {
        let manifest = TreeManifest::durable(store.clone(), initial_manifest_id);
        elapsed(|| {
            for file in manifest.files(&AlwaysMatcher::new()) {
                black_box(file).unwrap();
            }
        })
    });

    // Execute OP_COUNT insertions.
    bench("insert", || {
        let mut manifest = initial_manifest.clone();
        elapsed(|| {
            for (path, file_metadata) in op_entries.iter() {
                manifest.insert(path.to_owned(), *file_metadata).unwrap();
            }
        })
    });

    // Finalize a tree with OP_COUNT new leaves.
    bench("finalize", || {
        let mut manifest = initial_manifest.clone();
        for (path, file_metadata) in op_entries.iter() {
            manifest.insert(path.to_owned(), *file_metadata).unwrap();
        }
        elapsed(|| {
            for x in manifest.finalize(vec![&initial_manifest]).unwrap() {
                black_box(x);
            }
        })
    });

    // Remove the previously added files.
    bench("remove", || {
        let mut manifest = initial_manifest.clone();
        elapsed(|| {
            for (path, _) in initial_entries.iter().take(OP_COUNT) {
                manifest.remove(path).unwrap();
            }
        })
    });
}