/*
 * 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::io;
use std::io::Cursor;
use std::io::Read;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;

use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use pin_project::pin_project;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::ReadBuf;
use tokio_util::codec::FramedRead;

/// This trait represents the ability to transition a Reader (R) to a Drainer (Self). Both are
/// [AsyncRead], but we expect the Drainer (which implements this trait) to not wait longer for
/// more data to be produced.
pub trait DrainerFromReader<R> {
    fn from_reader(reader: R) -> Self;
}

/// This trait represents a AsyncRead that can be told to interrupt (and transition to draining).
pub trait InterruptNotifiable {
    fn notify_interrupt(self: Pin<&mut Self>);
}

/// An [AsyncRead] that can be interrupted. The underlying [AsyncRead] will be polled while the
/// [Future] provided in interrupt hasn't completed. When said future completes, then we will
/// proceed to "drain" the reader, which means reading the data that is there but not waiting for
/// any further data to get written.
#[pin_project]
pub struct InterruptibleAsyncRead<R, D> {
    state: InterruptibleAsyncReadState<R, D>,
}

enum InterruptibleAsyncReadState<R, D> {
    Reading(R),
    Draining(D),
}

impl<R, D> InterruptibleAsyncRead<R, D>
where
    D: DrainerFromReader<R>,
{
    pub fn new(reader: R) -> Self {
        Self {
            state: InterruptibleAsyncReadState::Reading(reader),
        }
    }
}

impl<R, D> AsyncRead for InterruptibleAsyncRead<R, D>
where
    R: AsyncRead + Unpin,
    D: AsyncRead + Unpin + DrainerFromReader<R>,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let mut this = self.project();

        match &mut this.state {
            InterruptibleAsyncReadState::Reading(reader) => Pin::new(reader).poll_read(cx, buf),
            InterruptibleAsyncReadState::Draining(drainer) => Pin::new(drainer).poll_read(cx, buf),
        }
    }
}

#[cfg(unix)]
mod unix_non_blocking_drainer {
    use std::os::unix::io::AsRawFd;
    use std::os::unix::prelude::RawFd;

    use nix::unistd;

    use super::*;

    /// Perform sync reads on an existing reader. We use this to bypass Tokio when draining an
    /// InterruptibleAsyncRead after it's been interrupted. This lets us ensure that even if Tokio
    /// hasn't completed `select()` on the pipe we want to drain, we'll still get to execute
    /// `read()` (and potentially get WouldBlock if there is nothing to read and the pipe isn't
    /// ready).
    pub struct UnixNonBlockingDrainer<R> {
        fd: RawFd,
        // Kept so this is dropped and closed properly.
        _owner: R,
    }

    impl<R> DrainerFromReader<R> for UnixNonBlockingDrainer<R>
    where
        R: AsRawFd + Send + 'static,
    {
        fn from_reader(reader: R) -> Self {
            UnixNonBlockingDrainer {
                fd: reader.as_raw_fd(),
                _owner: reader,
            }
        }
    }

    impl<R> AsyncRead for UnixNonBlockingDrainer<R> {
        fn poll_read(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<io::Result<()>> {
            Poll::Ready(
                match unistd::read(self.fd, buf.initialize_unfilled()).map_err(io::Error::from) {
                    Err(e) => {
                        if e.kind() == io::ErrorKind::WouldBlock {
                            tracing::debug!("Child did not close its pipe");
                            Ok(())
                        } else {
                            Err(e)
                        }
                    }
                    Ok(n) => {
                        buf.advance(n);
                        Ok(())
                    }
                },
            )
        }
    }
}

#[cfg(unix)]
pub use unix_non_blocking_drainer::*;

#[pin_project]
pub struct TimeoutDrainer<R> {
    state: TimeoutDrainerState,
    // To have a generic parameter like UnixNonBlockingDrainer does.
    _phantom: PhantomData<R>,
}

enum TimeoutDrainerState {
    Waiting(BoxFuture<'static, io::Result<Vec<u8>>>),
    Draining(Cursor<Vec<u8>>),
}

impl<R> DrainerFromReader<R> for TimeoutDrainer<R>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    fn from_reader(mut reader: R) -> Self {
        let fut = async move {
            let mut buff = Vec::new();
            {
                let do_read = reader.read_to_end(&mut buff);
                futures::pin_mut!(do_read);
                match tokio::time::timeout(Duration::from_secs(1), do_read).await {
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        tracing::debug!("Child did not close its pipe");
                    }
                    Ok(Ok(..)) => {}
                }
            }
            Ok(buff)
        };

        Self {
            state: TimeoutDrainerState::Waiting(fut.boxed()),
            _phantom: PhantomData,
        }
    }
}

impl<R> AsyncRead for TimeoutDrainer<R> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let mut this = self.project();

        loop {
            match &mut this.state {
                TimeoutDrainerState::Waiting(fut) => {
                    let bytes = futures::ready!(Pin::new(fut).poll(cx))?;
                    *this.state = TimeoutDrainerState::Draining(Cursor::new(bytes));
                }
                TimeoutDrainerState::Draining(bytes) => {
                    let n = Read::read(bytes, buf.initialize_unfilled())?;
                    buf.advance(n);
                    break Poll::Ready(Ok(()));
                }
            }
        }
    }
}

impl<R, D> InterruptNotifiable for InterruptibleAsyncRead<R, D>
where
    D: DrainerFromReader<R>,
{
    fn notify_interrupt(self: Pin<&mut Self>) {
        let this = self.project();

        take_mut::take(this.state, |state| match state {
            InterruptibleAsyncReadState::Reading(reader) => {
                InterruptibleAsyncReadState::Draining(D::from_reader(reader))
            }
            _ => unreachable!(),
        });
    }
}

impl<S1, S2> InterruptNotifiable for futures::stream::Select<S1, S2>
where
    S1: InterruptNotifiable,
    S2: InterruptNotifiable,
{
    fn notify_interrupt(self: Pin<&mut Self>) {
        let (s1, s2) = self.get_pin_mut();
        s1.notify_interrupt();
        s2.notify_interrupt();
    }
}

impl<S, F> InterruptNotifiable for futures::stream::Map<S, F>
where
    S: InterruptNotifiable,
{
    fn notify_interrupt(self: Pin<&mut Self>) {
        self.get_pin_mut().notify_interrupt();
    }
}

impl<T, D> InterruptNotifiable for FramedRead<T, D>
where
    T: InterruptNotifiable,
{
    fn notify_interrupt(self: Pin<&mut Self>) {
        self.get_pin_mut().notify_interrupt();
    }
}

impl<T> InterruptNotifiable for futures::stream::Empty<T> {
    fn notify_interrupt(self: Pin<&mut Self>) {}
}

impl<L, R> InterruptNotifiable for futures::future::Either<L, R>
where
    L: InterruptNotifiable,
    R: InterruptNotifiable,
{
    fn notify_interrupt(self: Pin<&mut Self>) {
        match self.as_pin_mut() {
            futures::future::Either::Left(s) => s.notify_interrupt(),
            futures::future::Either::Right(s) => s.notify_interrupt(),
        }
    }
}

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

    mod interruptible_async_read {
        use std::sync::Arc;
        use std::sync::Mutex;

        use assert_matches::assert_matches;
        use dupe::Dupe;
        use futures::stream::Stream;
        use tokio_util::codec::BytesCodec;
        use tokio_util::codec::FramedRead;

        use super::*;

        #[derive(Copy, Clone, Dupe)]
        enum StubAsyncReadState {
            Done,
            Pending,
            Ready(u8),
        }

        /// A stub AsyncRead implementation
        #[derive(Clone, Dupe)]
        #[pin_project]
        struct StubAsyncRead {
            state: Arc<Mutex<StubAsyncReadState>>,
            drainer: bool,
        }

        impl DrainerFromReader<StubAsyncRead> for StubAsyncRead {
            fn from_reader(reader: StubAsyncRead) -> Self {
                Self {
                    state: reader.state,
                    drainer: true,
                }
            }
        }

        impl StubAsyncRead {
            fn new() -> Self {
                Self {
                    state: Arc::new(Mutex::new(StubAsyncReadState::Pending)),
                    drainer: false,
                }
            }

            /// Push a byte to this. Doesn't check that this makes sense.
            fn push(&self, byte: u8) {
                let mut state = self.state.lock().unwrap();
                *state = StubAsyncReadState::Ready(byte);
            }

            /// Set this to done. Also doesn't check that this makes sense.
            fn done(&self) {
                let mut state = self.state.lock().unwrap();
                *state = StubAsyncReadState::Done;
            }
        }

        impl AsyncRead for StubAsyncRead {
            /// A poll_read implementation for our stub. Note that this never stores any wakers,
            /// that's OK because our tests do explicit poll!().
            fn poll_read(
                self: Pin<&mut Self>,
                _: &mut Context<'_>,
                buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                let this = self.project();
                let mut state = this.state.lock().unwrap();

                match *state {
                    StubAsyncReadState::Done => Poll::Ready(Ok(())),
                    StubAsyncReadState::Pending => {
                        if *this.drainer {
                            Poll::Ready(Ok(()))
                        } else {
                            Poll::Pending
                        }
                    }
                    StubAsyncReadState::Ready(b) => {
                        buf.put_slice(&[b]);
                        *state = StubAsyncReadState::Pending;
                        Poll::Ready(Ok(()))
                    }
                }
            }
        }

        #[pin_project]
        struct ReadAll<R, D> {
            #[pin]
            inner: FramedRead<InterruptibleAsyncRead<R, D>, BytesCodec>,

            buff: Vec<u8>,
        }

        impl<R, D> InterruptNotifiable for ReadAll<R, D>
        where
            D: DrainerFromReader<R>,
        {
            fn notify_interrupt(self: Pin<&mut Self>) {
                let this = self.project();
                this.inner.get_pin_mut().notify_interrupt()
            }
        }

        impl<R, D> Future for ReadAll<R, D>
        where
            R: AsyncRead + Unpin,
            D: AsyncRead + Unpin + DrainerFromReader<R>,
        {
            type Output = anyhow::Result<Vec<u8>>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let mut this = self.project();

                loop {
                    match futures::ready!(this.inner.as_mut().poll_next(cx)?) {
                        Some(bytes) => this.buff.extend(bytes.as_ref()),
                        None => return Poll::Ready(Ok(std::mem::take(this.buff))),
                    }
                }
            }
        }

        fn read_all<R, D>(read: InterruptibleAsyncRead<R, D>) -> ReadAll<R, D>
        where
            R: AsyncRead + Unpin,
            D: AsyncRead + Unpin + DrainerFromReader<R>,
        {
            ReadAll {
                inner: FramedRead::new(read, BytesCodec::new()),
                buff: Vec::new(),
            }
        }

        #[tokio::test]
        async fn test_drain_eof() {
            let read = StubAsyncRead::new();
            let interruptible = InterruptibleAsyncRead::<_, StubAsyncRead>::new(read.dupe());

            let drain = read_all(interruptible);
            futures::pin_mut!(drain);

            // No bytes, it's pending.
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Add a byte, still pending, because more may come.
            read.push(b'f');
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Mark our reader done, we expect this to be ready.
            read.done();
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Ready(Ok(ret)) => {
                assert_eq!(&ret, b"f");
            });
        }

        #[tokio::test]
        async fn test_drain_interrupt() {
            let read = StubAsyncRead::new();
            let interruptible = InterruptibleAsyncRead::<_, StubAsyncRead>::new(read.dupe());

            let drain = read_all(interruptible);
            futures::pin_mut!(drain);

            // No bytes, it's pending.
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Add a byte, still pending, because more may come.
            read.push(b'f');
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Close the reader.
            drain.as_mut().notify_interrupt();

            // Mark our reader done, we expect this to be ready.
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Ready(Ok(ret)) => {
                assert_eq!(&ret, b"f");
            });
        }

        #[tokio::test]
        async fn test_drain_finish() {
            let read = StubAsyncRead::new();
            let interruptible = InterruptibleAsyncRead::<_, StubAsyncRead>::new(read.dupe());

            let drain = read_all(interruptible);
            futures::pin_mut!(drain);

            // No bytes, it's pending.
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Add a byte, still pending, because more may come.
            read.push(b'f');
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Pending);

            // Close the reader.
            drain.as_mut().notify_interrupt();

            // But! Add more stuff.
            read.push(b'o');

            // We now expect this to read what's left.
            assert_matches!(futures::poll!(drain.as_mut()), Poll::Ready(Ok(ret)) => {
                assert_eq!(&ret, b"fo");
            });
        }
    }

    mod timeout_drainer {
        use bytes::Bytes;
        use futures::StreamExt;

        use super::*;

        #[tokio::test]
        async fn test_timeout_drainer() -> anyhow::Result<()> {
            // 64 bytes of a. Tokio allocates a 32 byte buffer for read_to_end so this is good to
            //    ensure we get 2 reads.
            let s = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";

            let reader = tokio_util::io::StreamReader::new(
                tokio_stream::iter(vec![io::Result::Ok(Bytes::from_static(s.as_bytes()))])
                    .chain(futures::stream::pending()),
            );

            let mut drainer = TimeoutDrainer::from_reader(reader);
            let mut buff = Vec::new();
            tokio::time::timeout(Duration::from_secs(5), drainer.read_to_end(&mut buff)).await??;
            assert_eq!(s.as_bytes(), buff.as_slice());

            Ok(())
        }
    }

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

        use buck2_util::process::async_background_command;

        use super::*;

        #[tokio::test]
        async fn test_unix_non_blocking_drainer() -> anyhow::Result<()> {
            // See above.
            let s = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";

            let mut command = async_background_command("sh");
            command
                .args(["-c", &format!("((sleep 10 && echo extra) &) && echo {}", s)])
                .stdin(Stdio::null())
                .stdout(Stdio::piped())
                .stderr(Stdio::null());
            let mut command = command.spawn()?;

            let stdout = command.stdout.take().unwrap();
            let mut drainer =
                unix_non_blocking_drainer::UnixNonBlockingDrainer::from_reader(stdout);

            command.wait().await?;

            let mut buff = String::new();
            drainer.read_to_string(&mut buff).await?;
            assert_eq!(format!("{}\n", s), buff.as_str());

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

mod interruptible_async_read;
pub mod status_decoder;

use std::borrow::Cow;
use std::io;
use std::path::Path;
use std::pin::Pin;
use std::process::Command;
use std::process::ExitStatus;
use std::process::Stdio;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;

use anyhow::Context as _;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_path::AbsPath;
use bytes::Bytes;
use futures::future::Future;
use futures::future::FutureExt;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use pin_project::pin_project;
use tokio::process::Child;
use tokio_util::codec::BytesCodec;
use tokio_util::codec::FramedRead;

use self::interruptible_async_read::InterruptNotifiable;
use self::interruptible_async_read::InterruptibleAsyncRead;
use self::status_decoder::DecodedStatus;
use self::status_decoder::DefaultStatusDecoder;
use self::status_decoder::StatusDecoder;

#[derive(Debug)]
pub enum GatherOutputStatus {
    /// Contains the exit code.
    Finished {
        exit_code: i32,
        execution_stats: Option<buck2_data::CommandExecutionStats>,
    },
    TimedOut(Duration),
    Cancelled,
    SpawnFailed(String),
}

impl From<DecodedStatus> for GatherOutputStatus {
    fn from(d: DecodedStatus) -> Self {
        match d {
            DecodedStatus::Status {
                exit_code,
                execution_stats,
            } => Self::Finished {
                exit_code,
                execution_stats,
            },
            DecodedStatus::SpawnFailed(v) => Self::SpawnFailed(v),
        }
    }
}

#[derive(Debug)]
pub enum CommandEvent {
    Stdout(Bytes),
    Stderr(Bytes),
    Exit(GatherOutputStatus),
}

enum StdioEvent {
    Stdout(Bytes),
    Stderr(Bytes),
}

impl From<StdioEvent> for CommandEvent {
    fn from(stdio: StdioEvent) -> Self {
        match stdio {
            StdioEvent::Stdout(bytes) => CommandEvent::Stdout(bytes),
            StdioEvent::Stderr(bytes) => CommandEvent::Stderr(bytes),
        }
    }
}

/// This stream will yield [CommandEvent] whenever we have something on stdout or stderr (this is
/// our stdio stream), and it'll finish up the stream with the exit status. This is basically like
/// a select, but with the exit guaranteed to come last.
#[pin_project]
struct CommandEventStream<Status, Stdio> {
    exit: Option<anyhow::Result<GatherOutputStatus>>,

    done: bool,

    #[pin]
    status: futures::future::Fuse<Status>,

    #[pin]
    stdio: futures::stream::Fuse<Stdio>,
}

impl<Status, Stdio> CommandEventStream<Status, Stdio>
where
    Status: Future,
    Stdio: Stream,
{
    fn new(status: Status, stdio: Stdio) -> Self {
        Self {
            exit: None,
            done: false,
            status: status.fuse(),
            stdio: stdio.fuse(),
        }
    }
}

impl<Status, Stdio> Stream for CommandEventStream<Status, Stdio>
where
    Status: Future<Output = anyhow::Result<GatherOutputStatus>>,
    Stdio: Stream<Item = anyhow::Result<StdioEvent>> + InterruptNotifiable,
{
    type Item = anyhow::Result<CommandEvent>;

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

        if *this.done {
            return Poll::Ready(None);
        }

        // This future is fused so it's guaranteed to be ready once. If it does, capture the exit
        // status, we'll return it later.
        if let Poll::Ready(status) = this.status.poll(cx) {
            *this.exit = Some(status);
            this.stdio.as_mut().get_pin_mut().notify_interrupt();
        }

        // This stream is also fused, so if it returns None, we'll know it's done for good and we'll
        // return the exit status if it's available.
        if let Some(stdio) = futures::ready!(this.stdio.poll_next(cx)) {
            return Poll::Ready(Some(stdio.map(|event| event.into())));
        }

        // If we got here that means the stream is done. If we have it we return, and if we don't
        // we report we're pending, because we'll have polled it already earlier.
        if let Some(exit) = this.exit.take() {
            *this.done = true;
            return Poll::Ready(Some(exit.map(CommandEvent::Exit)));
        }

        Poll::Pending
    }
}

pub async fn timeout_into_cancellation(
    timeout: Option<Duration>,
) -> anyhow::Result<GatherOutputStatus> {
    match timeout {
        Some(t) => {
            tokio::time::sleep(t).await;
            Ok(GatherOutputStatus::TimedOut(t))
        }
        None => futures::future::pending().await,
    }
}

pub fn stream_command_events<T>(
    child: io::Result<Child>,
    cancellation: T,
    decoder: impl StatusDecoder,
    kill_process: impl KillProcess,
    stream_stdio: bool,
) -> anyhow::Result<impl Stream<Item = anyhow::Result<CommandEvent>>>
where
    T: Future<Output = anyhow::Result<GatherOutputStatus>> + Send,
{
    let mut child = match child {
        Ok(child) => child,
        Err(e) => {
            let event = Ok(CommandEvent::Exit(GatherOutputStatus::SpawnFailed(
                e.to_string(),
            )));
            return Ok(futures::stream::once(futures::future::ready(event)).left_stream());
        }
    };

    let stdio = if stream_stdio {
        let stdout = child.stdout.take().context("Child stdout is not piped")?;
        let stderr = child.stderr.take().context("Child stderr is not piped")?;

        #[cfg(unix)]
        type Drainer<R> = self::interruptible_async_read::UnixNonBlockingDrainer<R>;

        // On Windows, for the time being we just give ourselves a timeout to finish reading.
        // Ideally this would perform a non-blocking read on self instead like we do on Unix.
        #[cfg(not(unix))]
        type Drainer<R> = self::interruptible_async_read::TimeoutDrainer<R>;

        let stdout = InterruptibleAsyncRead::<_, Drainer<_>>::new(stdout);
        let stderr = InterruptibleAsyncRead::<_, Drainer<_>>::new(stderr);
        let stdout = FramedRead::new(stdout, BytesCodec::new())
            .map(|data| anyhow::Ok(StdioEvent::Stdout(data?.freeze())));
        let stderr = FramedRead::new(stderr, BytesCodec::new())
            .map(|data| anyhow::Ok(StdioEvent::Stderr(data?.freeze())));

        futures::stream::select(stdout, stderr).left_stream()
    } else {
        futures::stream::empty().right_stream()
    };

    let status = async move {
        enum Outcome {
            Finished(ExitStatus),
            Cancelled(GatherOutputStatus),
        }

        // NOTE: This wrapping here is so that we release the borrow of `child` that stems from
        // `wait()` by the time we call kill_process a few lines down.
        let execute = async {
            let status = child.wait();
            futures::pin_mut!(status);
            futures::pin_mut!(cancellation);

            anyhow::Ok(match futures::future::select(status, cancellation).await {
                futures::future::Either::Left((status, _)) => Outcome::Finished(status?),
                futures::future::Either::Right((res, _)) => Outcome::Cancelled(res?),
            })
        };

        anyhow::Ok(match execute.await? {
            Outcome::Finished(status) => decoder.decode_status(status).await?.into(),
            Outcome::Cancelled(res) => {
                kill_process
                    .kill(&mut child)
                    .context("Failed to terminate child after timeout")?;

                decoder
                    .cancel()
                    .await
                    .context("Failed to cancel status decoder after timeout")?;

                // We just killed the child, so this should finish immediately. We should still call
                // this to release any process.
                child
                    .wait()
                    .await
                    .context("Failed to await child after kill")?;

                res
            }
        })
    };

    Ok(CommandEventStream::new(status, stdio).right_stream())
}

pub(crate) async fn decode_command_event_stream<S>(
    stream: S,
) -> anyhow::Result<(GatherOutputStatus, Vec<u8>, Vec<u8>)>
where
    S: Stream<Item = anyhow::Result<CommandEvent>>,
{
    futures::pin_mut!(stream);

    let mut stdout = Vec::<u8>::new();
    let mut stderr = Vec::<u8>::new();

    while let Some(event) = stream.try_next().await? {
        match event {
            CommandEvent::Stdout(bytes) => stdout.extend(&bytes),
            CommandEvent::Stderr(bytes) => stderr.extend(&bytes),
            CommandEvent::Exit(exit) => return Ok((exit, stdout, stderr)),
        }
    }

    Err(anyhow::Error::msg(
        "Stream did not yield CommandEvent::Exit",
    ))
}

pub async fn gather_output<T>(
    cmd: Command,
    cancellation: T,
) -> anyhow::Result<(GatherOutputStatus, Vec<u8>, Vec<u8>)>
where
    T: Future<Output = anyhow::Result<GatherOutputStatus>> + Send,
{
    let cmd = prepare_command(cmd);

    let child = spawn_retry_txt_busy(cmd, || tokio::time::sleep(Duration::from_millis(50))).await;
    let stream = stream_command_events(
        child,
        cancellation,
        DefaultStatusDecoder,
        DefaultKillProcess,
        true,
    )?;
    decode_command_event_stream(stream).await
}

/// Dependency injection for kill. We use this in testing.
pub trait KillProcess {
    fn kill(self, child: &mut Child) -> anyhow::Result<()>;
}

pub struct DefaultKillProcess;

impl KillProcess for DefaultKillProcess {
    fn kill(self, child: &mut Child) -> anyhow::Result<()> {
        let pid = match child.id() {
            Some(pid) => pid,
            None => {
                // Child just exited, so in this case we don't want to kill anything.
                return Ok(());
            }
        };
        tracing::info!("Killing process {}", pid);
        #[cfg(unix)]
        {
            if true {
                // On unix we want killpg, so we don't use the default impl.
                // We use `if true` here to do less conditional compilation
                // or conditional dependencies.
                return kill_process_impl(pid);
            }
        }
        // `start_kill` is just `std::process::Child::kill` on Windows.
        // Ignore the error because `kill` fails on Windows if the process has been terminated
        // even if we did not wait for it.
        let _ignore = child.start_kill();
        Ok(())
    }
}

#[cfg(unix)]
fn kill_process_impl(pid: u32) -> anyhow::Result<()> {
    use nix::sys::signal;
    use nix::sys::signal::Signal;
    use nix::unistd::Pid;

    let pid: i32 = pid.try_into().context("PID does not fit a i32")?;

    signal::killpg(Pid::from_raw(pid), Signal::SIGKILL)
        .with_context(|| format!("Failed to kill process {}", pid))
}

/// Unify the the behavior of using a relative path for the executable between Unix and Windows. On
/// UNIX, the path is understood to be relative to the cwd of the *spawned process*, whereas on
/// Windows, it's relative ot the cwd of the *spawning* process.
///
/// Here, we unify the two behaviors since we always run our subprocesses with a known cwd: we
/// check if the executable actually exists relative to said cwd, and if it does, we use that.
pub fn maybe_absolutize_exe<'a>(
    exe: &'a (impl AsRef<Path> + ?Sized),
    spawned_process_cwd: &'_ AbsPath,
) -> anyhow::Result<Cow<'a, Path>> {
    let exe = exe.as_ref();

    let abs = spawned_process_cwd.join(exe);
    if fs_util::try_exists(&abs).context("Error absolute-izing executable")? {
        return Ok(abs.into_path_buf().into());
    }

    Ok(exe.into())
}

pub fn prepare_command(mut cmd: Command) -> tokio::process::Command {
    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;
        cmd.process_group(0);
    }

    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    cmd.into()
}

/// fork-exec is a bit tricky in a busy process. We often have files open to writing just prior to
/// executing them (as we download from RE), and many processes being spawned concurrently. We do
/// close the fds properly before the exec, but what can happn is:
///
/// - Some thread forks
/// - We close the file. At this time we don't have it open, but the forked process does.
/// - We try to exec the file. This fails because the file is open for writing (by the forked
/// process).
/// - The forked process execs. At this point the file is closed (because everything is CLOEXEC).
///
/// The window during which the forked process holds the fd is small, so retrying a couple times
/// here should let us make this work.
///
/// The more correct solution for this here would be to start a fork server in a separate process
/// when we start.  However, until we get there, this should do the trick.
async fn spawn_retry_txt_busy<F, D>(
    mut cmd: tokio::process::Command,
    mut delay: F,
) -> io::Result<Child>
where
    F: FnMut() -> D,
    D: Future<Output = ()>,
{
    let mut attempts = 10;

    loop {
        let res = cmd.spawn();

        let res_errno = res.as_ref().map_err(|e| e.raw_os_error());
        let is_txt_busy = matches!(res_errno, Err(Some(libc::ETXTBSY)));

        if attempts == 0 || !is_txt_busy {
            return res;
        }

        delay().await;

        attempts -= 1;
    }
}

#[cfg(test)]
mod tests {
    use std::str;
    use std::sync::Arc;
    use std::sync::Mutex;
    use std::time::Instant;

    use assert_matches::assert_matches;
    use buck2_util::process::async_background_command;
    use buck2_util::process::background_command;
    use dupe::Dupe;

    use super::*;

    #[tokio::test]
    async fn test_gather_output() -> anyhow::Result<()> {
        let mut cmd = if cfg!(windows) {
            background_command("powershell")
        } else {
            background_command("sh")
        };
        cmd.args(["-c", "echo hello"]);

        let (status, stdout, stderr) = gather_output(cmd, futures::future::pending()).await?;
        assert!(matches!(status, GatherOutputStatus::Finished { exit_code, .. } if exit_code == 0));
        assert_eq!(str::from_utf8(&stdout)?.trim(), "hello");
        assert_eq!(stderr, b"");

        Ok(())
    }

    #[tokio::test]
    async fn test_gather_does_not_wait_for_children() -> anyhow::Result<()> {
        // If we wait for sleep, this will time out.
        let mut cmd = if cfg!(windows) {
            background_command("powershell")
        } else {
            background_command("sh")
        };
        if cfg!(windows) {
            cmd.args([
                "-c",
                "Start-Job -ScriptBlock {sleep 10} | Out-Null; echo hello",
            ]);
        } else {
            cmd.args(["-c", "(sleep 10 &) && echo hello"]);
        }

        let timeout = if cfg!(windows) { 5 } else { 1 };
        let (status, stdout, stderr) = gather_output(
            cmd,
            timeout_into_cancellation(Some(Duration::from_secs(timeout))),
        )
        .await?;
        assert!(matches!(status, GatherOutputStatus::Finished { exit_code, .. } if exit_code == 0));
        assert_eq!(str::from_utf8(&stdout)?.trim(), "hello");
        assert_eq!(stderr, b"");

        Ok(())
    }

    #[tokio::test]
    async fn test_gather_output_timeout() -> anyhow::Result<()> {
        let now = Instant::now();

        let mut cmd = if cfg!(windows) {
            background_command("powershell")
        } else {
            background_command("sh")
        };
        cmd.args(["-c", "echo hello; sleep 10; echo bye"]);

        let timeout = if cfg!(windows) { 5 } else { 1 };
        let (status, stdout, stderr) = gather_output(
            cmd,
            timeout_into_cancellation(Some(Duration::from_secs(timeout))),
        )
        .await?;
        assert!(matches!(status, GatherOutputStatus::TimedOut(..)));
        assert_eq!(str::from_utf8(&stdout)?.trim(), "hello");
        assert_eq!(stderr, b"");

        assert!(now.elapsed() < Duration::from_secs(9)); // Lots of leeway here.

        Ok(())
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn test_spawn_retry_txt_busy() -> anyhow::Result<()> {
        use futures::future;
        use tokio::fs::OpenOptions;
        use tokio::io::AsyncWriteExt;

        let tempdir = tempfile::tempdir()?;
        let bin = tempdir.path().join("bin");

        let mut file = OpenOptions::new()
            .mode(0o755)
            .write(true)
            .create(true)
            .open(&bin)
            .await?;

        file.write_all(b"#!/bin/bash\ntrue\n").await?;

        let cmd = async_background_command(&bin);
        let mut child = spawn_retry_txt_busy(cmd, {
            let mut file = Some(file);
            move || {
                file.take();
                future::ready(())
            }
        })
        .await?;

        let status = child.wait().await?;
        assert_eq!(status.code(), Some(0));

        Ok(())
    }

    #[tokio::test]
    async fn test_spawn_retry_other_error() -> anyhow::Result<()> {
        let tempdir = tempfile::tempdir()?;
        let bin = tempdir.path().join("bin"); // Does not actually exist

        let cmd = async_background_command(&bin);
        let res = spawn_retry_txt_busy(cmd, || async { panic!("Should not be called!") }).await;
        assert!(res.is_err());

        Ok(())
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn test_kill_terminates_process_group() -> anyhow::Result<()> {
        use std::str::FromStr;

        use nix::errno::Errno;
        use nix::sys::signal;
        use nix::unistd::Pid;

        // This command will spawn 2 subprocesses (subshells) and print the PID of the 2nd shell.
        let mut cmd = background_command("sh");
        cmd.arg("-c").arg("( ( echo $$ && sleep 1000 ) )");
        let (_status, stdout, _stderr) =
            gather_output(cmd, timeout_into_cancellation(Some(Duration::from_secs(1)))).await?;
        let pid = i32::from_str(std::str::from_utf8(&stdout)?.trim())?;

        for _ in 0..10 {
            // This does rely on no PID reuse but the odds of PIDs wrapping around all the way to the
            // same PID we just used before we issue this kill seem low. So, we expect this to error
            // out.
            if matches!(signal::kill(Pid::from_raw(pid), None), Err(e) if e == Errno::ESRCH) {
                return Ok(());
            }

            // This is awkward but unfortunately the process does not immediately disappear.
            tokio::time::sleep(Duration::from_secs(1)).await;
        }

        Err(anyhow::anyhow!("PID did not exit: {}", pid))
    }

    #[tokio::test]
    async fn test_stream_command_events_ends() -> anyhow::Result<()> {
        let mut cmd = if cfg!(windows) {
            background_command("powershell")
        } else {
            background_command("sh")
        };
        cmd.args(["-c", "exit 0"]);

        let child = prepare_command(cmd).spawn();
        let mut events = stream_command_events(
            child,
            futures::future::pending(),
            DefaultStatusDecoder,
            DefaultKillProcess,
            true,
        )?
        .boxed();
        assert_matches!(events.next().await, Some(Ok(CommandEvent::Exit(..))));
        assert_matches!(futures::poll!(events.next()), Poll::Ready(None));
        Ok(())
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn test_signal_exit_code() -> anyhow::Result<()> {
        use nix::sys::signal::Signal;

        let mut cmd = background_command("sh");
        cmd.arg("-c").arg("kill -KILL \"$$\"");
        let (status, _stdout, _stderr) = gather_output(cmd, futures::future::pending()).await?;

        assert_matches!(
            status,
            GatherOutputStatus::Finished { exit_code, .. } if exit_code == 128 + Signal::SIGKILL as i32
        );

        Ok(())
    }

    #[tokio::test]
    async fn timeout_kills_before_dropping_decoder() -> anyhow::Result<()> {
        struct Kill {
            killed: Arc<Mutex<bool>>,
        }

        impl KillProcess for Kill {
            fn kill(self, child: &mut Child) -> anyhow::Result<()> {
                *self.killed.lock().unwrap() = true;

                // We still need to kill the process. On Windows in particular our test will hang
                // if we do not.
                DefaultKillProcess.kill(child)
            }
        }

        struct Decoder {
            killed: Arc<Mutex<bool>>,
            cancelled: Arc<Mutex<bool>>,
        }

        #[async_trait::async_trait]
        impl StatusDecoder for Decoder {
            async fn decode_status(self, _status: ExitStatus) -> anyhow::Result<DecodedStatus> {
                panic!("Should not be called in this test since we timeout")
            }

            async fn cancel(self) -> anyhow::Result<()> {
                assert!(*self.killed.lock().unwrap());
                *self.cancelled.lock().unwrap() = true;
                Ok(())
            }
        }

        let killed = Arc::new(Mutex::new(false));
        let cancelled = Arc::new(Mutex::new(false));

        let mut cmd = if cfg!(windows) {
            background_command("powershell")
        } else {
            background_command("sh")
        };
        cmd.args(["-c", "sleep 10000"]);

        let mut cmd = prepare_command(cmd);
        let child = cmd.spawn();

        let stream = stream_command_events(
            child,
            timeout_into_cancellation(Some(Duration::from_secs(1))),
            Decoder {
                killed: killed.dupe(),
                cancelled: cancelled.dupe(),
            },
            Kill {
                killed: killed.dupe(),
            },
            true,
        )?;

        let (status, _stdout, _stderr) = decode_command_event_stream(stream).await?;
        assert!(matches!(status, GatherOutputStatus::TimedOut(..)));

        assert!(*killed.lock().unwrap());
        assert!(*cancelled.lock().unwrap());

        Ok(())
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn test_no_stdio_stream_command_events() -> anyhow::Result<()> {
        let mut cmd = background_command("sh");
        cmd.args(["-c", "echo hello"]);

        let mut cmd = prepare_command(cmd);
        let tempdir = tempfile::tempdir()?;
        let stdout = tempdir.path().join("stdout");
        cmd.stdout(std::fs::File::create(stdout.clone())?);

        let child = cmd.spawn();
        let mut events = stream_command_events(
            child,
            futures::future::pending(),
            DefaultStatusDecoder,
            DefaultKillProcess,
            false,
        )?
        .boxed();
        assert_matches!(events.next().await, Some(Ok(CommandEvent::Exit(..))));
        assert_matches!(futures::poll!(events.next()), Poll::Ready(None));

        assert_matches!(tokio::fs::read_to_string(stdout).await?.as_str(), "hello\n");

        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::process::ExitStatus;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_miniperf_proto::MiniperfOutput;

pub enum DecodedStatus {
    /// An actual status.
    Status {
        exit_code: i32,
        execution_stats: Option<buck2_data::CommandExecutionStats>,
    },

    /// Spawn failed, provide the error.
    SpawnFailed(String),
}

#[async_trait]
pub trait StatusDecoder {
    /// Status decoders receive the exit status of the command we ran, but they might also obtain
    /// information out of band to obtain a different exit status.
    async fn decode_status(self, status: ExitStatus) -> anyhow::Result<DecodedStatus>;

    /// Notify this decoder that it will not be used.
    async fn cancel(self) -> anyhow::Result<()>;
}

pub struct DefaultStatusDecoder;

#[async_trait]
impl StatusDecoder for DefaultStatusDecoder {
    async fn decode_status(self, status: ExitStatus) -> anyhow::Result<DecodedStatus> {
        Ok(DecodedStatus::Status {
            exit_code: default_decode_exit_code(status),
            execution_stats: None,
        })
    }

    async fn cancel(self) -> anyhow::Result<()> {
        Ok(())
    }
}

pub fn default_decode_exit_code(status: ExitStatus) -> i32 {
    let exit_code;

    #[cfg(unix)]
    {
        // Shell convention on UNIX is to return 128 + signal number on a signal exit,
        // so we emulate this here.
        use std::os::unix::process::ExitStatusExt;
        exit_code = status.code().or_else(|| Some(128 + status.signal()?));
    }

    #[cfg(not(unix))]
    {
        exit_code = status.code();
    }

    exit_code.unwrap_or(-1)
}

pub struct MiniperfStatusDecoder {
    out_path: AbsNormPathBuf,
}

impl MiniperfStatusDecoder {
    pub fn new(out_path: AbsNormPathBuf) -> Self {
        Self { out_path }
    }
}

#[async_trait]
impl StatusDecoder for MiniperfStatusDecoder {
    async fn decode_status(self, status: ExitStatus) -> anyhow::Result<DecodedStatus> {
        if !status.success() {
            return Ok(DecodedStatus::Status {
                exit_code: default_decode_exit_code(status),
                execution_stats: None,
            });
        }

        let status = tokio::fs::read(&self.out_path).await.with_context(|| {
            format!(
                "Error reading miniperf output at `{}`",
                self.out_path.display()
            )
        })?;

        tokio::fs::remove_file(&self.out_path)
            .await
            .with_context(|| format!("Error removing miniperf output at `{}`", self.out_path))?;

        let status = bincode::deserialize::<MiniperfOutput>(&status)
            .with_context(|| format!("Invalid miniperf output at `{}`", self.out_path.display()))?;

        match status.raw_exit_code {
            #[allow(unused_variables)]
            Ok(v) => {
                #[cfg(unix)]
                {
                    use std::os::unix::process::ExitStatusExt;
                    let exit_code = default_decode_exit_code(ExitStatus::from_raw(v));
                    let execution_stats =
                        status
                            .counters
                            .map(|counters| buck2_data::CommandExecutionStats {
                                cpu_instructions_user: Some(
                                    counters.user_instructions.adjusted_count(),
                                ),
                                cpu_instructions_kernel: Some(
                                    counters.kernel_instructions.adjusted_count(),
                                ),
                            });

                    if let Err(e) = execution_stats.as_ref() {
                        // TODO @torozco: report this in the event log? Might be verbose for little
                        // value.
                        tracing::debug!("Miniperf stats not available: {}", e);
                    }

                    Ok(DecodedStatus::Status {
                        exit_code,
                        execution_stats: execution_stats.ok(),
                    })
                }

                #[cfg(not(unix))]
                {
                    Err(anyhow::anyhow!("Attempted to use Miniperf output off-UNIX"))
                }
            }
            Err(e) => Ok(DecodedStatus::SpawnFailed(e)),
        }
    }

    async fn cancel(self) -> anyhow::Result<()> {
        let res = tokio::fs::remove_file(&self.out_path).await;

        match res {
            Ok(()) => Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(e) => Err(anyhow::Error::from(e).context(format!(
                "Error removing miniperf output at `{}`",
                self.out_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::os::unix::io::FromRawFd;
use std::os::unix::io::RawFd;
use std::os::unix::net::UnixStream as StdUnixStream;
use std::sync::Arc;

use anyhow::Context as _;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_forkserver_proto::forkserver_server;
use buck2_grpc::DuplexChannel;
use tokio::net::UnixStream;

use super::service::UnixForkserverService;

pub async fn run_forkserver(
    fd: RawFd,
    log_reload_handle: Arc<dyn LogConfigurationReloadHandle>,
    state_dir: AbsNormPathBuf,
) -> anyhow::Result<()> {
    // SAFETY: At worst, we just read (or close) the wrong FD.
    let io = UnixStream::from_std(unsafe { StdUnixStream::from_raw_fd(fd) })
        .expect("Failed to create io");

    let io = {
        let (read, write) = tokio::io::split(io);
        DuplexChannel::new(read, write)
    };

    let service = UnixForkserverService::new(log_reload_handle, &state_dir)
        .context("Failed to create UnixForkserverService")?;

    let router = tonic::transport::Server::builder().add_service(
        forkserver_server::ForkserverServer::new(service)
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX),
    );

    buck2_grpc::spawn_oneshot(io, router)
        .into_join_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.
 */

