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

use std::time::Duration;

#[allow(unused_variables)]
pub async fn try_terminate_process_gracefully(pid: i32, timeout: Duration) -> anyhow::Result<()> {
    #[cfg(unix)]
    {
        unix::try_terminate_process_gracefully(pid, timeout).await
    }

    #[cfg(not(unix))]
    {
        Err(anyhow::anyhow!(
            "Graceful process termination is not implemented for non-unix target family."
        ))
    }
}

#[cfg(unix)]
mod unix {
    use std::ops::ControlFlow;
    use std::time::Duration;
    use std::time::Instant;

    use nix::sys::signal::Signal;
    use tracing::info;
    use tracing::warn;

    pub(crate) async fn try_terminate_process_gracefully(
        pid: i32,
        timeout: Duration,
    ) -> anyhow::Result<()> {
        let pid = nix::unistd::Pid::from_raw(pid);
        match send_signal_and_wait_for_shutdown(Signal::SIGTERM, pid, Some(timeout)).await {
            ControlFlow::Continue(_) => {}
            ControlFlow::Break(StoppedWaiting::Success) => {
                return Ok(());
            }
            ControlFlow::Break(StoppedWaiting::UnexpectedError(e)) => {
                return Err(e);
            }
        }
        warn!(
            "Failed to gracefully terminate process `{}`, sending SIGKILL.",
            pid,
        );
        match send_signal_and_wait_for_shutdown(Signal::SIGKILL, pid, None).await {
            ControlFlow::Continue(_) => {
                unreachable!("Not expected when there is no timeout specified.")
            }
            ControlFlow::Break(StoppedWaiting::Success) => Ok(()),
            ControlFlow::Break(StoppedWaiting::UnexpectedError(e)) => Err(e),
        }
    }

    enum StoppedWaiting {
        Success,
        UnexpectedError(anyhow::Error),
    }

    fn check_result(pid: nix::unistd::Pid, result: nix::Result<()>) -> ControlFlow<StoppedWaiting> {
        match result {
            // Signal is sent successfully, now we need to wait for a process to terminate.
            Ok(_) => ControlFlow::Continue(()),
            // There is no such process, our desired outcome.
            Err(nix::errno::Errno::ESRCH) => ControlFlow::Break(StoppedWaiting::Success),
            Err(e) => ControlFlow::Break(StoppedWaiting::UnexpectedError(anyhow::anyhow!(
                "Unexpected error while waiting for process `{}` to terminate (`{}`)",
                pid,
                e
            ))),
        }
    }

    async fn send_signal_and_wait_for_shutdown(
        signal: nix::sys::signal::Signal,
        pid: nix::unistd::Pid,
        timeout: Option<Duration>,
    ) -> ControlFlow<StoppedWaiting> {
        let start = Instant::now();
        info!("Sending signal `{}` to process `{}`.", signal, pid,);
        // Actually send a signal.
        check_result(pid, nix::sys::signal::kill(pid, signal))?;
        loop {
            if let Some(timeout) = timeout {
                if start.elapsed() > timeout {
                    break;
                }
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
            // Just check the process, signal already sent.
            check_result(pid, nix::sys::signal::kill(pid, None))?;
        }
        ControlFlow::Continue(())
    }
}

#[cfg(unix)]
#[cfg(test)]
mod tests {
    use std::process::Stdio;

    use tempfile;
    use tokio::fs::File;
    use tokio::io::AsyncBufReadExt;
    use tokio::io::AsyncWriteExt;
    use tokio::io::BufReader;

    use super::*;

    #[tokio::test]
    async fn test_graceful_termination() -> anyhow::Result<()> {
        let dir = tempfile::tempdir()?;

        // We cannot start child script directly else it will enter a zombie state
        // when trying to kill it and termination would be stuck.
        let mut parent_file = File::create(dir.path().join("parent.sh")).await?;
        let parent_script = r#"
            sh child.sh &
            while [ ! -f trap_registered ]; do sleep 1; done # wait for a child script to register a trap
            echo $!
        "#;
        parent_file.write_all(parent_script.as_bytes()).await?;

        let mut child_file = File::create(dir.path().join("child.sh")).await?;
        let child_script = r#"
            sigterm()
            {
                touch gracefully_terminated
            }
            trap sigterm TERM
            touch trap_registered

            sleep 100 &
            wait
        "#;
        child_file.write_all(child_script.as_bytes()).await?;

        let mut command = tokio::process::Command::new("sh");
        command
            .args(["parent.sh"])
            .current_dir(dir.path())
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::null());

        let mut parent = command.spawn()?;
        parent.wait().await?;

        let pid = {
            let stdout = parent.stdout.take().unwrap();
            let mut reader = BufReader::new(stdout).lines();
            let pid_str = reader.next_line().await?;
            pid_str.unwrap().parse().unwrap()
        };

        try_terminate_process_gracefully(pid, Duration::from_secs(1)).await?;

        assert!(tokio::fs::try_exists(dir.path().join("gracefully_terminated")).await?);

        Ok(())
    }

    #[tokio::test]
    async fn test_sigkill_when_graceful_termination_fails() -> anyhow::Result<()> {
        let dir = tempfile::tempdir()?;

        // We cannot start child script directly else it will enter a zombie state
        // when trying to kill it and termination would be stuck.
        let mut parent_file = File::create(dir.path().join("parent.sh")).await?;
        let parent_script = r#"
            sh child.sh &
            while [ ! -f trap_registered ]; do sleep 1; done # wait for a child script to register a trap
            echo $!
        "#;
        parent_file.write_all(parent_script.as_bytes()).await?;

        let mut child_file = File::create(dir.path().join("child.sh")).await?;
        let child_script = r#"
            sigterm()
            {
                sleep 3  # block termination
                touch gracefully_terminated
            }
            trap sigterm TERM
            touch trap_registered

            sleep 100 &
            wait
        "#;
        child_file.write_all(child_script.as_bytes()).await?;

        let mut command = tokio::process::Command::new("sh");
        command
            .args(["parent.sh"])
            .current_dir(dir.path())
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::null());

        let mut parent = command.spawn()?;
        parent.wait().await?;

        let pid = {
            let stdout = parent.stdout.take().unwrap();
            let mut reader = BufReader::new(stdout).lines();
            let pid_str = reader.next_line().await?;
            pid_str.unwrap().parse().unwrap()
        };

        try_terminate_process_gracefully(pid, Duration::from_secs(1)).await?;

        assert!(!tokio::fs::try_exists(dir.path().join("gracefully_terminated")).await?);

        Ok(())
    }
}


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

//! Common core components of buck2

// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]
#![feature(fs_try_exists)]
#![feature(io_error_more)]
#![feature(is_sorted)]
#![feature(trait_alias)]
#![feature(never_type)]
#![feature(pattern)]

#[cfg(test)]
#[macro_use]
extern crate maplit;

pub mod buckd_connection;
pub mod cas_digest;
pub mod client_utils;
pub mod convert;
pub mod daemon_dir;
pub mod dice;
pub mod error_report;
pub mod events;
pub mod executor_config;
pub mod external_symlink;
pub mod file_ops;
pub mod find_buildfile;
pub mod home_buck_tmp;
pub mod http;
pub mod ignores;
pub mod invocation_paths;
pub mod invocation_roots;
pub mod io;
pub mod kill_util;
pub mod legacy_configs;
pub mod liveliness_observer;
pub mod local_resource_state;
pub mod memory;
pub mod package_boundary;
pub mod package_listing;
pub mod pattern;
pub mod result;
pub mod sqlite;
pub mod target_aliases;
pub mod temp_path;


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

use std::sync::Arc;

use async_trait::async_trait;
use dupe::Dupe;
use thiserror::Error;
use tokio::sync::OwnedRwLockWriteGuard;
use tokio::sync::RwLock;

#[derive(Debug, Error, Copy, Clone, Dupe)]
#[error("LivelinessObserver reports this session is shutting down")]
struct NotAlive;

/// A LivelinessObserver can be passed to notify callees that they should stop work and return
/// early.
///
/// LivelinessObserver exposes a single method: `while_alive`. This method returns a future: when
/// the future resolves, the callee should stop its work and return to the caller.
///
/// This is usually implemented by using `select` to poll both the `while_alive` future and an
/// actual work future. Users should be mindful of not accidentally dropping handles to ongoing
/// work by dropping the "cancelled" side of the select.
///
/// How exactly the callee reports to the caller that it interrupted its work is not a concern of
/// this trait.
#[async_trait]
pub trait LivelinessObserver: Send + Sync {
    /// Pending while we are alive. Ready when we aren't.
    async fn while_alive(&self);
}

impl dyn LivelinessObserver {
    pub async fn is_alive(&self) -> bool {
        futures::poll!(self.while_alive()).is_pending()
    }
}

/// A LivelinessObserver with an implementation backed by an RW Lock. While the lock is held with
/// write access, this LivelinessObserver is alive.
///
/// The way it works is as follows:
///
/// - The LivelinessGuard holds a RW lock with write access.
/// - The `while_alive()` implementation attempts to acquire the lock with read access.
/// - This means that while the guard hasn't been dropped, the manager is considered alive.
///
/// We also allow the LivelinessGuard be "forgotten", which drops it but still forces `while_alive`
/// to stay pending (we do this by tracking this with a state flag in the mutex).
///
/// In an ideal world, this would be a newtype, but that means it needs to contain an `Arc<RwLock>`
/// to support `try_write_owned()`, and now we have 2 Arcs unnecessarily.
type LivelinessObserverForGuard = RwLock<LivelinessObserverState>;

pub struct LivelinessGuard {
    guard: OwnedRwLockWriteGuard<LivelinessObserverState>,

    // A reference to the underlying manager to support `cancel`.
    manager: Arc<LivelinessObserverForGuard>,
}

impl LivelinessGuard {
    pub fn create() -> (Arc<dyn LivelinessObserver>, LivelinessGuard) {
        let manager = Arc::new(LivelinessObserverForGuard::new(
            LivelinessObserverState::AliveWhenLocked,
        ));

        let guard = manager
            .dupe()
            .try_write_owned()
            .expect("This lock was just created");

        (manager.dupe() as _, LivelinessGuard { guard, manager })
    }

    /// Declare that this liveliness manager is no longer alive. Dropping the guard does the same,
    /// but this allows potentially restoring it later.
    pub fn cancel(self) -> CancelledLivelinessGuard {
        CancelledLivelinessGuard {
            manager: self.manager,
        }
    }

    /// Declare that the underlying liveliness manager should stay alive forever, even when we drop
    /// this guard.
    pub fn forget(mut self) {
        *self.guard = LivelinessObserverState::ForeverAlive;
    }
}

#[derive(Debug)]
pub struct CancelledLivelinessGuard {
    manager: Arc<LivelinessObserverForGuard>,
}

impl CancelledLivelinessGuard {
    /// If the lock is available, re-acquire it, thus allowing things to stay alive again.
    ///
    /// Note that callees may have already observed the "not alive" state.
    ///
    /// TODO: We should probably remove this API: it's a bit of a footgun.
    pub fn restore(self) -> Option<LivelinessGuard> {
        let guard = self.manager.dupe().try_write_owned().ok()?;
        Some(LivelinessGuard {
            guard,
            manager: self.manager,
        })
    }
}

/// The state of this liveliness manager. By default, it's alive if and only if the LivelinessGuard
/// is holding the lock, but if `forget` was called, then we'll record that even upon a successful
/// lock acquisition, the manager should be considered alive.
#[derive(Debug)]
enum LivelinessObserverState {
    AliveWhenLocked,
    ForeverAlive,
}

#[async_trait]
impl LivelinessObserver for LivelinessObserverForGuard {
    async fn while_alive(&self) {
        match *self.read().await {
            LivelinessObserverState::AliveWhenLocked => {}
            LivelinessObserverState::ForeverAlive => futures::future::pending().await,
        }
    }
}

/// Always alive.
pub struct NoopLivelinessObserver;

impl NoopLivelinessObserver {
    pub fn create() -> Arc<dyn LivelinessObserver> {
        Arc::new(Self) as _
    }
}

#[async_trait]
impl LivelinessObserver for NoopLivelinessObserver {
    async fn while_alive(&self) {
        futures::future::pending().await
    }
}

#[async_trait]
impl LivelinessObserver for Arc<dyn LivelinessObserver> {
    async fn while_alive(&self) {
        self.as_ref().while_alive().await
    }
}

pub struct LivelinessAnd<A, B> {
    a: A,
    b: B,
}

#[async_trait]
impl<A, B> LivelinessObserver for LivelinessAnd<A, B>
where
    A: LivelinessObserver,
    B: LivelinessObserver,
{
    async fn while_alive(&self) {
        let a = self.a.while_alive();
        let b = self.b.while_alive();
        futures::pin_mut!(a);
        futures::pin_mut!(b);
        futures::future::select(a, b).await;
    }
}

pub trait LivelinessObserverExt: Sized {
    fn and<B>(self, b: B) -> LivelinessAnd<Self, B>;
}

impl<T> LivelinessObserverExt for T
where
    T: LivelinessObserver + Sized,
{
    fn and<B>(self, b: B) -> LivelinessAnd<Self, B> {
        LivelinessAnd { a: self, b }
    }
}

#[async_trait]
impl LivelinessObserver for more_futures::cancellable_future::CancellationObserver {
    async fn while_alive(&self) {
        self.dupe().await
    }
}

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

    #[tokio::test]
    async fn test_guard_is_alive() {
        let (manager, guard) = LivelinessGuard::create();
        assert!(manager.is_alive().await);
        drop(guard);
        assert!(!manager.is_alive().await);
    }

    #[tokio::test]
    async fn test_and() {
        let (manager_a, guard) = LivelinessGuard::create();
        let manager_b = NoopLivelinessObserver::create();

        let manager = manager_a.and(manager_b);
        let manager = &manager as &dyn LivelinessObserver;

        assert!(manager.is_alive().await);
        drop(guard);
        assert!(!manager.is_alive().await);
    }

    #[tokio::test]
    async fn test_cancel_restore_forget() {
        let (manager, guard) = LivelinessGuard::create();
        assert!(manager.is_alive().await);

        let cancelled = guard.cancel();
        assert!(!manager.is_alive().await);

        let restored = cancelled.restore().expect("This is not currently held");
        assert!(manager.is_alive().await);

        restored.forget();
        assert!(manager.is_alive().await);
    }
}


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

use std::sync::Arc;

use buck2_core::target::label::ConfiguredTargetLabel;
use derivative::Derivative;
use tokio::sync::mpsc;
use tokio::sync::mpsc::UnboundedReceiver;
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::Mutex;

#[derive(Debug, PartialEq)]
pub struct EnvironmentVariable {
    pub key: String,
    pub value: String,
}

/// Resource represented by a list of environment variable key-value pairs.
#[derive(Debug, PartialEq)]
pub struct LocalResource(pub Vec<EnvironmentVariable>);

/// RAII handle for resource spec, returns spec to the pool on drop.
pub struct LocalResourceHolder {
    // Optionality is only needed so we can move out the spec on drop.
    spec: Option<LocalResource>,
    sender: UnboundedSender<LocalResource>,
}

impl Drop for LocalResourceHolder {
    fn drop(&mut self) {
        let _ignored = self.sender.send(
            self.spec
                .take()
                .expect("Should only be absent in already dropped object."),
        );
    }
}

impl AsRef<LocalResource> for LocalResourceHolder {
    fn as_ref(&self) -> &LocalResource {
        self.spec
            .as_ref()
            .expect("Should only be absent in already dropped object.")
    }
}

/// Blocking resource pool to manage access to prepared local resources.
#[derive(Clone, Debug, Derivative)]
#[derivative(PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct LocalResourceState {
    // Set of resources of same type should be uniquely identified by configured target label providing `LocalResourceInfo`.
    // This is the assumption for equiality, ordering and hash implementations.
    source_target: ConfiguredTargetLabel,
    #[derivative(
        Hash = "ignore",
        PartialEq = "ignore",
        PartialOrd = "ignore",
        Ord = "ignore"
    )]
    owning_pid: Option<i32>,
    #[derivative(
        Hash = "ignore",
        PartialEq = "ignore",
        PartialOrd = "ignore",
        Ord = "ignore"
    )]
    sender: UnboundedSender<LocalResource>,
    #[derivative(
        Hash = "ignore",
        PartialEq = "ignore",
        PartialOrd = "ignore",
        Ord = "ignore"
    )]
    receiver: Arc<Mutex<UnboundedReceiver<LocalResource>>>,
}