use std::ffi::OsStr;
use std::io;
use std::os::unix::io::AsRawFd;
use std::os::unix::process::CommandExt;
use std::process::Stdio;

use anyhow::Context;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_util::process::background_command;
use tokio::net::UnixStream;
use tokio::process::Command;

use crate::client::ForkserverClient;

pub async fn launch_forkserver(
    exe: impl AsRef<OsStr>,
    args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    state_dir: &AbsNormPath,
) -> anyhow::Result<ForkserverClient> {
    let (client_io, server_io) =
        UnixStream::pair().context("Failed to create fork server channel")?;

    let server_io = server_io
        .into_std()
        .context("Failed to convert server_io to std")?;

    let exe = exe.as_ref();

    let mut command = background_command(exe);
    command
        .stdin(Stdio::null())
        .stdout(Stdio::inherit()) // TODO
        .stderr(Stdio::inherit()) // TODO
        .arg0("(buck2-forkserver)")
        .args(args)
        .arg("--fd")
        .arg(server_io.as_raw_fd().to_string())
        .arg("--state-dir")
        .arg(state_dir.as_path());

    let fds = [server_io.as_raw_fd()];

    unsafe {
        command.pre_exec(move || {
            // Clear CLOEXEC on the 2 FDs we intend to pass. It's better to set CLOEXEC and
            // clear it specifically here since that ensure that if we spawn anything else
            // concurrently, those FDs will be released on exec.

            for fd in &fds {
                let flags = libc::fcntl(*fd, libc::F_GETFD);
                if flags < 0 {
                    return Err(io::Error::last_os_error());
                }

                if libc::fcntl(*fd, libc::F_SETFD, flags & !libc::FD_CLOEXEC) < 0 {
                    return Err(io::Error::last_os_error());
                }
            }

            Ok(())
        });
    }

    let child = Command::from(command)
        .spawn()
        .with_context(|| format!("Failed to start Forkserver `{}`", exe.to_string_lossy()))?;

    let channel = buck2_grpc::make_channel(client_io, "forkserver")
        .await
        .context("Error connecting to Forkserver")?;

    Ok(ForkserverClient::new(child, channel))
}


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