impl LocalResourceState {
    pub fn new(
        source_target: ConfiguredTargetLabel,
        owning_pid: Option<i32>,
        specs: Vec<LocalResource>,
    ) -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();
        for spec in specs {
            sender.send(spec).expect(
                "Not expected send to fail when channel is not closed and receiver is not dropped.",
            );
        }
        LocalResourceState {
            source_target,
            owning_pid,
            sender,
            receiver: Arc::new(Mutex::new(receiver)),
        }
    }

    pub fn source_target(&self) -> &ConfiguredTargetLabel {
        &self.source_target
    }

    /// ID of process which actually is holding a pool of resources.
    /// SIGTERM is sent to this process to free the resources.
    pub fn owning_pid(&self) -> Option<i32> {
        self.owning_pid
    }

    pub async fn acquire_resource(&self) -> LocalResourceHolder {
        let spec = {
            let mut guard = self.receiver.lock().await;
            Some(guard.recv().await.unwrap())
        };
        LocalResourceHolder {
            spec,
            sender: self.sender.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use buck2_core::configuration::data::ConfigurationData;
    use buck2_core::target::label::ConfiguredTargetLabel;

    use super::EnvironmentVariable;
    use crate::local_resource_state::LocalResource;
    use crate::local_resource_state::LocalResourceState;

    #[tokio::test]
    async fn test_canary() -> anyhow::Result<()> {
        let target =
            ConfiguredTargetLabel::testing_parse("foo//bar:baz", ConfigurationData::testing_new());
        let specs = vec![
            LocalResource(vec![EnvironmentVariable {
                key: "FOO".to_owned(),
                value: "foo".to_owned(),
            }]),
            LocalResource(vec![EnvironmentVariable {
                key: "BAR".to_owned(),
                value: "bar".to_owned(),
            }]),
        ];

        let state = LocalResourceState::new(target, Some(0), specs);
        let handle = tokio::spawn(async move {
            {
                let _holder1 = state.acquire_resource().await;
                let _holder2 = state.acquire_resource().await;
            }
            for _ in 0..10 {
                let _x = state.acquire_resource().await;
            }
        });
        handle.await?;
        Ok(())
    }
}


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

//! Utilities for interacting with the jemalloc heap used by buck2.
//!
//! In order to make use of jemalloc's heap dump or profiling utilities, you must set the MALLOC_CONF environment
//! variable to a suitable value prior to launching the daemon, such as:
//!  `export MALLOC_CONF=prof:true,prof_final:false,prof_prefix:/tmp/jeprof`
//! This turns on the profiler (`prof:true`), tells the profile to not take heap dump on process exit,
//! (`prof_final:false`), and write dumps to `/tmp/jeprof` if a file argument isn't given to `mallctl`
//! (`prof_prefix:/tmp/jeprof`).

#[cfg(fbcode_build)]
mod imp {
    use std::env;

    /// Output the current state of the heap to the filename specified.
    /// Intended to be used for debugging purposes.
    /// Requires MALLOC_CONF=prof:true to be set in environment variables
    /// when run, though must be built without MALLOC_CONF=prof:true.
    pub fn write_heap_to_file(filename: &str) -> anyhow::Result<()> {
        if !memory::is_using_jemalloc() {
            return Err(anyhow::anyhow!(
                "not using jemalloc; are you building with @mode/dev or @mode/dbgo?"
            ));
        }

        let prof_enabled: bool = memory::mallctl_read("opt.prof")?;
        if !prof_enabled {
            if env::var_os("MALLOC_CONF").is_some() {
                return Err(anyhow::anyhow!(
                    "the environment variable MALLOC_CONF is set, but profiling is not enabled. MALLOC_CONF must contain prof:true to enable the profiler"
                ));
            }

            return Err(anyhow::anyhow!(
                "profiling is not enabled for this process; you must set the environment variable MALLOC_CONF to contain at least prof:true in order to profile"
            ));
        }

        eprintln!("dumping heap to: {:?}", filename);
        memory::mallctl_write("prof.dump", filename)?;
        Ok(())
    }

    /// Dump allocator stats from JEMalloc. Intended for debug purposes
    pub fn allocator_stats(options: &str) -> anyhow::Result<String> {
        allocator_stats::malloc_stats(options)
    }

    /// Enables background threads for jemalloc. See [here](http://jemalloc.net/jemalloc.3.html#background_thread) for
    /// jemalloc documentation. When set, jemalloc will use background threads to purge unused dirty pages instead of
    /// doing it synchronously.
    ///
    /// This function has no effect if not using jemalloc.
    pub fn enable_background_threads() -> anyhow::Result<()> {
        if memory::is_using_jemalloc() {
            memory::mallctl_write("background_thread", true)?;
        }
        Ok(())
    }

    pub fn has_jemalloc_stats() -> bool {
        memory::is_using_jemalloc()
    }
}

#[cfg(not(fbcode_build))]
mod imp {

    pub fn write_heap_to_file(_filename: &str) -> anyhow::Result<()> {
        // TODO(swgillespie) the `jemalloc_ctl` crate is probably capable of doing this
        // and we already link against it
        Err(anyhow::anyhow!(
            "not implemented: heap dump for Cargo builds"
        ))
    }

    pub fn allocator_stats(_: &str) -> anyhow::Result<String> {
        Err(anyhow::anyhow!(
            "not implemented: allocator stats  for Cargo builds"
        ))
    }

    pub fn enable_background_threads() -> anyhow::Result<()> {
        Ok(())
    }

    pub fn has_jemalloc_stats() -> bool {
        false
    }
}

pub use imp::allocator_stats;
pub use imp::enable_background_threads;
pub use imp::has_jemalloc_stats;
pub use imp::write_heap_to_file;


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

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

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::name::CellName;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::cells::paths::CellRelativePathBuf;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use ref_cast::RefCast;

use crate::legacy_configs::dice::HasLegacyConfigs;
use crate::legacy_configs::LegacyBuckConfigs;
use crate::result::SharedResult;

#[derive(PartialEq, Allocative)]
pub struct PackageBoundaryExceptions(HashMap<CellName, CellPackageBoundaryExceptions>);

#[derive(PartialEq, Allocative)]
struct CellPackageBoundaryExceptions {
    // The reason we avoid a trie is that there's not a convenient `TrieSet` implementation to use,
    // and tries will likely have worse performance because most exception paths are very short.
    // Instead, we use a HashMap of first directory of the path to the rest of the path.
    prefix_to_subpaths: HashMap<FileNameBuf, Vec<ForwardRelativePathBuf>>,
    // Sometimes we want to say everything is allowed
    allow_everything: bool,
}

impl CellPackageBoundaryExceptions {
    fn new(s: &str) -> anyhow::Result<Self> {
        let mut prefix_to_subpaths = HashMap::new();
        let mut allow_everything = false;
        for path_str in s.split(',') {
            let path_str = path_str.trim();

            if path_str == "." {
                allow_everything = true;
            } else {
                let path = ForwardRelativePath::new(path_str)?;
                // path.split_first() only returns None if the path is empty.
                // In the case of the buckconfig `project.package_boundary_exceptions`,
                // we only get an empty path if there is an extra newline, in which case
                // we can just ignore.
                if let Some((prefix, subpath)) = path.split_first() {
                    let subpaths = prefix_to_subpaths
                        .entry(prefix.to_owned())
                        .or_insert_with(Vec::new);
                    subpaths.push(subpath.to_owned());
                }
            }
        }
        Ok(Self {
            prefix_to_subpaths,
            allow_everything,
        })
    }

    fn get_package_boundary_exception_path(
        &self,
        path: &CellRelativePath,
    ) -> Option<CellRelativePathBuf> {
        if self.allow_everything {
            return Some(CellRelativePathBuf::unchecked_new("".to_owned()));
        }
        let path: &ForwardRelativePath = path.as_ref();
        if let Some((package_prefix, package_subpath)) = path.split_first() {
            if let Some(subpaths) = self.prefix_to_subpaths.get(package_prefix) {
                return subpaths
                    .iter()
                    .find(|p| package_subpath.starts_with(p))
                    .map(|p| {
                        CellRelativePath::new(<&ForwardRelativePath>::from(package_prefix)).join(p)
                    });
            }
        }
        None
    }
}

impl PackageBoundaryExceptions {
    fn new(configs: &LegacyBuckConfigs) -> anyhow::Result<Self> {
        Ok(Self(
            configs
                .iter()
                .filter_map(|(name, cell_configs)| {
                    cell_configs
                        .get("project", "package_boundary_exceptions")
                        .map(|v| {
                            let e = CellPackageBoundaryExceptions::new(v).with_context(
                                || format!("Error parsing `project.package_boundary_exceptions` key from cell `{}`", name)
                            )?;
                            Ok((name, e))
                        })
                })
                .collect::<anyhow::Result<_>>()?,
        ))
    }

    /// Returns the package boundary exception path that covers this path, if it exists
    pub fn get_package_boundary_exception_path(&self, path: &CellPath) -> Option<CellPath> {
        if let Some(exceptions) = self.0.get(&path.cell()) {
            exceptions
                .get_package_boundary_exception_path(path.path())
                .map(|p| CellPath::new(path.cell(), p))
        } else {
            None
        }
    }

    pub fn contains(&self, path: &CellPath) -> bool {
        self.get_package_boundary_exception_path(path).is_some()
    }
}

#[async_trait]
pub trait HasPackageBoundaryExceptions {
    async fn get_package_boundary_exceptions(&self)
    -> SharedResult<Arc<PackageBoundaryExceptions>>;

    async fn get_package_boundary_exception(
        &self,
        path: CellPathRef<'async_trait>,
    ) -> SharedResult<bool>;
}

#[async_trait]
impl HasPackageBoundaryExceptions for DiceComputations {
    async fn get_package_boundary_exceptions(
        &self,
    ) -> SharedResult<Arc<PackageBoundaryExceptions>> {
        #[derive(Hash, Eq, PartialEq, Clone, Dupe, Display, Debug, Allocative)]
        #[display(fmt = "{:?}", self)]
        struct PackageBoundaryExceptionsKey;

        #[async_trait]
        impl Key for PackageBoundaryExceptionsKey {
            type Value = SharedResult<Arc<PackageBoundaryExceptions>>;

            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellations: &CancellationContext,
            ) -> Self::Value {
                Ok(Arc::new(PackageBoundaryExceptions::new(
                    &ctx.get_legacy_configs().await?,
                )?))
            }

            fn validity(x: &Self::Value) -> bool {
                x.is_ok()
            }

            fn equality(x: &Self::Value, y: &Self::Value) -> bool {
                match (x, y) {
                    (Ok(x), Ok(y)) => x == y,
                    _ => false,
                }
            }
        }

        self.compute(&PackageBoundaryExceptionsKey).await?
    }

    async fn get_package_boundary_exception(
        &self,
        path: CellPathRef<'async_trait>,
    ) -> SharedResult<bool> {
        #[derive(Hash, Eq, PartialEq, Clone, Display, Debug, RefCast, Allocative)]
        #[repr(transparent)]
        struct PackageBoundaryExceptionKey(CellPath);

        #[async_trait]
        impl Key for PackageBoundaryExceptionKey {
            type Value = SharedResult<bool>;

            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellations: &CancellationContext,
            ) -> Self::Value {
                Ok(ctx
                    .get_package_boundary_exceptions()
                    .await?
                    .contains(&self.0))
            }

            fn validity(x: &Self::Value) -> bool {
                x.is_ok()
            }

            fn equality(x: &Self::Value, y: &Self::Value) -> bool {
                match (x, y) {
                    (Ok(x), Ok(y)) => x == y,
                    _ => false,
                }
            }
        }

        self.compute(&PackageBoundaryExceptionKey(path.to_owned()))
            .await?
    }
}

#[cfg(test)]
mod tests {
    use buck2_core::cells::paths::CellRelativePath;

    use super::*;

    fn get_package_boundary_exception_path(
        exceptions: &CellPackageBoundaryExceptions,
        s: &str,
    ) -> Option<CellRelativePathBuf> {
        exceptions.get_package_boundary_exception_path(CellRelativePath::unchecked_new(s))
    }

    fn package_boundary_allowlist_path(s: &str) -> Option<CellRelativePathBuf> {
        Some(CellRelativePathBuf::unchecked_new(s.to_owned()))
    }

    #[test]
    fn test_package_boundary_exceptions() {
        let exceptions =
            CellPackageBoundaryExceptions::new("foo/bar/foo,foo/bar/bar,foo/bar/baz,foo/baz,baz")
                .unwrap();

        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/foo"),
            package_boundary_allowlist_path("foo/bar/foo"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/bar"),
            package_boundary_allowlist_path("foo/bar/bar"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/baz"),
            package_boundary_allowlist_path("foo/bar/baz"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/baz/test"),
            package_boundary_allowlist_path("foo/bar/baz"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/baz/util"),
            package_boundary_allowlist_path("foo/baz"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/baz/util"),
            package_boundary_allowlist_path("foo/baz"),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "baz/foo/bar"),
            package_boundary_allowlist_path("baz"),
        );

        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/scripts"),
            None,
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/scripts"),
            None,
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "qux"),
            None,
        );
    }

    #[test]
    fn test_package_boundary_dot() {
        let exceptions = CellPackageBoundaryExceptions::new("").unwrap();
        assert!(get_package_boundary_exception_path(&exceptions, "foo/bar/foo").is_none());

        let exceptions = CellPackageBoundaryExceptions::new(",foo").unwrap();
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/foo"),
            package_boundary_allowlist_path("foo"),
        );
        assert_eq!(
            exceptions
                .get_package_boundary_exception_path(CellRelativePath::unchecked_new("bar/foo")),
            None,
        );

        let exceptions = CellPackageBoundaryExceptions::new(".").unwrap();
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "foo/bar/foo"),
            package_boundary_allowlist_path(""),
        );
        assert_eq!(
            get_package_boundary_exception_path(&exceptions, "bar/foo"),
            package_boundary_allowlist_path(""),
        );
    }
}


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

use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::sync::Arc;

use allocative::Allocative;
use dice::DiceError;
use dupe::Dupe;

/// SharedError is a simple, cloneable Error wrapper. It holds the inner error in an Arc to support Clone.
///
/// Propagation of errors via `?` convert automatically between error types, but propagating results directly
/// as different error types requires use of `shared_error()` or `unshared_error()`.
///
/// ```
/// use buck2_common::result::*;
///
/// fn foo() -> Result<(), anyhow::Error> {
///     bar()?;
///     bar().unshared_error()
/// }
///
/// fn bar() -> Result<(), SharedError> {
///     foo()?;
///     io()?;
///     foo().shared_error()
/// }
///
/// fn io() -> Result<(), std::io::Error> {
///     Ok(())
/// }
///
/// ```
///
/// Caveat: Each std::error::Error implementation requires its own explicit `impl From<E> for SharedError`, so we add those as needed.
#[derive(Clone, Dupe, Allocative)]
pub struct SharedError(Arc<anyhow::Error>);

impl SharedError {
    pub fn new(e: impl Into<anyhow::Error>) -> SharedError {
        SharedError(Arc::new(e.into()))
    }

    pub fn inner(&self) -> &anyhow::Error {
        &self.0
    }
}

impl Display for SharedError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        // Normally when displaying an anyhow, we format with {:#}.
        // However, because we also give up source(), that can lead
        // to O(n^2) output. See the test at the bottom of this file.
        Display::fmt(&self.0, f)
    }
}

/// Because `anyhow::Error` overrides `Debug` to display, we override `Debug` too.
impl Debug for SharedError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        Debug::fmt(&self.0, f)
    }
}

impl std::error::Error for SharedError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.0.source()
    }

    fn cause(&self) -> Option<&dyn std::error::Error> {
        #[allow(deprecated)]
        self.0.cause()
    }
}

impl From<&SharedError> for SharedError {
    fn from(v: &SharedError) -> Self {
        v.dupe()
    }
}

impl From<anyhow::Error> for SharedError {
    fn from(err: anyhow::Error) -> Self {
        SharedError(Arc::new(err))
    }
}

impl From<Arc<anyhow::Error>> for SharedError {
    fn from(err: Arc<anyhow::Error>) -> Self {
        SharedError(err)
    }
}

impl From<std::io::Error> for SharedError {
    fn from(err: std::io::Error) -> Self {
        SharedError(Arc::new(err.into()))
    }
}

impl From<tokio::task::JoinError> for SharedError {
    fn from(err: tokio::task::JoinError) -> Self {
        SharedError(Arc::new(err.into()))
    }
}

impl From<DiceError> for SharedError {
    fn from(err: DiceError) -> Self {
        SharedError(Arc::new(err.into()))
    }
}

pub type SharedResult<T> = Result<T, SharedError>;

/// shared_error() can convert some non-shared Result to Result.
///
/// The `?` operator will automatically convert non-shared Error to SharedError, so shared_error()
/// is usually unnecessary except when propagating a Result directly as a return value.
pub trait ToSharedResultExt<T> {
    fn shared_error(self) -> SharedResult<T>;
}

impl<T, E> ToSharedResultExt<T> for Result<T, E>
where
    SharedError: From<E>,
{
    fn shared_error(self) -> SharedResult<T> {
        Ok(self?)
    }
}

/// unshared_error() can be used to convert a SharedResult to a normal anyhow::Result. The inner error will still be shared.
///
/// The `?` operator will automatically convert SharedError to non-shared Error in the same way,
/// so unshared_error() is only necessary when propagating a SharedResult directly.
pub trait ToUnsharedResultExt<T, E> {
    fn unshared_error(self) -> Result<T, E>;
}

impl<T> ToUnsharedResultExt<T, anyhow::Error> for SharedResult<T> {
    fn unshared_error(self) -> Result<T, anyhow::Error> {
        self.map_err(|e| e.into())
    }
}

/// Like `downcast_ref()`, but if the error is a [`SharedError`], attempt to downcast
/// its `inner()` error instead.
pub fn shared_downcast_ref<'a, E: std::error::Error + Display + Debug + Send + Sync + 'static>(
    error: &'a (dyn std::error::Error + 'static),
) -> Option<&'a E> {
    match error.downcast_ref::<E>() {
        Some(e) => Some(e),
        None => match error.downcast_ref::<SharedError>() {
            Some(shared_err) => shared_err.inner().downcast_ref::<E>(),
            None => None,
        },
    }
}

/// Traverse `anyhow::Error` in `SharedError` recursively until the context `E` is found.
pub fn recursive_shared_downcast_ref<E>(error: &anyhow::Error) -> Option<&E>
where
    E: Display + Debug + Send + Sync + 'static,
{
    let mut err = error;
    loop {
        match err.downcast_ref::<E>() {
            Some(e) => return Some(e),
            None => match err.downcast_ref::<SharedError>() {
                Some(shared_err) => {
                    err = shared_err.inner();
                }
                None => return None,
            },
        }
    }
}

pub trait MayProvideAnyhowError {
    fn as_anyhow(&self) -> Option<&anyhow::Error>;
}

impl<T> MayProvideAnyhowError for Result<T, SharedError> {
    fn as_anyhow(&self) -> Option<&anyhow::Error> {
        self.as_ref().err().map(|e| e.inner())
    }
}

impl<T> MayProvideAnyhowError for anyhow::Result<T> {
    fn as_anyhow(&self) -> Option<&anyhow::Error> {
        self.as_ref().err()
    }
}

impl MayProvideAnyhowError for SharedError {
    fn as_anyhow(&self) -> Option<&anyhow::Error> {
        Some(self.inner())
    }
}

impl MayProvideAnyhowError for anyhow::Error {
    fn as_anyhow(&self) -> Option<&anyhow::Error> {
        Some(self)
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    // Which ~VALUE words are in the string
    // Use {||} to make sure we only get real signal
    fn parts(x: &str) -> Vec<&str> {
        let mut res: Vec<&str> = x
            .split_whitespace()
            .filter_map(|x| x.strip_prefix('~'))
            .collect();
        res.sort_unstable();
        res
    }

    #[test]
    fn test_shared_error_display() {
        // With anyhow/context its easy to either lose the original error (if you forget {:#}),
        // and also easy to get O(n^2) instances if you do alternative formatting, AND
        // anyhow looks at Error::source (which it does with alternative formatting).

        let x = anyhow::anyhow!(" ~A ");
        assert_eq!(parts(&format!("{}", x)), &["A"]);
        assert_eq!(parts(&format!("{:#}", x)), &["A"]);

        let x = x.context(" ~B ");
        assert_eq!(parts(&format!("{}", x)), &["B"]);
        assert_eq!(parts(&format!("{:#}", x)), &["A", "B"]);

        let x: SharedError = x.into();
        assert_eq!(parts(&format!("{}", x)), &["B"]);
        assert_eq!(parts(&format!("{:#}", x)), &["A", "B"]);

        let x: anyhow::Error = x.into();
        assert_eq!(parts(&format!("{}", x)), &["B"]);
        assert_eq!(parts(&format!("{:#}", x)), &["A", "B"]);

        let x = x.context(" ~C ");
        assert_eq!(parts(&format!("{}", x)), &["C"]);
        assert_eq!(parts(&format!("{:#}", x)), &["A", "B", "C"]);
    }

    #[derive(thiserror::Error, Debug)]
    #[error("Test error")]
    struct TestError {
        source: Option<anyhow::Error>,
    }

    #[derive(thiserror::Error, Debug)]
    #[error("Inner error")]
    struct TestInnerError;

    #[test]
    fn test_anyhow_chain_works() {
        let err: anyhow::Error = TestError { source: None }.into();
        let err_with_source: anyhow::Error = TestError {
            source: Some(TestInnerError {}.into()),
        }
        .into();

        let shared_err: anyhow::Error = SharedError::from(err).into();
        let shared_err_with_source: anyhow::Error = SharedError::from(err_with_source).into();

        let found_err = shared_err
            .chain()
            .find_map(|e| shared_downcast_ref::<TestError>(e));
        assert!(found_err.is_some());

        let found_err = shared_err_with_source
            .chain()
            .find_map(|e| shared_downcast_ref::<TestInnerError>(e));
        assert!(found_err.is_some());
    }

    #[test]
    fn test_anyhow_recursive_downcast() {
        #[derive(thiserror::Error, Debug)]
        #[error("Bottom")]
        struct ContextBottom;

        #[derive(thiserror::Error, Debug)]
        #[error("Middle")]
        struct ContextMiddle;

        #[derive(thiserror::Error, Debug)]
        #[error("Top")]
        struct ContextTop;

        #[derive(thiserror::Error, Debug)]
        #[error("None")]
        struct ContextNone;

        // Construct an error stack such that
        //
        // anyhow::Error - ContextTop
        //     v
        // SharedError
        //     v
        // anyhow::Error - ContextMiddle
        //     v
        // SharedError
        //     v
        // SharedError
        //     v
        // anyhow::Error - ContextBottom
        let error_stack = anyhow::Error::from(SharedError::new(
            anyhow::Error::from(SharedError::new(SharedError::from(
                anyhow::anyhow!("bottom").context(ContextBottom {}),
            )))
            .context(ContextMiddle {}),
        ))
        .context(ContextTop {});

        assert!(recursive_shared_downcast_ref::<ContextBottom>(&error_stack).is_some());
        assert!(recursive_shared_downcast_ref::<ContextMiddle>(&error_stack).is_some());
        assert!(recursive_shared_downcast_ref::<ContextTop>(&error_stack).is_some());
        assert!(recursive_shared_downcast_ref::<ContextNone>(&error_stack).is_none());
    }
}


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

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

use anyhow::Context;
use itertools::Itertools;
use parking_lot::Mutex;
use rusqlite::Connection;

/// A generic sqlite table for storing string key-value pairs.
pub struct KeyValueSqliteTable {
    table_name: String,
    connection: Arc<Mutex<Connection>>,
}

impl KeyValueSqliteTable {
    pub fn new(table_name: String, connection: Arc<Mutex<Connection>>) -> Self {
        Self {
            table_name,
            connection,
        }
    }

    pub fn create_table(&self) -> anyhow::Result<()> {
        let sql = format!(
            "CREATE TABLE {} (
                key     TEXT PRIMARY KEY NOT NULL,
                value   TEXT NOT NULL
            )",
            self.table_name
        );
        tracing::trace!(sql = %sql, "creating table");
        self.connection
            .lock()
            .execute(&sql, [])
            .with_context(|| format!("creating sqlite table {}", self.table_name))?;
        Ok(())
    }

    pub fn insert_all(&self, map: HashMap<String, String>) -> anyhow::Result<()> {
        let sql = format!(
            "INSERT OR REPLACE INTO {} (key, value) VALUES {}",
            self.table_name,
            // According to rusqlite docs this is the recommended way to generate the right
            // number of query placeholders for multi-row insertions.
            map.iter().map(|_| "(?, ?)").join(", ")
        );
        tracing::trace!(sql = %sql, map = ?map, "inserting into table");

        self.connection
            .lock()
            .execute(
                &sql,
                rusqlite::params_from_iter(map.into_iter().flat_map(|(k, v)| [k, v])),
            )
            .with_context(|| format!("inserting into sqlite table {}", self.table_name))?;
        Ok(())
    }

    pub fn read_all(&self) -> anyhow::Result<HashMap<String, String>> {
        let sql = format!("SELECT key, value FROM {}", self.table_name);
        tracing::trace!(sql = %sql, "read all from table");
        let connection = self.connection.lock();
        let mut stmt = connection.prepare(&sql)?;
        let map = stmt
            .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
            .collect::<Result<HashMap<String, String>, _>>()
            .with_context(|| format!("reading from sqlite table {}", self.table_name))?;
        Ok(map)
    }