mod command;
mod launch;
mod service;

pub use command::run_forkserver;
pub use launch::launch_forkserver;


/*
 * 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::OsStr;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Write;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;

use anyhow::Context as _;
use buck2_common::convert::ProstDurationExt;
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::abs_path::AbsPath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_forkserver_proto::forkserver_server::Forkserver;
use buck2_forkserver_proto::CommandRequest;
use buck2_forkserver_proto::RequestEvent;
use buck2_forkserver_proto::SetLogFilterRequest;
use buck2_forkserver_proto::SetLogFilterResponse;
use buck2_grpc::to_tonic;
use buck2_util::process::background_command;
use futures::future::select;
use futures::future::FutureExt;
use futures::stream::Stream;
use futures::stream::StreamExt;
use rand::distributions::Alphanumeric;
use rand::distributions::DistString;
use tonic::Request;
use tonic::Response;
use tonic::Status;
use tonic::Streaming;

use crate::convert::encode_event_stream;
use crate::run::maybe_absolutize_exe;
use crate::run::prepare_command;
use crate::run::status_decoder::DefaultStatusDecoder;
use crate::run::status_decoder::MiniperfStatusDecoder;
use crate::run::stream_command_events;
use crate::run::timeout_into_cancellation;
use crate::run::DefaultKillProcess;
use crate::run::GatherOutputStatus;

// Not quite BoxStream: it has to be Sync (...)
type RunStream =
    Pin<Box<dyn Stream<Item = Result<buck2_forkserver_proto::CommandEvent, Status>> + Send>>;

pub struct UnixForkserverService {
    log_reload_handle: Arc<dyn LogConfigurationReloadHandle>,

    /// State for Miniperf.
    miniperf: Option<MiniperfContainer>,
}

impl UnixForkserverService {
    pub fn new(
        log_reload_handle: Arc<dyn LogConfigurationReloadHandle>,
        state_dir: &AbsNormPath,
    ) -> anyhow::Result<Self> {
        let miniperf = MiniperfContainer::new(state_dir)?;

        Ok(Self {
            log_reload_handle,
            miniperf,
        })
    }
}

#[async_trait::async_trait]
impl Forkserver for UnixForkserverService {
    type RunStream = RunStream;

    async fn run(
        &self,
        req: Request<Streaming<RequestEvent>>,
    ) -> Result<Response<Self::RunStream>, Status> {
        to_tonic(async move {
            let mut stream = req.into_inner();

            let msg = stream
                .message()
                .await?
                .and_then(|m| m.data)
                .and_then(|m| m.into_command_request())
                .context("RequestEvent was not a CommandRequest!")?;

            let cancel = async move {
                stream
                    .message()
                    .await?
                    .and_then(|m| m.data)
                    .and_then(|m| m.into_cancel_request())
                    .context("RequestEvent was not a CancelRequest!")?;

                anyhow::Ok(GatherOutputStatus::Cancelled)
            };

            let CommandRequest {
                exe,
                argv,
                env,
                cwd,
                timeout,
                enable_miniperf,
                std_redirects,
            } = msg;

            let exe = OsStr::from_bytes(&exe);
            let cwd = OsStr::from_bytes(&cwd.as_ref().context("Missing cwd")?.path);
            let cwd = AbsPath::new(Path::new(cwd)).context("Inalid cwd")?;
            let argv = argv.iter().map(|a| OsStr::from_bytes(a));
            let timeout = timeout
                .map(|t| t.try_into_duration())
                .transpose()
                .context("Invalid timeout")?;

            let exe = maybe_absolutize_exe(exe, cwd)?;

            let (mut cmd, miniperf_output) = match (enable_miniperf, &self.miniperf) {
                (true, Some(miniperf)) => {
                    let mut cmd = background_command(miniperf.miniperf.as_path());
                    let output_path = miniperf.allocate_output_path();
                    cmd.arg(output_path.as_path());
                    cmd.arg(exe.as_ref());
                    (cmd, Some(output_path))
                }
                _ => (background_command(exe.as_ref()), None),
            };

            cmd.current_dir(cwd);
            cmd.args(argv);

            {
                use buck2_forkserver_proto::env_directive::Data;

                for directive in env {
                    match directive.data.context("EnvDirective is missing data")? {
                        Data::Clear(..) => {
                            cmd.env_clear();
                        }
                        Data::Set(var) => {
                            cmd.env(OsStr::from_bytes(&var.key), OsStr::from_bytes(&var.value));
                        }
                        Data::Remove(var) => {
                            cmd.env_remove(OsStr::from_bytes(&var.key));
                        }
                    }
                }
            }

            let mut cmd = prepare_command(cmd);
            let stream_stdio = std_redirects.is_none();
            if let Some(std_redirects) = std_redirects {
                cmd.stdout(File::create(OsStr::from_bytes(&std_redirects.stdout))?);
                cmd.stderr(File::create(OsStr::from_bytes(&std_redirects.stderr))?);
            }
            let child = cmd.spawn();

            let timeout = timeout_into_cancellation(timeout);

            let cancellation = select(timeout.boxed(), cancel.boxed()).map(|r| r.factor_first().0);

            let stream = match miniperf_output {
                Some(out) => stream_command_events(
                    child,
                    cancellation,
                    MiniperfStatusDecoder::new(out),
                    DefaultKillProcess,
                    stream_stdio,
                )?
                .left_stream(),
                None => stream_command_events(
                    child,
                    cancellation,
                    DefaultStatusDecoder,
                    DefaultKillProcess,
                    stream_stdio,
                )?
                .right_stream(),
            };

            let stream = encode_event_stream(stream);
            Ok(Box::pin(stream) as _)
        })
        .await
    }

    async fn set_log_filter(
        &self,
        req: Request<SetLogFilterRequest>,
    ) -> Result<Response<SetLogFilterResponse>, Status> {
        self.log_reload_handle
            .update_log_filter(&req.get_ref().log_filter)
            .context("Error updating forkserver filter")
            .map_err(|e| Status::invalid_argument(format!("{:#}", e)))?;

        Ok(Response::new(SetLogFilterResponse {}))
    }
}

struct MiniperfContainer {
    /// The Miniperf binary
    miniperf: AbsNormPathBuf,

    /// The directory where Miniperf outputs go.
    output_dir: AbsNormPathBuf,
}

impl MiniperfContainer {
    fn new(forkserver_state_dir: &AbsNormPath) -> anyhow::Result<Option<Self>> {
        let miniperf_bin: Option<&'static [u8]>;

        #[cfg(all(fbcode_build, target_os = "linux"))]
        {
            miniperf_bin = Some(include_bytes!("miniperf.bin").as_slice());
        }

        #[cfg(not(all(fbcode_build, target_os = "linux")))]
        {
            miniperf_bin = None;
        }

        let miniperf_bin = match miniperf_bin {
            Some(m) => m,
            None => return Ok(None),
        };

        let miniperf = forkserver_state_dir.join(ForwardRelativePath::unchecked_new("miniperf"));
        let output_dir = forkserver_state_dir.join(ForwardRelativePath::unchecked_new("out"));

        fs_util::remove_all(&miniperf)?;
        fs_util::remove_all(&output_dir)?;
        fs_util::create_dir_all(&output_dir)?;

        let mut opts = OpenOptions::new();
        opts.create_new(true);
        opts.write(true);

        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            opts.mode(0o755);
        }

        let mut miniperf_writer = opts
            .open(miniperf.as_path())
            .with_context(|| format!("Error opening: `{}`", miniperf.display()))?;

        miniperf_writer
            .write_all(miniperf_bin)
            .and_then(|()| miniperf_writer.flush())
            .with_context(|| format!("Error writing miniperf to `{}`", miniperf.display()))?;

        Ok(Some(Self {
            miniperf,
            output_dir,
        }))
    }

    fn allocate_output_path(&self) -> AbsNormPathBuf {
        let name = Alphanumeric.sample_string(&mut rand::thread_rng(), 16);
        self.output_dir
            .join(ForwardRelativePath::unchecked_new(&name))
    }
}


/*
 * 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::io;
use std::path::PathBuf;

fn main() -> io::Result<()> {
    let proto_files = &["forkserver.proto"];

    let data_include = if let Ok(value) = env::var("BUCK_HACK_DATA_PROTOC_INCLUDE") {
        let path = PathBuf::from(value);
        path.parent().unwrap().to_str().unwrap().to_owned()
    } else {
        "../buck2_data".to_owned()
    };

    buck2_protoc_dev::configure()
        .setup_protoc()
        .type_attribute(
            "buck.forkserver.RequestEvent.data",
            "#[derive(::derive_more::From, ::gazebo::variants::VariantName, ::gazebo::variants::UnpackVariants)]",
        )
        .type_attribute(
            "buck.forkserver.EnvDirective.data",
            "#[derive(::derive_more::From, ::gazebo::variants::VariantName, ::gazebo::variants::UnpackVariants)]",
        )
        .extern_path(".buck.data", "::buck2_data")
        .compile(proto_files, &[".", &data_include])
}


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

tonic::include_proto!("buck.forkserver");


/*
 * 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::io;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use anyhow::Context as _;
use futures::future;
use pin_project::pin_project;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::io::ReadBuf;
use tonic::transport::server::Connected;
use tonic::transport::Channel;
use tonic::transport::Endpoint;
use tonic::transport::Uri;
use tower::service_fn;

#[pin_project]
pub struct DuplexChannel<R, W> {
    #[pin]
    read: R,
    #[pin]
    write: W,
}

impl<R, W> DuplexChannel<R, W> {
    pub fn new(read: R, write: W) -> Self {
        Self { read, write }
    }
}

impl<R, W> AsyncWrite for DuplexChannel<R, W>
where
    W: AsyncWrite,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        self.project().write.poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().write.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().write.poll_shutdown(cx)
    }
}

impl<R, W> AsyncRead for DuplexChannel<R, W>
where
    R: AsyncRead,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        self.project().read.poll_read(cx, buf)
    }
}

impl<R, W> Connected for DuplexChannel<R, W> {
    type ConnectInfo = ();

    fn connect_info(&self) {}
}

/// Create a channel using a pre-existing I/O instance. This will not support reconnecting since
/// there is no way to establish connections here. We're just using one that already exists.
pub async fn make_channel<T>(io: T, name: &str) -> anyhow::Result<Channel>
where
    T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
{
    let mut io = Some(io);

    // NOTE: The uri here is only used to populate the requests we send. We don't actually connect
    // anywhere since we already have an I/O channel on hand.
    let channel = Endpoint::try_from(format!("http://{}.invalid", name))
        .context("Invalid endpoint")?
        .connect_with_connector(service_fn(move |_: Uri| {
            let io = io.take().context("Cannot reconnect after connection loss");
            future::ready(io)
        }))
        .await
        .context("Failed to create channel")?;

    Ok(channel)
}


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

mod channel;
mod server;
mod util;

pub use channel::make_channel;
pub use channel::DuplexChannel;
pub use server::spawn_oneshot;
pub use server::ServerHandle;
pub use util::to_tonic;


/*
 * 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 anyhow::Context;
use futures::future;
use futures::stream;
use futures::stream::StreamExt;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tonic::transport::server::Router;
use tonic::transport::server::Routes;
use tower::layer::Layer;
use tower::Service;

use self::connection_with_extra::ConnectionWithExtra;
use self::drop_notifier::DropNotifier;

pub struct ServerHandle {
    channel: DropNotifier,
    handle: JoinHandle<anyhow::Result<()>>,
}

impl ServerHandle {
    /// Tell the server to shutdown and wait for it to exit.
    pub async fn shutdown(self) -> anyhow::Result<()> {
        self.channel.notify_now();
        self.handle.await.context("Failed to join task")?
    }

    /// Obtain the JoinHandle to the task driving the server, without asking the server to
    /// shutdown.
    pub fn into_join_handle(self) -> JoinHandle<anyhow::Result<()>> {
        self.channel.cancel();
        self.handle
    }
}

pub fn spawn_oneshot<T, L>(io: T, router: Router<L>) -> ServerHandle
where
    T: AsyncRead + AsyncWrite + Send + Unpin + 'static + tonic::transport::server::Connected,

    L: Layer<Routes> + Send + 'static,
    L::Service: Service<
            http::Request<tonic::transport::Body>,
            Response = http::Response<tonic::body::BoxBody>,
        > + Clone
        + Send
        + 'static,
    <<L as Layer<Routes>>::Service as Service<http::Request<tonic::transport::Body>>>::Future:
        Send + 'static,
    <<L as Layer<Routes>>::Service as Service<http::Request<tonic::transport::Body>>>::Error:
        Into<Box<dyn std::error::Error + Send + Sync>> + Send,
{
    // We reserve 2 slots here: one for the connection and one for the ServerHandle's
    // shutdown.
    let (send, mut recv) = mpsc::channel(2);

    let io = ConnectionWithExtra {
        inner: io,
        notifier: DropNotifier::new(send.clone()),
    };

    let fut = async move {
        // NOTE: We immediately yield the connection we have on hand, then just report no further
        // connections going forward. This is important to ensure that the server doesn't just
        // exist when it sees no connections (it's a little weird that the server exists even with
        // open connections if they're not in the middle of a request, but thats's what it is).
        let incoming = stream::once(future::ready(io))
            .chain(stream::once(future::pending()))
            .map(std::io::Result::Ok);

        router
            .serve_with_incoming_shutdown(incoming, async move {
                let _ignored = recv.recv().await;
            })
            .await
            .context("Server exited with an error")?;

        Ok(())
    };

    let handle = tokio::task::spawn(fut);

    ServerHandle {
        channel: DropNotifier::new(send),
        handle,
    }
}

mod drop_notifier {
    use tokio::sync::mpsc::Sender;

    /// Upon being dropped, this notifies the sender is contains.
    pub struct DropNotifier {
        sender: Option<Sender<()>>,
    }

    impl DropNotifier {
        pub fn new(sender: Sender<()>) -> Self {
            Self {
                sender: Some(sender),
            }
        }

        /// Notify the sender this contains (this is equivalent to dropping the DropNotifier).
        pub fn notify_now(self) {
            drop(self)
        }

        /// Drop this DropNotifier without notifying. This is useful in cases where "defusing" the
        /// DropNotifier is desired.
        pub fn cancel(mut self) {
            self.sender.take();
        }
    }

    impl Drop for DropNotifier {
        fn drop(&mut self) {
            if let Some(sender) = self.sender.as_ref() {
                let _ = sender.try_send(());
            }
        }
    }
}

mod connection_with_extra {
    use std::io;
    use std::pin::Pin;
    use std::task::Context;
    use std::task::Poll;

    use pin_project::pin_project;
    use tokio::io::AsyncRead;
    use tokio::io::AsyncWrite;
    use tokio::io::ReadBuf;
    use tonic::transport::server::Connected;

    use super::DropNotifier;

    /// A connection (AsyncWrite + AsyncRead + Connected), with a field for a DropNotifier. This
    /// acts a connection, but lets us observe when it's dropped.
    #[pin_project]
    pub struct ConnectionWithExtra<I> {
        #[pin]
        pub inner: I,
        pub notifier: DropNotifier,
    }

    impl<I> AsyncWrite for ConnectionWithExtra<I>
    where
        I: AsyncWrite,
    {
        fn poll_write(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<Result<usize, io::Error>> {
            self.project().inner.poll_write(cx, buf)
        }

        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
            self.project().inner.poll_flush(cx)
        }

        fn poll_shutdown(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), io::Error>> {
            self.project().inner.poll_shutdown(cx)
        }
    }

    impl<I> AsyncRead for ConnectionWithExtra<I>
    where
        I: AsyncRead,
    {
        fn poll_read(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<io::Result<()>> {
            self.project().inner.poll_read(cx, buf)
        }
    }

    impl<I> Connected for ConnectionWithExtra<I>
    where
        I: Connected,
    {
        type ConnectInfo = I::ConnectInfo;

        fn connect_info(&self) -> Self::ConnectInfo {
            self.inner.connect_info()
        }
    }
}


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

pub async fn to_tonic<F, T>(fut: F) -> Result<tonic::Response<T>, tonic::Status>
where
    F: Future<Output = anyhow::Result<T>>,
{
    match fut.await {
        Ok(r) => Ok(tonic::Response::new(r)),
        Err(e) => Err(tonic::Status::unknown(format!("{:#}", e))),
    }
}


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

fn main() -> io::Result<()> {
    let proto_files = &["install.proto"];

    buck2_protoc_dev::configure()
        .setup_protoc()
        .compile(proto_files, &["."])
}


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

tonic::include_proto!("install");


/*
 * 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 buck2_util::late_binding::LateBinding;
use starlark::eval::Evaluator;

use crate::path::StarlarkPath;

pub static STARLARK_PATH_FROM_BUILD_CONTEXT: LateBinding<
    for<'a> fn(&Evaluator<'_, 'a>) -> anyhow::Result<StarlarkPath<'a>>,
> = LateBinding::new("STARLARK_PATH_FROM_BUILD_CONTEXT");


/*
 * 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 buck2_util::late_binding::LateBinding;
use starlark::environment::GlobalsBuilder;

/// Functions defined in `buck2_bxl` crate,
/// which are used to create the context for `.bxl` evaluation.
pub struct BxlFunctions {
    pub register_cli_args_struct: fn(&mut GlobalsBuilder),
    pub register_bxl_function: fn(&mut GlobalsBuilder),
    pub register_artifact_function: fn(&mut GlobalsBuilder),
    pub register_label_function: fn(&mut GlobalsBuilder),
    pub register_target_function: fn(&mut GlobalsBuilder),
    pub register_instant_function: fn(&mut GlobalsBuilder),
    pub register_error_handling_function: fn(&mut GlobalsBuilder),
}

/// Default version used in the `buck2` binary.
pub static BXL_FUNCTIONS: LateBinding<BxlFunctions> = LateBinding::new("BXL_FUNCTIONS");


/*
 * 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 buck2_core::target::label::TargetLabel;
use buck2_util::late_binding::LateBinding;
use starlark::eval::Evaluator;

pub static COERCE_TARGET_LABEL: LateBinding<
    fn(&mut Evaluator, &str) -> anyhow::Result<TargetLabel>,
> = LateBinding::new("COERCE_TARGET_LABEL");


/*
 * 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 starlark::environment::FrozenModule;
use starlark::environment::Module;
use starlark::eval::Evaluator;

/// Provides a starlark Evaluator.
pub trait StarlarkEvaluatorProvider {
    /// Creates an Evaluator for a module. The evaluator will be configured for instrumenting/profiling/debugging
    /// as appropriate.
    fn make<'v, 'a>(&mut self, module: &'v Module) -> anyhow::Result<Evaluator<'v, 'a>>;

    fn evaluation_complete(&mut self, eval: &mut Evaluator) -> anyhow::Result<()>;

    fn visit_frozen_module(&mut self, module: Option<&FrozenModule>) -> anyhow::Result<()>;
}

/// Trivial provider that just constructs an Evaluator. Useful for tests (but not necessarily limited to them).
pub struct StarlarkPassthroughProvider;

impl StarlarkEvaluatorProvider for StarlarkPassthroughProvider {
    fn make<'v, 'a>(&mut self, module: &'v Module) -> anyhow::Result<Evaluator<'v, 'a>> {
        Ok(Evaluator::new(module))
    }

    fn evaluation_complete(&mut self, _eval: &mut Evaluator) -> anyhow::Result<()> {
        Ok(())
    }

    fn visit_frozen_module(&mut self, _module: Option<&FrozenModule>) -> anyhow::Result<()> {
        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::sync::Arc;

use allocative::Allocative;
use buck2_core::bzl::ImportPath;
use buck2_util::collections::ordered_map::OrderedMap;
use derivative::Derivative;
use dupe::Dupe;
use starlark::codemap::FileSpan;
use starlark::environment::FrozenModule;
use starlark::eval::FileLoader;

use crate::path::OwnedStarlarkModulePath;
use crate::path::StarlarkModulePath;

#[derive(Default, Clone, Allocative, Debug)]
pub struct LoadedModules {
    pub map: OrderedMap<OwnedStarlarkModulePath, LoadedModule>,
}

impl LoadedModules {
    pub fn imports(&self) -> impl Iterator<Item = &ImportPath> {
        self.map.values().map(|module| {
            *module
                .path()
                .unpack_load_file()
                .expect("imports should only be bzl files")
        })
    }
}

pub trait LoadResolver {
    fn resolve_load(
        &self,
        path: &str,
        location: Option<&FileSpan>,
    ) -> anyhow::Result<OwnedStarlarkModulePath>;
}

pub struct ModuleDeps(pub Vec<LoadedModule>);

impl ModuleDeps {
    pub fn get_loaded_modules(&self) -> LoadedModules {
        let mut map = OrderedMap::with_capacity(self.0.len());
        for dep in &*self.0 {
            map.insert(dep.path().to_owned(), dep.dupe());
        }
        LoadedModules { map }
    }
}

#[derive(Clone, Dupe, Allocative, Debug)]
pub struct LoadedModule(Arc<LoadedModuleData>);

#[derive(Derivative, Allocative)]
#[derivative(Debug)]
struct LoadedModuleData {
    path: OwnedStarlarkModulePath,
    #[derivative(Debug = "ignore")]
    loaded_modules: LoadedModules,
    #[derivative(Debug = "ignore")]
    env: FrozenModule,
}

impl LoadedModule {
    pub fn new(
        path: OwnedStarlarkModulePath,
        loaded_modules: LoadedModules,
        env: FrozenModule,
    ) -> Self {
        Self(Arc::new(LoadedModuleData {
            path,
            loaded_modules,
            env,
        }))
    }

    pub fn loaded_modules(&self) -> &LoadedModules {
        &self.0.loaded_modules
    }

    pub fn imports(&self) -> impl Iterator<Item = &ImportPath> {
        self.0.loaded_modules.imports()
    }

    pub fn path(&self) -> StarlarkModulePath<'_> {
        self.0.path.borrow()
    }

    pub fn env(&self) -> &FrozenModule {
        &self.0.env
    }
}

pub struct InterpreterFileLoader {
    loaded_modules: LoadedModules,
    info: Arc<dyn LoadResolver>,
}

impl InterpreterFileLoader {
    pub fn new(loaded_modules: LoadedModules, info: Arc<dyn LoadResolver>) -> Self {
        Self {
            loaded_modules,
            info,
        }
    }
}

fn to_diagnostic(err: &anyhow::Error, id: &str) -> anyhow::Error {
    anyhow::anyhow!("UnknownError in {}: {}", id, err)
}

impl InterpreterFileLoader {
    /// Used for looking up modules by id.
    fn find_module(&self, id: StarlarkModulePath) -> anyhow::Result<&FrozenModule> {
        match self.loaded_modules.map.get(&id) {
            Some(v) => Ok(&v.0.env),
            None => Err(to_diagnostic(
                &anyhow::anyhow!(
                    "Should have had an env for {}. had <{:?}>",
                    id,
                    self.loaded_modules.map.keys().collect::<Vec<_>>()
                ),
                &id.to_string(),
            )),
        }
    }
}

impl FileLoader for InterpreterFileLoader {
    /// The Interpreter will call this to resolve and load imports for load()
    /// statements.
    fn load(&self, path: &str) -> anyhow::Result<FrozenModule> {
        match self.info.resolve_load(path, None) {
            Ok(import) => Ok(self.find_module(import.borrow())?.dupe()),
            Err(e) => Err(to_diagnostic(&e, path)),
        }
    }
}

#[cfg(test)]
mod tests {
    use starlark::environment::Module;

    use super::*;

    struct TestLoadResolver {}

    impl LoadResolver for TestLoadResolver {
        fn resolve_load(
            &self,
            path: &str,
            _location: Option<&FileSpan>,
        ) -> anyhow::Result<OwnedStarlarkModulePath> {
            match path {
                "//some/package:import.bzl" => Ok(OwnedStarlarkModulePath::LoadFile(
                    ImportPath::testing_new("root//some/package:import.bzl"),
                )),
                "cell1//next/package:import.bzl" => Ok(OwnedStarlarkModulePath::LoadFile(
                    ImportPath::testing_new("cell1//next/package:import.bzl"),
                )),
                "alias2//last/package:import.bzl" => Ok(OwnedStarlarkModulePath::LoadFile(
                    ImportPath::testing_new("cell2//last/package:import.bzl"),
                )),
                _ => Err(anyhow::anyhow!("error")),
            }
        }
    }

    fn resolver() -> Arc<TestLoadResolver> {
        Arc::new(TestLoadResolver {})
    }

    fn env(name: StarlarkModulePath) -> FrozenModule {
        let m = Module::new();
        m.set("name", m.heap().alloc(name.to_string()));
        m.freeze().unwrap()
    }

    fn loaded_modules() -> LoadedModules {
        let mut loaded_modules = LoadedModules::default();
        let resolver = resolver();

        let mut insert = |path| {
            let import_path = resolver.resolve_load(path, None).unwrap();
            let module = LoadedModule::new(
                import_path.clone(),
                LoadedModules::default(),
                env(import_path.borrow()),
            );
            loaded_modules.map.insert(import_path, module);
        };

        // Insert the same things that the TestLoadResolver supports.
        insert("//some/package:import.bzl");
        insert("cell1//next/package:import.bzl");
        insert("alias2//last/package:import.bzl");

        loaded_modules
    }

    #[test]
    fn no_resolution() -> anyhow::Result<()> {
        let path = "some//random:file.bzl".to_owned();
        let loader = InterpreterFileLoader::new(loaded_modules(), resolver());
        match loader.load(&path) {
            Ok(_) => panic!("Expected load failure for {}", path),
            Err(_) => {
                // TODO: verify the error is correct
            }
        }
        Ok(())
    }

    #[test]
    fn missing_in_loaded_modules() -> anyhow::Result<()> {
        let path = "cell1//next/package:import.bzl".to_owned();
        let resolver = resolver();
        let id = resolver.resolve_load(&path, None)?;

        let mut loaded_modules = loaded_modules();
        loaded_modules.map.remove(&id);
        let loader = InterpreterFileLoader::new(loaded_modules, resolver);
        match loader.load(&path) {
            Ok(_) => panic!("Expected load failure for {}", path),
            Err(_) => {
                // TODO: verify the error is correct
            }
        }
        Ok(())
    }

    #[test]
    fn valid_load() -> anyhow::Result<()> {
        let path = "cell1//next/package:import.bzl".to_owned();
        let resolver = resolver();
        let id = resolver.resolve_load(&path, None)?.to_string();

        let loader = InterpreterFileLoader::new(loaded_modules(), resolver);
        let loaded = loader.load(&path)?;

        let v = loaded.get("name").unwrap();
        assert_eq!(v.value().unpack_str(), Some(id.as_str()));

        Ok(())
    }

    #[test]
    fn valid_find() -> anyhow::Result<()> {
        let path = "cell1//next/package:import.bzl".to_owned();
        let resolver = resolver();
        let resolved = resolver.resolve_load(&path, None)?;
        let borrow = resolved.borrow();

        let loader = InterpreterFileLoader::new(loaded_modules(), resolver);
        let found = loader.find_module(borrow)?;

        let v = found.get("name").unwrap();
        assert_eq!(v.value().unpack_str(), Some(borrow.to_string().as_str()));

        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 dupe::Dupe;
use starlark::syntax::Dialect;
use starlark::syntax::DialectTypes;

#[derive(Copy, Clone, Dupe, Debug, Eq, PartialEq, Hash)]
pub enum StarlarkFileType {
    Bzl,
    Bxl,
    Buck,
    Package,
}

/// What type of file are we parsing - a `.bzl` file, `.bxl` file, or a `BUCK`/`TARGETS` file.
impl StarlarkFileType {
    pub fn dialect(&self, disable_starlark_types: bool) -> Dialect {
        let buck_dialect: Dialect = Dialect {
            enable_def: false,
            enable_lambda: true,
            enable_load: true,
            enable_keyword_only_arguments: false,
            enable_types: DialectTypes::Disable,
            enable_load_reexport: false,
            enable_top_level_stmt: false,
            ..Dialect::Standard
        };
        let package_dialect: Dialect = Dialect {
            enable_def: false,
            enable_lambda: true,
            enable_load: true,
            enable_keyword_only_arguments: false,
            enable_types: DialectTypes::Disable,
            enable_load_reexport: false,
            enable_top_level_stmt: false,
            ..Dialect::Standard
        };
        let bzl_dialect: Dialect = Dialect {
            enable_def: true,
            enable_lambda: true,
            enable_load: true,
            enable_keyword_only_arguments: true,
            enable_types: if disable_starlark_types {
                DialectTypes::ParseOnly
            } else {
                DialectTypes::Enable
            },
            enable_load_reexport: false,
            enable_top_level_stmt: true,
            ..Dialect::Standard
        };
        let bxl_dialect: Dialect = Dialect {
            enable_def: true,
            enable_lambda: true,
            enable_load: true,
            enable_keyword_only_arguments: true,
            enable_types: if disable_starlark_types {
                DialectTypes::ParseOnly
            } else {
                DialectTypes::Enable
            },
            enable_load_reexport: false,
            enable_top_level_stmt: true,
            ..Dialect::Standard
        };

        match self {
            Self::Bzl => bzl_dialect,
            Self::Buck => buck_dialect,
            Self::Package => package_dialect,
            Self::Bxl => bxl_dialect,
        }
    }
}


/*
 * 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 allocative::Allocative;
use async_trait::async_trait;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::legacy_configs::dice::HasLegacyConfigs;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_common::result::SharedResult;
use buck2_common::result::ToUnsharedResultExt;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::paths::CellRelativePathBuf;
use buck2_core::cells::CellAliasResolver;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;

use crate::package_imports::PackageImplicitImports;

#[derive(PartialEq, Allocative)]
pub struct ImplicitImportPaths {
    pub root_import: Option<ImportPath>,
    pub package_imports: PackageImplicitImports,
}

impl ImplicitImportPaths {
    pub fn parse(
        config: &LegacyBuckConfig,
        cell_name: BuildFileCell,
        cell_alias_resolver: &CellAliasResolver,
    ) -> anyhow::Result<ImplicitImportPaths> {
        // Oddly, the root import is defined to use a more path-like representation than
        // normal imports. e.g. it uses `cell//path/to/file.bzl` instead of
        // `cell//path/to:file.bzl`.
        let root_import = config
            .get("buildfile", "includes")
            .map(|i| {
                let (cell_alias, path): (&str, &str) = i.split_once("//").unwrap_or(("", i));
                let path = CellRelativePathBuf::try_from(path.to_owned())?;
                let path = CellPath::new(cell_alias_resolver.resolve(cell_alias)?, path.to_buf());

                // root imports are only going to be used by a top-level module in the cell they
                // are defined, so we can set the build_file_cell early.
                ImportPath::new_with_build_file_cells(path, cell_name)
            })
            .map_or(Ok(None), |e: anyhow::Result<ImportPath>| e.map(Some))?;
        let package_imports = PackageImplicitImports::new(
            cell_name,
            cell_alias_resolver.dupe(),
            config.get("buildfile", "package_includes"),
        )?;
        Ok(ImplicitImportPaths {
            root_import,
            package_imports,
        })
    }

    pub fn root_import(&self) -> Option<&ImportPath> {
        self.root_import.as_ref()
    }
}

#[async_trait]
pub trait HasImportPaths {
    async fn import_paths_for_cell(
        &self,
        cell_name: BuildFileCell,
    ) -> anyhow::Result<Arc<ImplicitImportPaths>>;
}

#[async_trait]
impl HasImportPaths for DiceComputations {
    async fn import_paths_for_cell(
        &self,
        cell_name: BuildFileCell,
    ) -> anyhow::Result<Arc<ImplicitImportPaths>> {
        #[derive(Debug, Eq, PartialEq, Hash, Clone, derive_more::Display, Allocative)]
        #[display(fmt = "{}", cell_name)]
        struct ImportPathsKey {
            cell_name: BuildFileCell,
        }

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

            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellation: &CancellationContext,
            ) -> Self::Value {
                let config = ctx
                    .get_legacy_config_for_cell(self.cell_name.name())
                    .await?;
                let cell_resolver = ctx.get_cell_resolver().await?;
                let cell_alias_resolver = cell_resolver
                    .get(self.cell_name.name())?
                    .cell_alias_resolver();
                Ok(Arc::new(ImplicitImportPaths::parse(
                    &config,
                    self.cell_name,
                    cell_alias_resolver,
                )?))
            }

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

        self.compute(&ImportPathsKey { cell_name })
            .await?
            .unshared_error()
    }
}


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

//! Implements Buck's handling of target patterns and parsing of build files.

#![feature(pattern)]
#![feature(try_blocks)]
#![feature(never_type)]
#![feature(box_patterns)]
// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

#[macro_use]
extern crate starlark;

pub mod build_context;
pub mod bxl;
pub mod coerce;
pub mod dice;
pub mod extra;
pub mod factory;
pub mod file_loader;
pub mod file_type;
pub mod functions;
pub mod import_paths;
pub mod load_module;
pub mod package_imports;
pub mod parse_import;
pub mod path;
pub mod prelude_path;
pub mod print_handler;
pub mod starlark_debug;
pub mod starlark_profiler;
pub mod starlark_promise;
pub mod types;


/*
 * 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 async_trait::async_trait;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::package::PackageLabel;
use buck2_util::late_binding::LateBinding;
use dice::DiceComputations;
use starlark::environment::Globals;

use crate::file_loader::LoadedModule;
use crate::file_loader::ModuleDeps;
use crate::file_type::StarlarkFileType;
use crate::path::PackageFilePath;
use crate::path::StarlarkModulePath;
use crate::prelude_path::PreludePath;

#[async_trait]
pub trait InterpreterCalculationImpl: Send + Sync + 'static {
    async fn get_loaded_module(
        &self,
        ctx: &DiceComputations,
        path: StarlarkModulePath<'_>,
    ) -> anyhow::Result<LoadedModule>;

    async fn get_module_deps(
        &self,
        ctx: &DiceComputations,
        package: PackageLabel,
        build_file_cell: BuildFileCell,
    ) -> anyhow::Result<ModuleDeps>;

    /// Return `None` if the PACKAGE file doesn't exist.
    async fn get_package_file_deps(
        &self,
        ctx: &DiceComputations,
        package: &PackageFilePath,
    ) -> anyhow::Result<Option<Vec<ImportPath>>>;

    async fn global_env_for_file_type(
        &self,
        ctx: &DiceComputations,
        file_type: StarlarkFileType,
    ) -> anyhow::Result<Globals>;

    async fn prelude_import(&self, ctx: &DiceComputations) -> anyhow::Result<Option<PreludePath>>;
}

pub static INTERPRETER_CALCULATION_IMPL: LateBinding<&'static dyn InterpreterCalculationImpl> =
    LateBinding::new("INTERPRETER_CALCULATION_IMPL");

#[async_trait]
pub trait InterpreterCalculation {
    /// Returns the LoadedModule for a given starlark file. This is cached on the dice graph.
    async fn get_loaded_module(&self, path: StarlarkModulePath<'_>)
    -> anyhow::Result<LoadedModule>;

    async fn get_loaded_module_from_import_path(
        &self,
        path: &ImportPath,
    ) -> anyhow::Result<LoadedModule> {
        self.get_loaded_module(StarlarkModulePath::LoadFile(path))
            .await
    }
}

#[async_trait]
impl InterpreterCalculation for DiceComputations {
    async fn get_loaded_module(
        &self,
        path: StarlarkModulePath<'_>,
    ) -> anyhow::Result<LoadedModule> {
        INTERPRETER_CALCULATION_IMPL
            .get()?
            .get_loaded_module(self, path)
            .await
    }
}