    pub fn get(&self, key: &str) -> anyhow::Result<Option<String>> {
        let sql = format!("SELECT value FROM {} WHERE key = ?", self.table_name);
        tracing::trace!(sql = %sql, key = %key, "read from table");
        let connection = self.connection.lock();
        let mut stmt = connection.prepare(&sql)?;
        let row = stmt
            .query_map([key], |row| row.get(0))?
            .next()
            .transpose()
            .with_context(|| format!("reading `{}` from sqlite table {}", key, self.table_name))?;
        Ok(row)
    }
}

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

    use buck2_core::fs::project::ProjectRootTemp;
    use buck2_core::fs::project_rel_path::ProjectRelativePath;

    use super::*;

    #[test]
    fn test_key_value_sqlite_table() {
        let fs = ProjectRootTemp::new().unwrap();
        let connection = Connection::open(
            fs.path()
                .resolve(ProjectRelativePath::unchecked_new("test.db")),
        )
        .unwrap();
        let table =
            KeyValueSqliteTable::new("metadata".to_owned(), Arc::new(Mutex::new(connection)));

        table.create_table().unwrap();

        let expected = HashMap::from([
            ("foo".to_owned(), "foo".to_owned()),
            ("bar".to_owned(), "bar".to_owned()),
        ]);
        table.insert_all(expected.clone()).unwrap();

        let actual = table.read_all().unwrap();
        assert_eq!(expected, actual);

        assert_eq!(table.get("foo").unwrap().as_deref(), Some("foo"));
        assert_eq!(table.get("baz").unwrap(), None);
    }
}


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

use allocative::Allocative;
use async_trait::async_trait;
use buck2_core::cells::name::CellName;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::package::PackageLabel;
use buck2_core::target_aliases::TargetAliasResolver;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use indexmap::IndexSet;
use itertools::Itertools;
use more_futures::cancellation::CancellationContext;

use crate::dice::cells::HasCellResolver;
use crate::legacy_configs::dice::HasLegacyConfigs;
use crate::legacy_configs::LegacyBuckConfig;
use crate::result::SharedResult;

#[derive(thiserror::Error, Debug)]
enum AliasResolutionError {
    #[error("No [alias] section in buckconfig")]
    MissingAliasSection,
    #[error("[alias] section does not contain the requested alias")]
    NotAnAlias,
    #[error("[alias] section produced a dangling chain: `{}`", .0.iter().join(" -> "))]
    AliasChainBroken(Vec<String>),
    #[error("cycle detected in alias resolution [{} -> {1}]", .0.iter().join(" -> "))]
    AliasCycle(Vec<String>, String),
}

#[derive(Dupe, Clone, Allocative)]
pub struct BuckConfigTargetAliasResolver {
    config: LegacyBuckConfig,
}

impl PartialEq for BuckConfigTargetAliasResolver {
    fn eq(&self, other: &BuckConfigTargetAliasResolver) -> bool {
        // `TargetAliasResolver` only uses `alias` section of buckconfig,
        // comparing only this section is enough.
        // Please update this code if `TargetAliasResolver` uses other buckconfigs.
        let self_aliases = self.config.get_section("alias");
        let other_aliases = other.config.get_section("alias");
        match (self_aliases, other_aliases) {
            (Some(self_aliases), Some(other_aliases)) => self_aliases.compare(other_aliases),
            (None, None) => true,
            (None, Some(_)) | (Some(_), None) => false,
        }
    }
}

impl TargetAliasResolver for BuckConfigTargetAliasResolver {
    fn get<'a>(&'a self, name: &str) -> anyhow::Result<Option<&'a str>> {
        match self.resolve_alias(name) {
            Ok(a) => Ok(Some(a)),
            Err(AliasResolutionError::MissingAliasSection | AliasResolutionError::NotAnAlias) => {
                Ok(None)
            }
            Err(
                e @ AliasResolutionError::AliasChainBroken(..)
                | e @ AliasResolutionError::AliasCycle(..),
            ) => Err(anyhow::Error::from(e).context(format!("Error resolving alias `{}`", name))),
        }
    }
}

impl BuckConfigTargetAliasResolver {
    pub fn new(config: LegacyBuckConfig) -> Self {
        Self { config }
    }

    /// Resolves an alias in the `[alias]` section. Aliases can refer to other aliases. Any
    /// string containing ":" is considered to be the end of the alias resolution.
    fn resolve_alias<'a>(&'a self, alias: &str) -> Result<&'a str, AliasResolutionError> {
        if alias.contains(':') {
            return Err(AliasResolutionError::NotAnAlias);
        }

        let mut alias = alias;

        let section = self.config.get_section("alias");
        let mut stack = IndexSet::<&str>::new();
        loop {
            if stack.contains(alias) {
                return Err(AliasResolutionError::AliasCycle(
                    stack.into_iter().map(|s| s.to_owned()).collect(),
                    alias.to_owned(),
                ));
            }

            let new_alias = match &section {
                Some(section) => match section.get(alias) {
                    Some(v) => {
                        stack.insert(alias);
                        v.as_str()
                    }
                    None => {
                        if stack.is_empty() {
                            return Err(AliasResolutionError::NotAnAlias);
                        } else {
                            let chain = stack
                                .into_iter()
                                .chain(std::iter::once(alias))
                                .map(|e| e.to_owned())
                                .collect();
                            return Err(AliasResolutionError::AliasChainBroken(chain));
                        }
                    }
                },
                None => return Err(AliasResolutionError::MissingAliasSection),
            };

            if new_alias.contains(':') {
                return Ok(new_alias);
            }
            alias = new_alias;
        }
    }
}

#[async_trait]
pub trait HasTargetAliasResolver {
    async fn target_alias_resolver_for_cell(
        &self,
        cell_name: CellName,
    ) -> anyhow::Result<BuckConfigTargetAliasResolver>;

    async fn target_alias_resolver_for_working_dir(
        &self,
        working_dir: &ProjectRelativePath,
    ) -> anyhow::Result<BuckConfigTargetAliasResolver>;
}

#[derive(Debug, Display, Hash, PartialEq, Eq, Clone, Allocative)]
struct TargetAliasResolverKey {
    cell_name: CellName,
}

#[async_trait]
impl Key for TargetAliasResolverKey {
    type Value = SharedResult<BuckConfigTargetAliasResolver>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> SharedResult<BuckConfigTargetAliasResolver> {
        let legacy_configs = ctx.get_legacy_config_for_cell(self.cell_name).await?;
        Ok(legacy_configs.target_alias_resolver())
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Ok(x), Ok(y)) => x == y,
            _ => false,
        }
    }
}

#[async_trait]
impl HasTargetAliasResolver for DiceComputations {
    async fn target_alias_resolver_for_cell(
        &self,
        cell_name: CellName,
    ) -> anyhow::Result<BuckConfigTargetAliasResolver> {
        Ok(self
            .compute(&TargetAliasResolverKey { cell_name })
            .await??)
    }

    async fn target_alias_resolver_for_working_dir(
        &self,
        working_dir: &ProjectRelativePath,
    ) -> anyhow::Result<BuckConfigTargetAliasResolver> {
        let cell_resolver = self.get_cell_resolver().await?;
        let working_dir =
            PackageLabel::from_cell_path(cell_resolver.get_cell_path(&working_dir)?.as_ref());
        let cell_name = working_dir.as_cell_path().cell();
        self.target_alias_resolver_for_cell(cell_name).await
    }
}

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

    use crate::legacy_configs;
    use crate::target_aliases::AliasResolutionError;
    use crate::target_aliases::BuckConfigTargetAliasResolver;

    #[test]
    fn test_aliases() -> anyhow::Result<()> {
        let config = legacy_configs::testing::parse(
            &[(
                "/config",
                indoc!(
                    r#"
            [alias]
              baz = foo
              foo = //:foo
              bar = foo
              bar2 = bar
              cycle1 = cycle2
              cycle2 = cycle3
              cycle3 = cycle1
              chain1 = chain2
              chain2 = chain3

        "#
                ),
            )],
            "/config",
        )?;

        let target_alias_resolver = BuckConfigTargetAliasResolver::new(config);

        assert_eq!("//:foo", target_alias_resolver.resolve_alias("foo")?);
        assert_eq!("//:foo", target_alias_resolver.resolve_alias("bar")?);
        assert_eq!("//:foo", target_alias_resolver.resolve_alias("bar2")?);
        assert_eq!("//:foo", target_alias_resolver.resolve_alias("baz")?);

        assert_matches!(
            target_alias_resolver.resolve_alias("missing"),
            Err(AliasResolutionError::NotAnAlias)
        );

        assert_matches!(
            target_alias_resolver.resolve_alias("chain1"),
            Err(e) => {
                let err = format!("{:#}", e);
                let expected = "chain1 -> chain2 -> chain3";
                assert!(
                    err.contains(expected),
                    "expected error to contain `{}`, got `{}`",
                    expected,
                    err
                );
            }
        );

        assert_matches!(
            target_alias_resolver.resolve_alias("cycle1"),
            Err(e) => {
                let err = format!("{:#}", e);
                let expected = "cycle1 -> cycle2 -> cycle3 -> cycle1";
                assert!(
                    err.contains(expected),
                    "expected error to contain `{}`, got `{}`",
                    expected,
                    err
                );
            }
        );

        Ok(())
    }
}


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

use std::env;
use std::mem;

use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::file_name::FileNameBuf;
use rand::Rng;

/// Temporary path.
///
/// Just a path, but whatever is written to that path, is removed on drop.
pub struct TempPath {
    /// `None` when explicitly "closed".
    path: Option<AbsNormPathBuf>,
}

impl TempPath {
    pub fn new() -> anyhow::Result<TempPath> {
        TempPath::new_in(AbsNormPath::new(&env::temp_dir())?)
    }

    pub fn new_path(path: AbsNormPathBuf) -> TempPath {
        TempPath { path: Some(path) }
    }

    pub fn new_in(temp_dir: &AbsNormPath) -> anyhow::Result<TempPath> {
        let mut name = String::with_capacity(10);
        for _ in 0..10 {
            name.push(rand::thread_rng().gen_range('a'..='z'));
        }
        let path = temp_dir.join(FileNameBuf::try_from(name)?);
        Ok(TempPath { path: Some(path) })
    }

    pub fn path(&self) -> &AbsNormPath {
        self.path.as_deref().unwrap()
    }

    /// Delete temp file explicitly, because this returns error
    /// and `drop` can only panic or ignore error.
    pub fn close(mut self) -> anyhow::Result<()> {
        let path = mem::take(&mut self.path).unwrap();
        fs_util::remove_all(path)?;
        Ok(())
    }
}

impl Drop for TempPath {
    fn drop(&mut self) {
        if let Some(path) = mem::take(&mut self.path) {
            // Ignore error.
            drop(fs_util::remove_all(path));
        }
    }
}

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

    use crate::temp_path::TempPath;

    #[test]
    fn test_temp_path() {
        let temp_path = TempPath::new().unwrap();
        let path = temp_path.path().to_path_buf();

        assert!(!fs::try_exists(&path).unwrap());

        fs::write(&path, "hello").unwrap();

        assert!(fs::try_exists(&path).unwrap(), "Sanity check");

        temp_path.close().unwrap();

        assert!(!fs::try_exists(&path).unwrap());
    }
}


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

//! Core dice computations relating to cells

use allocative::Allocative;
use async_trait::async_trait;
use buck2_core::cells::CellResolver;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dupe::Dupe;

#[async_trait]
pub trait HasCellResolver {
    async fn get_cell_resolver(&self) -> anyhow::Result<CellResolver>;

    async fn is_cell_resolver_key_set(&self) -> anyhow::Result<bool>;
}

pub trait SetCellResolver {
    fn set_cell_resolver(&mut self, cell_resolver: CellResolver) -> anyhow::Result<()>;

    fn set_none_cell_resolver(&mut self) -> anyhow::Result<()>;
}

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "{:?}", self)]
struct CellResolverKey;

impl InjectedKey for CellResolverKey {
    type Value = Option<CellResolver>;

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Some(x), Some(y)) => x == y,
            (None, None) => true,
            (_, _) => false,
        }
    }
}

#[async_trait]
impl HasCellResolver for DiceComputations {
    async fn get_cell_resolver(&self) -> anyhow::Result<CellResolver> {
        self.compute(&CellResolverKey).await?.ok_or_else(|| {
            panic!("Tried to retrieve CellResolverKey from the graph, but key has None value")
        })
    }

    async fn is_cell_resolver_key_set(&self) -> anyhow::Result<bool> {
        Ok(self.compute(&CellResolverKey).await?.is_some())
    }
}

impl SetCellResolver for DiceTransactionUpdater {
    fn set_cell_resolver(&mut self, cell_resolver: CellResolver) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(CellResolverKey, Some(cell_resolver))])?)
    }

    fn set_none_cell_resolver(&mut self) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(CellResolverKey, None)])?)
    }
}


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

use std::any::Any;
use std::fmt::Debug;

use allocative::Allocative;
use async_trait::async_trait;
use buck2_util::cycle_detector::CycleDescriptor;
use buck2_util::cycle_detector::LazyCycleDetector;
use buck2_util::cycle_detector::LazyCycleDetectorGuard;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dice::UserCycleDetector;
use dice::UserCycleDetectorGuard;
use futures::Future;
use more_futures::cancellation::CancellationContext;
use tracing::debug;

/// Additional requirement for a CycleDescriptor to be used for defining a Dice UserCycleDetector through
/// the CycleDetectorAdapter. Simply requires converting the Dice Key to the CycleDescriptor::Key type.
pub trait CycleAdapterDescriptor: CycleDescriptor {
    /// Will be provided a &dyn Any for a Dice Key implementation.
    fn to_key(key: &dyn Any) -> Option<Self::Key>;
}

/// This allows using the LazyCycleDetector as a Dice UserCycleDetector. All it needs is an implementation of the normal
/// CycleDescriptor for the LazyCycleDetector and a CycleAdapterDescriptor to convert dice keys to the CycleDescriptor::Key.
#[derive(Debug)]
pub struct CycleDetectorAdapter<D: CycleAdapterDescriptor> {
    inner: LazyCycleDetector<D>,
}

#[async_trait]
pub trait CycleGuard<E> {
    async fn guard_this<R: Send, Fut: Future<Output = R> + Send>(
        ctx: &DiceComputations,
        fut: Fut,
    ) -> anyhow::Result<Result<R, E>>;
}

#[async_trait]
impl<D: CycleAdapterDescriptor> CycleGuard<D::Error> for D {
    /// Use this to wrap futures waiting on dependencies where you want to detect cycles. All dice keys involved
    /// in the cycle must be supported by the CycleAdapterDescriptor implementation. If the cycle is FooKey(1) ->
    /// BarKey(a) -> FooKey(1) and the descriptor only support FooKey, the cycle won't be detected.
    ///
    /// The guard_this generally just needs to appear on at least one of the edges in the cycle. For example,
    /// FooKey/BarKey example above, a descriptor that supports both keys but only guards on the BarKey->FooKey
    /// edge would probably be okay (that edge would detect the cycle and get an error, which would
    /// then propagate up to the waiting fookey).
    ///
    /// It's probably the case that the keys in the cycle should treat the cycle error case as invalid (in the sense
    /// of Dice Key::validity()).
    async fn guard_this<R: Send, Fut: Future<Output = R> + Send>(
        ctx: &DiceComputations,
        fut: Fut,
    ) -> anyhow::Result<Result<R, D::Error>> {
        match ctx.cycle_guard::<CycleAdapterGuard<D>>()? {
            Some(v) => match v.guard.guard_this(fut).await {
                Ok(Ok(v)) => Ok(Ok(v)),
                v => {
                    // The cycle detector either hit an error or it detected a cycle. In either case, we
                    // want to make sure that dice doesn't cache this node. To do that, we add a dep on our
                    // PoisonedDueToDetectedCycle key. We shouldn't hit a dice error, but we know we're already
                    // returning an error so just ignore it.
                    let _unused = ctx.compute(&PoisonedDueToDetectedCycleKey).await;
                    v
                }
            },
            None => Ok(Ok(fut.await)),
        }
    }
}

/// This is a simple type we can use that will mark any dice node that depends on it as invalid.
/// This is used to ensure we don't cache an error from the cycle detector (the cycle detector allows
/// flow of data that is potentially not tracked by dice, and while we may be able to identify those
/// and fix them it'll still be fragile and its best to just make sure they aren't cached).
#[derive(Allocative, Debug, Display, Clone, PartialEq, Eq, Hash)]
#[display(fmt = "poisoned_due_to_detected_cycle")]
struct PoisonedDueToDetectedCycleKey;

#[async_trait]
impl Key for PoisonedDueToDetectedCycleKey {
    type Value = ();

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        true
    }

    fn validity(_x: &Self::Value) -> bool {
        false
    }
}

impl<D: CycleAdapterDescriptor> CycleDetectorAdapter<D> {
    pub fn new() -> Self {
        Self {
            inner: LazyCycleDetector::new(),
        }
    }
}

impl<D: CycleAdapterDescriptor> CycleDescriptor for CycleDetectorAdapter<D> {
    type Key = D::Key;
    type Error = D::Error;

    fn cycle_error(cycle: Vec<&Self::Key>) -> Self::Error {
        D::cycle_error(cycle)
    }
}

impl<D: CycleAdapterDescriptor> UserCycleDetector for CycleDetectorAdapter<D> {
    fn start_computing_key(&self, key: &dyn Any) -> Option<Box<dyn UserCycleDetectorGuard>> {
        match D::to_key(key) {
            None => None,
            Some(v) => Some(Box::new(CycleAdapterGuard {
                guard: self.inner.start(v),
            })),
        }
    }

    fn finished_computing_key(&self, key: &dyn Any) {
        if let Some(v) = D::to_key(key) {
            debug!("finish computing key {}", v);
            self.inner.finish(v);
        }
    }
}

pub struct CycleAdapterGuard<D: CycleAdapterDescriptor> {
    guard: LazyCycleDetectorGuard<D>,
}

impl<D: CycleAdapterDescriptor> UserCycleDetectorGuard for CycleAdapterGuard<D> {
    fn add_edge(&self, key: &dyn Any) {
        if let Some(k) = D::to_key(key) {
            self.guard.add_edge(k);
        }
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

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

/// Allows having multiple cycle detectors (e.g. for different keys/phases) as the one Dice UserCycleDetector.
pub struct PairDiceCycleDetector<A: UserCycleDetector, B: UserCycleDetector>(pub A, pub B);

impl<A: UserCycleDetector, B: UserCycleDetector> UserCycleDetector for PairDiceCycleDetector<A, B> {
    fn start_computing_key(&self, key: &dyn Any) -> Option<Box<dyn UserCycleDetectorGuard>> {
        // Right now, only one of the inner detectors is allowed to claim a key. We could feasibly change that, but it's a bit trickier.
        if let Some(v) = self.0.start_computing_key(key) {
            return Some(v);
        }
        if let Some(v) = self.1.start_computing_key(key) {
            return Some(v);
        }
        None
    }

    fn finished_computing_key(&self, key: &dyn Any) {
        self.0.finished_computing_key(key);
        self.1.finished_computing_key(key);
    }
}

#[cfg(test)]
mod tests {
    use std::any::Any;
    use std::fmt::Debug;
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering;

    use dice::UserCycleDetector;
    use dice::UserCycleDetectorGuard;

    use crate::dice::cycles::PairDiceCycleDetector;

    #[test]
    fn pair_cycle_detector() {
        struct TestingGuard;

        impl UserCycleDetectorGuard for TestingGuard {
            fn add_edge(&self, _key: &dyn Any) {
                unreachable!("testing")
            }

            fn as_any(&self) -> &dyn Any {
                unreachable!("testing")
            }

            fn type_name(&self) -> &'static str {
                unreachable!("testing")
            }
        }

        #[derive(Debug, Default)]
        struct ReceivesStartAndFinish {
            got_start: AtomicBool,
            got_finish: AtomicBool,
        }

        impl UserCycleDetector for ReceivesStartAndFinish {
            fn start_computing_key(
                &self,
                _key: &dyn Any,
            ) -> Option<Box<dyn UserCycleDetectorGuard>> {
                self.got_start.store(true, Ordering::SeqCst);
                Some(Box::new(TestingGuard))
            }

            fn finished_computing_key(&self, _key: &dyn Any) {
                assert!(self.got_start.load(Ordering::SeqCst));
                self.got_finish.store(true, Ordering::SeqCst);
            }
        }

        #[derive(Debug, Default)]
        struct ReceivesOnlyFinish {
            got_finish: AtomicBool,
        }

        impl UserCycleDetector for ReceivesOnlyFinish {
            fn start_computing_key(
                &self,
                _key: &dyn Any,
            ) -> Option<Box<dyn UserCycleDetectorGuard>> {
                panic!("shouldn't be called")
            }

            fn finished_computing_key(&self, _key: &dyn Any) {
                self.got_finish.store(true, Ordering::SeqCst);
            }
        }

        let detector = PairDiceCycleDetector(
            ReceivesStartAndFinish::default(),
            ReceivesOnlyFinish::default(),
        );

        assert!(detector.start_computing_key(&()).is_some());

        detector.finished_computing_key(&());

        assert!(detector.0.got_start.load(Ordering::SeqCst));
        assert!(detector.0.got_finish.load(Ordering::SeqCst));
        assert!(detector.1.got_finish.load(Ordering::SeqCst));
    }
}


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

//! Core data on dice

use std::sync::Arc;

use dice::DiceData;
use dice::DiceDataBuilder;
use dupe::Dupe;

use crate::cas_digest::CasDigestConfig;
use crate::io::IoProvider;

pub trait HasIoProvider {
    // TODO(bobyf) can we make this not an arc
    fn get_io_provider(&self) -> Arc<dyn IoProvider>;
}

pub trait SetIoProvider {
    fn set_io_provider(&mut self, fs: Arc<dyn IoProvider>);
}

impl HasIoProvider for DiceData {
    fn get_io_provider(&self) -> Arc<dyn IoProvider> {
        self.get::<Arc<dyn IoProvider>>()
            .expect("project filesystem should be set")
            .dupe()
    }
}

impl SetIoProvider for DiceDataBuilder {
    fn set_io_provider(&mut self, fs: Arc<dyn IoProvider>) {
        self.set(fs)
    }
}

pub mod testing {
    use buck2_core::fs::project::ProjectRootTemp;

    use super::*;
    use crate::io::fs::FsIoProvider;

    pub trait SetTestingIoProvider {
        fn set_testing_io_provider(&mut self, fs: &ProjectRootTemp);
    }

    impl SetTestingIoProvider for DiceDataBuilder {
        fn set_testing_io_provider(&mut self, fs: &ProjectRootTemp) {
            self.set_io_provider(Arc::new(FsIoProvider::new(
                fs.path().dupe(),
                CasDigestConfig::testing_default(),
            )))
        }
    }
}


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

use std::collections::HashSet;
use std::fmt::Debug;
use std::hash::Hash;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::cell_root_path::CellRootPathBuf;
use buck2_core::cells::name::CellName;
use buck2_core::cells::unchecked_cell_rel_path::UncheckedCellRelativePath;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_events::dispatch::console_message;
use cmp_any::PartialEqAny;
use derivative::Derivative;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::Key;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;

use crate::dice::cells::HasCellResolver;
use crate::dice::data::HasIoProvider;
use crate::dice::file_ops::keys::FileOpsKey;
use crate::dice::file_ops::keys::FileOpsValue;
use crate::file_ops::FileOps;
use crate::file_ops::RawDirEntry;
use crate::file_ops::RawPathMetadata;
use crate::file_ops::ReadDirOutput;
use crate::file_ops::SimpleDirEntry;
use crate::ignores::all_cells::AllCellIgnores;
use crate::ignores::all_cells::HasAllCellIgnores;
use crate::io::IoProvider;
use crate::result::SharedResult;
use crate::result::ToSharedResultExt;
use crate::result::ToUnsharedResultExt;

pub trait HasFileOps<'c> {
    type T: FileOps;
    fn file_ops(&'c self) -> Self::T;
}

impl<'c> HasFileOps<'c> for DiceComputations {
    type T = DiceFileOps<'c>;
    fn file_ops(&'c self) -> DiceFileOps<'c> {
        DiceFileOps(self)
    }
}

// TODO(cjhopman, bobyf): This FileToken can go away once Dice has support for
// transient values.
/// This is used as the "result" of a read_file computation so that we don't
/// need to store the file content's in dice's cache.
#[derive(Clone, Dupe, Allocative)]
struct FileToken(Arc<CellPath>);

impl FileToken {
    async fn read_if_exists(&self, fs: &dyn FileOps) -> anyhow::Result<Option<String>> {
        fs.read_file_if_exists((*self.0).as_ref()).await
    }
}

#[derive(Clone, Dupe, Allocative)]
pub struct DiceFileOps<'c>(#[allocative(skip)] pub &'c DiceComputations);

pub mod keys {
    use std::sync::Arc;

    use allocative::Allocative;
    use derive_more::Display;
    use dupe::Dupe;

    use crate::file_ops::FileOps;

    #[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
    #[display(fmt = "{:?}", self)]
    pub struct FileOpsKey();

    #[derive(Dupe, Clone, Allocative)]
    pub struct FileOpsValue(#[allocative(skip)] pub Arc<dyn FileOps>);
}

async fn get_default_file_ops(dice: &DiceComputations) -> SharedResult<Arc<dyn FileOps>> {
    #[derive(Clone, Dupe, Derivative, Allocative)]
    #[derivative(PartialEq)]
    struct DiceFileOpsDelegate {
        // Safe to ignore because `io` does not change during the lifetime of the daemon.
        #[derivative(PartialEq = "ignore")]
        io: Arc<dyn IoProvider>,
        cells: CellResolver,
        ignores: Arc<AllCellIgnores>,
    }

    impl DiceFileOpsDelegate {
        fn resolve(&self, path: CellPathRef) -> anyhow::Result<ProjectRelativePathBuf> {
            let cell_root = self.resolve_cell_root(path.cell())?;
            Ok(cell_root.project_relative_path().join(path.path()))
        }

        fn resolve_cell_root(&self, cell: CellName) -> anyhow::Result<CellRootPathBuf> {
            Ok(self.cells.get(cell).unwrap().path().to_buf())
        }

        fn get_cell_path(&self, path: &ProjectRelativePath) -> anyhow::Result<CellPath> {
            self.cells.get_cell_path(path)
        }

        fn io_provider(&self) -> &dyn IoProvider {
            self.io.as_ref()
        }
    }

    #[async_trait]
    impl FileOps for DiceFileOpsDelegate {
        async fn read_file_if_exists(
            &self,
            path: CellPathRef<'async_trait>,
        ) -> anyhow::Result<Option<String>> {
            // TODO(cjhopman): error on ignored paths, maybe.
            let project_path = self.resolve(path)?;
            self.io_provider().read_file_if_exists(project_path).await
        }

        async fn read_dir(&self, path: CellPathRef<'async_trait>) -> anyhow::Result<ReadDirOutput> {
            // TODO(cjhopman): This should also probably verify that the parent chain is not ignored.
            self.ignores
                .check_ignored(path.cell(), UncheckedCellRelativePath::new(path.path()))?
                .into_result()
                .with_context(|| format!("Error checking whether dir `{}` is ignored", path))?;

            let project_path = self.resolve(path)?;
            let mut entries = self
                .io_provider()
                .read_dir(project_path)
                .await
                .with_context(|| format!("Error listing dir `{}`", path))?;

            // Make sure entries are deterministic, since read_dir isn't.
            entries.sort_by(|a, b| a.file_name.cmp(&b.file_name));

            let is_ignored = |file_name: &str| {
                let mut cell_relative_path_buf;
                let cell_relative_path: &str = if path.path().is_empty() {
                    file_name
                } else {
                    cell_relative_path_buf =
                        String::with_capacity(path.path().as_str().len() + 1 + file_name.len());
                    cell_relative_path_buf.push_str(path.path().as_str());
                    cell_relative_path_buf.push('/');
                    cell_relative_path_buf.push_str(file_name);
                    &cell_relative_path_buf
                };

                let cell_relative_path =
                    UncheckedCellRelativePath::unchecked_new(cell_relative_path);
                let is_ignored = self
                    .ignores
                    .check_ignored(path.cell(), cell_relative_path)?
                    .is_ignored();
                anyhow::Ok(is_ignored)
            };

            // Filter out any entries that are ignored.
            let mut included_entries = Vec::new();
            for e in entries {
                let RawDirEntry {
                    file_type,
                    file_name,
                } = e;

                if !is_ignored(&file_name)? {
                    let file_name = match FileNameBuf::try_from_or_get_back(file_name) {
                        Ok(file_name) => file_name,
                        Err(file_name) => {
                            console_message(format!(
                                "File name `{file_name}` is not valid. \
                                    Add the path to `project.ignore` to mute this message",
                            ));
                            continue;
                        }
                    };
                    included_entries.push(SimpleDirEntry {
                        file_name,
                        file_type,
                    });
                }
            }

            Ok(ReadDirOutput {
                included: included_entries.into(),
            })
        }

        async fn read_path_metadata_if_exists(
            &self,
            path: CellPathRef<'async_trait>,
        ) -> anyhow::Result<Option<RawPathMetadata>> {
            let project_path = self.resolve(path)?;

            let res = self
                .io_provider()
                .read_path_metadata_if_exists(project_path)
                .await
                .with_context(|| format!("Error accessing metadata for path `{}`", path))?;
            res.map(|meta| meta.try_map(|path| Ok(Arc::new(self.get_cell_path(&path)?))))
                .transpose()
        }

        async fn is_ignored(&self, path: CellPathRef<'async_trait>) -> anyhow::Result<bool> {
            Ok(self
                .ignores
                .check_ignored(path.cell(), UncheckedCellRelativePath::new(path.path()))?
                .is_ignored())
        }

        fn eq_token(&self) -> PartialEqAny {
            PartialEqAny::new(self)
        }
    }

    #[async_trait]
    impl Key for FileOpsKey {
        type Value = SharedResult<FileOpsValue>;
        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            let cells = ctx.get_cell_resolver().await?;
            let io = ctx.global_data().get_io_provider();

            let ignores = ctx.new_all_cell_ignores().await?;

            Ok(FileOpsValue(Arc::new(DiceFileOpsDelegate {
                io,
                cells,
                ignores,
            })))
        }

        fn equality(x: &Self::Value, y: &Self::Value) -> bool {
            match (x, y) {
                (Ok(x), Ok(y)) => *x.0 == *y.0,
                _ => false,
            }
        }

        fn validity(x: &Self::Value) -> bool {
            x.is_ok()
        }
    }

    Ok(dice.compute(&FileOpsKey()).await??.0)
}

#[derive(Allocative)]
pub struct FileChangeTracker {
    files_to_dirty: HashSet<ReadFileKey>,
    dirs_to_dirty: HashSet<ReadDirKey>,
    paths_to_dirty: HashSet<PathMetadataKey>,
}

impl FileChangeTracker {
    pub fn new() -> Self {
        Self {
            files_to_dirty: Default::default(),
            dirs_to_dirty: Default::default(),
            paths_to_dirty: Default::default(),
        }
    }

    pub fn write_to_dice(self, ctx: &mut DiceTransactionUpdater) -> anyhow::Result<()> {
        ctx.changed(self.files_to_dirty)?;
        ctx.changed(self.dirs_to_dirty)?;
        ctx.changed(self.paths_to_dirty)?;

        Ok(())
    }

    fn file_contents_modify(&mut self, path: CellPath) {
        self.files_to_dirty
            .insert(ReadFileKey(Arc::new(path.clone())));
        self.paths_to_dirty.insert(PathMetadataKey(path));
    }

    pub fn file_added_or_removed(&mut self, path: CellPath) {
        let parent = path.parent();

        self.file_contents_modify(path.clone());
        if let Some(parent) = parent {
            // The above can be None (validly!) if we have a cell we either create or delete.
            // That never happens in established repos, but if you are setting one up, it's not uncommon.
            // Since we don't include paths in different cells, the fact we don't dirty the parent
            // (which is in an enclosing cell) doesn't matter.
            self.dirs_to_dirty.insert(ReadDirKey(parent.to_owned()));
        }
    }

    pub fn dir_added_or_removed(&mut self, path: CellPath) {
        self.paths_to_dirty.insert(PathMetadataKey(path.clone()));
        if let Some(parent) = path.parent() {
            let parent = parent.to_owned();
            // The above can be None (validly!) if we have a cell we either create or delete.
            // That never happens in established repos, but if you are setting one up, it's not uncommon.
            // Since we don't include paths in different cells, the fact we don't dirty the parent
            // (which is in an enclosing cell) doesn't matter.
            self.dirs_to_dirty
                .extend([ReadDirKey(path), ReadDirKey(parent)]);
        }
    }

    pub fn file_changed(&mut self, path: CellPath) {
        self.file_contents_modify(path)
    }

    pub fn file_removed(&mut self, path: CellPath) {
        self.file_added_or_removed(path)
    }

    pub fn file_added(&mut self, path: CellPath) {
        self.file_added_or_removed(path)
    }

    pub fn dir_changed(&mut self, path: CellPath) {
        self.paths_to_dirty.insert(PathMetadataKey(path.clone()));
        self.dirs_to_dirty.insert(ReadDirKey(path));
    }

    pub fn dir_added(&mut self, path: CellPath) {
        self.dir_added_or_removed(path)
    }

    pub fn dir_removed(&mut self, path: CellPath) {
        self.dir_added_or_removed(path)
    }
}

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
struct ReadFileKey(Arc<CellPath>);

#[async_trait]
impl Key for ReadFileKey {
    type Value = FileToken;
    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        FileToken(self.0.dupe())
    }

    fn equality(_: &Self::Value, _: &Self::Value) -> bool {
        false
    }
}

#[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
struct ReadDirKey(CellPath);

#[async_trait]
impl Key for ReadDirKey {
    type Value = SharedResult<ReadDirOutput>;
    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        get_default_file_ops(ctx)
            .await?
            .read_dir(self.0.as_ref())
            .await
            .shared_error()
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Ok(x), Ok(y)) => x == y,
            _ => false,
        }
    }

    fn validity(x: &Self::Value) -> bool {
        x.is_ok()
    }
}

#[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
struct PathMetadataKey(CellPath);

#[async_trait]
impl Key for PathMetadataKey {
    type Value = SharedResult<Option<RawPathMetadata>>;
    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        let res = get_default_file_ops(ctx)
            .await?
            .read_path_metadata_if_exists(self.0.as_ref())
            .await?;

        match res {
            Some(RawPathMetadata::Symlink {
                at: ref path,
                to: _,
            }) => {
                ctx.compute(&ReadFileKey(path.dupe())).await?;
            }
            _ => (),
        };

        Ok(res)
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Ok(x), Ok(y)) => x == y,
            _ => false,
        }
    }

    fn validity(x: &Self::Value) -> bool {
        x.is_ok()
    }
}

#[async_trait]
impl<'c> FileOps for DiceFileOps<'c> {
    async fn read_file_if_exists(
        &self,
        path: CellPathRef<'async_trait>,
    ) -> anyhow::Result<Option<String>> {
        let path = path.to_owned();
        let file_ops = get_default_file_ops(self.0).await?;

        self.0
            .compute(&ReadFileKey(Arc::new(path)))
            .await?
            .read_if_exists(&*file_ops)
            .await
    }

    async fn read_dir(&self, path: CellPathRef<'async_trait>) -> anyhow::Result<ReadDirOutput> {
        self.0
            .compute(&ReadDirKey(path.to_owned()))
            .await?
            .unshared_error()
    }

    async fn read_path_metadata_if_exists(
        &self,
        path: CellPathRef<'async_trait>,
    ) -> anyhow::Result<Option<RawPathMetadata>> {
        self.0
            .compute(&PathMetadataKey(path.to_owned()))
            .await?
            .unshared_error()
    }

    async fn is_ignored(&self, path: CellPathRef<'async_trait>) -> anyhow::Result<bool> {
        let file_ops = get_default_file_ops(self.0).await?;
        file_ops.is_ignored(path).await
    }

    fn eq_token(&self) -> PartialEqAny {
        // We do not store this on DICE, so we don't care about equality.
        // Also we cannot do `PartialEqAny` here because `Self` is not `'static`.
        PartialEqAny::always_false()
    }
}

pub mod testing {
    pub use super::keys::FileOpsKey;
}


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

//! Common dice operations

pub mod cells;
pub mod cycles;
pub mod data;
pub mod file_ops;


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

use std::pin::Pin;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use allocative::Allocative;
use bytes::Bytes;
use dupe::Dupe;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::task::Poll;
use futures::Stream;
use hyper::Request;
use hyper::Response;
use pin_project::pin_project;

use crate::http::HttpClient;
use crate::http::HttpError;

#[derive(Allocative)]
pub struct HttpNetworkStats {
    pub downloaded_bytes: u64,
}

/// Wrapper for the HTTP client, storing collected metrics such as the number
/// of bytes downloaded by the client's GET requests during its entire lifetime.
#[derive(Allocative, Clone, Dupe)]
pub struct CountingHttpClient {
    inner: Arc<dyn HttpClient>,
    bytes_downloaded: Arc<AtomicU64>,
}

impl CountingHttpClient {
    pub fn new(client: Arc<dyn HttpClient>) -> Self {
        Self {
            inner: client,
            bytes_downloaded: Arc::new(AtomicU64::new(0)),
        }
    }

    pub fn get_updated_http_network_stats(&self) -> HttpNetworkStats {
        HttpNetworkStats {
            downloaded_bytes: self.bytes_downloaded.load(Ordering::Relaxed),
        }
    }

    pub fn supports_vpnless(&self) -> bool {
        self.inner.supports_vpnless()
    }
}

#[pin_project]
struct CountingStream<S> {
    #[pin]
    inner: S,
    bytes_read: Arc<AtomicU64>,
}

impl<S> CountingStream<S> {
    fn new(stream: S, bytes_read: Arc<AtomicU64>) -> Self {
        Self {
            inner: stream,
            bytes_read,
        }
    }
}

impl<S, E> Stream for CountingStream<S>
where
    S: Stream<Item = Result<Bytes, E>>,
{
    type Item = Result<Bytes, E>;

    fn poll_next(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        match this.inner.as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(bytes))) => {
                this.bytes_read
                    .fetch_add(bytes.len() as u64, Ordering::Relaxed);
                Poll::Ready(Some(Ok(bytes)))
            }
            x => x,
        }
    }
}

#[async_trait::async_trait]
impl HttpClient for CountingHttpClient {
    async fn request(
        &self,
        request: Request<Bytes>,
    ) -> Result<Response<BoxStream<hyper::Result<Bytes>>>, HttpError> {
        let resp = self.inner.request(request).await?;
        let resp = resp.map(|body| CountingStream::new(body, self.bytes_downloaded.dupe()).boxed());
        Ok(resp)
    }

    fn supports_vpnless(&self) -> bool {
        CountingHttpClient::supports_vpnless(self)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::Ordering;

    use futures::StreamExt;
    use httptest::matchers::*;
    use httptest::responders;
    use httptest::Expectation;

    use super::*;
    use crate::http::secure_client::SecureHttpClient;
    use crate::http::tls_config_with_system_roots;

    #[tokio::test]
    async fn test_count_response_size() -> anyhow::Result<()> {
        let test_server = httptest::Server::run();
        test_server.expect(
            Expectation::matching(request::method_path("GET", "/foo"))
                .times(2)
                // Response body is 100 bytes in size.
                .respond_with(responders::status_code(200).body(vec![0; 100])),
        );

        let client = SecureHttpClient::new(tls_config_with_system_roots()?, 10);
        let counting_http_client = CountingHttpClient::new(Arc::new(client));

        let mut resp = counting_http_client
            .get(&test_server.url_str("/foo"))
            .await?;
        // Consume the stream so we trigger a count.
        while (resp.body_mut().next().await).is_some() {}
        assert_eq!(
            100,
            counting_http_client
                .bytes_downloaded
                .load(Ordering::Relaxed)
        );

        let mut resp = counting_http_client
            .get(&test_server.url_str("/foo"))
            .await?;
        // Consume the stream so we trigger a count.
        while (resp.body_mut().next().await).is_some() {}
        assert_eq!(
            200,
            counting_http_client
                .bytes_downloaded
                .load(Ordering::Relaxed)
        );

        Ok(())
    }
}


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

use std::ffi::OsString;

use anyhow::Context as _;
use buck2_core::is_fbcode_build;

/// Find TLS certs.
///
/// Return `None` in Cargo or open source builds.
/// Return `Err` if certificates cannot be found in internal buck2 builds.
pub fn find_tls_cert() -> anyhow::Result<Option<OsString>> {
    if is_fbcode_build() {
        find_meta_internal_tls_cert().map(Some)
    } else {
        Ok(None)
    }
}

/// Error in open source or Cargo builds.
pub fn find_meta_internal_tls_cert() -> anyhow::Result<OsString> {
    let cert;

    #[cfg(fbcode_build)]
    {
        cert = ::find_certs::find_tls_cert();
    }

    #[cfg(not(fbcode_build))]
    {
        if buck2_core::is_open_source() {
            cert = Err(anyhow::anyhow!(
                "Cannot access internal Meta certs in open source builds"
            ));
        } else {
            cert = Err(anyhow::anyhow!("Cannot access certs in Cargo builds"));
        }
    }

    cert.context("Error finding a cert")
}