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

use std::time::SystemTime;

use anyhow::Context as _;
use buck2_core::cells::name::CellName;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use gazebo::prelude::*;

use super::LocalResourceType;
use super::PrepareForLocalExecutionResult;
use super::RequiredLocalResources;
use crate::convert;
use crate::data::ArgHandle;
use crate::data::ArgValue;
use crate::data::ArgValueContent;
use crate::data::ConfiguredTarget;
use crate::data::ConfiguredTargetHandle;
use crate::data::DeclaredOutput;
use crate::data::DisplayMetadata;
use crate::data::EnvHandle;
use crate::data::ExecuteRequest2;
use crate::data::ExecutionResult2;
use crate::data::ExecutionStatus;
use crate::data::ExecutionStream;
use crate::data::ExecutorConfigOverride;
use crate::data::ExternalRunnerSpec;
use crate::data::ExternalRunnerSpecValue;
use crate::data::Output;
use crate::data::TestExecutable;
use crate::data::TestResult;
use crate::data::TestStatus;
use crate::protocol::convert::host_sharing_requirements_from_grpc;
use crate::protocol::convert::host_sharing_requirements_to_grpc;

impl TryFrom<buck2_test_proto::DisplayMetadata> for DisplayMetadata {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::DisplayMetadata) -> Result<Self, Self::Error> {
        use buck2_test_proto::display_metadata::*;
        use buck2_test_proto::Testing;

        let res = match s.item.context("Missing `item`")? {
            Item::Listing(Listing { suite }) => Self::Listing(suite),
            Item::Testing(Testing { suite, testcases }) => Self::Testing { suite, testcases },
        };

        Ok(res)
    }
}

impl TryInto<buck2_test_proto::DisplayMetadata> for DisplayMetadata {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::DisplayMetadata, Self::Error> {
        use buck2_test_proto::display_metadata::*;
        use buck2_test_proto::Testing;

        let item = match self {
            Self::Listing(suite) => Item::Listing(Listing { suite }),
            Self::Testing { suite, testcases } => Item::Testing(Testing { suite, testcases }),
        };

        Ok(buck2_test_proto::DisplayMetadata { item: Some(item) })
    }
}

impl TryFrom<buck2_test_proto::ExecutionStream> for ExecutionStream {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExecutionStream) -> Result<Self, Self::Error> {
        use buck2_test_proto::execution_stream::*;

        Ok(match s.item.context("Missing `item`")? {
            Item::Inline(bytes) => Self::Inline(bytes),
        })
    }
}

impl TryInto<buck2_test_proto::ExecutionStream> for ExecutionStream {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExecutionStream, Self::Error> {
        use buck2_test_proto::execution_stream::*;

        let item = match self {
            Self::Inline(bytes) => Item::Inline(bytes),
        };

        Ok(buck2_test_proto::ExecutionStream { item: Some(item) })
    }
}

impl TryFrom<buck2_test_proto::ExecutionStatus> for ExecutionStatus {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExecutionStatus) -> Result<Self, Self::Error> {
        use buck2_test_proto::execution_status::*;

        Ok(match s.status.context("Missing `status`")? {
            Status::Finished(exitcode) => Self::Finished { exitcode },
            Status::TimedOut(duration) => Self::TimedOut {
                duration: convert::to_std_duration(duration)?,
            },
        })
    }
}

impl TryInto<buck2_test_proto::ExecutionStatus> for ExecutionStatus {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExecutionStatus, Self::Error> {
        use buck2_test_proto::execution_status::*;

        let status = match self {
            Self::Finished { exitcode } => Status::Finished(exitcode),
            Self::TimedOut { duration } => Status::TimedOut(duration.try_into()?),
        };

        Ok(buck2_test_proto::ExecutionStatus {
            status: Some(status),
        })
    }
}

impl TryFrom<buck2_test_proto::ConfiguredTargetHandle> for ConfiguredTargetHandle {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ConfiguredTargetHandle) -> Result<Self, Self::Error> {
        let handle = s.id.try_into().context("Invalid `id`")?;
        Ok(Self(handle))
    }
}

impl TryInto<buck2_test_proto::ConfiguredTargetHandle> for ConfiguredTargetHandle {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ConfiguredTargetHandle, Self::Error> {
        Ok(buck2_test_proto::ConfiguredTargetHandle {
            id: self.0.try_into().context("Invalid `handle`")?,
        })
    }
}

impl TryFrom<buck2_test_proto::ConfiguredTarget> for ConfiguredTarget {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ConfiguredTarget) -> Result<Self, Self::Error> {
        let buck2_test_proto::ConfiguredTarget {
            handle,
            cell,
            package,
            target,
            configuration,
            package_project_relative_path,
        } = s;

        Ok(Self {
            handle: handle
                .context("Missing `handle`")?
                .try_into()
                .context("Invalid `handle`")?,
            cell,
            package,
            target,
            configuration,
            package_project_relative_path: ForwardRelativePathBuf::try_from(
                package_project_relative_path,
            )?,
        })
    }
}

impl TryInto<buck2_test_proto::ConfiguredTarget> for ConfiguredTarget {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ConfiguredTarget, Self::Error> {
        Ok(buck2_test_proto::ConfiguredTarget {
            handle: Some(self.handle.try_into().context("Invalid `handle`")?),
            cell: self.cell,
            package: self.package,
            target: self.target,
            configuration: self.configuration,
            package_project_relative_path: self.package_project_relative_path.as_str().to_owned(),
        })
    }
}

impl TryFrom<i32> for TestStatus {
    type Error = anyhow::Error;

    fn try_from(s: i32) -> Result<Self, Self::Error> {
        let s = buck2_test_proto::TestStatus::from_i32(s).context("Invalid `status`")?;

        Ok(match s {
            buck2_test_proto::TestStatus::NotSet => {
                return Err(anyhow::Error::msg("Missing `status`"));
            }
            buck2_test_proto::TestStatus::Pass => TestStatus::PASS,
            buck2_test_proto::TestStatus::Fail => TestStatus::FAIL,
            buck2_test_proto::TestStatus::Skip => TestStatus::SKIP,
            buck2_test_proto::TestStatus::Omitted => TestStatus::OMITTED,
            buck2_test_proto::TestStatus::Fatal => TestStatus::FATAL,
            buck2_test_proto::TestStatus::Timeout => TestStatus::TIMEOUT,
            buck2_test_proto::TestStatus::Unknown => TestStatus::UNKNOWN,
            buck2_test_proto::TestStatus::Rerun => TestStatus::RERUN,
            buck2_test_proto::TestStatus::ListingSuccess => TestStatus::LISTING_SUCCESS,
            buck2_test_proto::TestStatus::ListingFailed => TestStatus::LISTING_FAILED,
        })
    }
}

impl TryInto<i32> for TestStatus {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<i32, Self::Error> {
        Ok(match self {
            TestStatus::PASS => buck2_test_proto::TestStatus::Pass,
            TestStatus::FAIL => buck2_test_proto::TestStatus::Fail,
            TestStatus::SKIP => buck2_test_proto::TestStatus::Skip,
            TestStatus::OMITTED => buck2_test_proto::TestStatus::Omitted,
            TestStatus::FATAL => buck2_test_proto::TestStatus::Fatal,
            TestStatus::TIMEOUT => buck2_test_proto::TestStatus::Timeout,
            TestStatus::UNKNOWN => buck2_test_proto::TestStatus::Unknown,
            TestStatus::RERUN => buck2_test_proto::TestStatus::Rerun,
            TestStatus::LISTING_SUCCESS => buck2_test_proto::TestStatus::ListingSuccess,
            TestStatus::LISTING_FAILED => buck2_test_proto::TestStatus::ListingFailed,
        } as i32)
    }
}

impl TryFrom<buck2_test_proto::TestResult> for TestResult {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::TestResult) -> Result<Self, Self::Error> {
        let buck2_test_proto::TestResult {
            target,
            name,
            status,
            msg,
            duration,
            details,
        } = s;

        let duration = duration
            .map(convert::to_std_duration)
            .transpose()
            .context("For `duration`")?;

        Ok(Self {
            target: target
                .context("Missing `target`")?
                .try_into()
                .context("Invalid `target`")?,
            name,
            status: status.try_into().context("Invalid `status`")?,
            msg: msg.map(|m| m.msg),
            duration,
            details,
        })
    }
}

impl TryInto<buck2_test_proto::TestResult> for TestResult {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::TestResult, Self::Error> {
        use buck2_test_proto::test_result::*;

        Ok(buck2_test_proto::TestResult {
            target: Some(self.target.try_into().context("Invalid `target`")?),
            name: self.name,
            status: self.status.try_into().context("Invalid `status`")?,
            details: self.details,
            msg: self.msg.map(|msg| OptionalMsg { msg }),
            duration: self.duration.try_map(|d| d.try_into())?,
        })
    }
}

impl TryFrom<buck2_test_proto::ExternalRunnerSpec> for ExternalRunnerSpec {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExternalRunnerSpec) -> Result<Self, Self::Error> {
        let buck2_test_proto::ExternalRunnerSpec {
            target,
            test_type,
            command,
            env,
            labels,
            contacts,
            oncall,
            working_dir_cell,
        } = s;

        Ok(Self {
            target: target
                .context("Missing `target`")?
                .try_into()
                .context("Invalid `target`")?,
            test_type,
            command: command
                .into_try_map(|x| x.try_into())
                .context("Invalid `command`")?,
            env: env
                .into_iter()
                .map(|(k, v)| Ok((k, v.try_into().unwrap())))
                .collect::<Result<_, Self::Error>>()?,
            labels,
            contacts,
            oncall,
            working_dir_cell: CellName::unchecked_new(&working_dir_cell)?,
        })
    }
}

impl TryInto<buck2_test_proto::ExternalRunnerSpec> for ExternalRunnerSpec {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExternalRunnerSpec, Self::Error> {
        let ExternalRunnerSpec {
            target,
            test_type,
            command,
            env,
            labels,
            contacts,
            oncall,
            working_dir_cell,
        } = self;
        Ok(buck2_test_proto::ExternalRunnerSpec {
            target: Some(target.try_into().context("Invalid `target`")?),
            test_type,
            command: command
                .into_try_map(|x| x.try_into())
                .context("Invalid `command`")?,
            env: env
                .into_iter()
                .map(|(k, v)| Ok((k, v.try_into().unwrap())))
                .collect::<Result<_, Self::Error>>()?,
            labels,
            contacts,
            oncall,
            working_dir_cell: working_dir_cell.as_str().to_owned(),
        })
    }
}

impl TryFrom<buck2_test_proto::ExternalRunnerSpecValue> for ExternalRunnerSpecValue {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExternalRunnerSpecValue) -> Result<Self, Self::Error> {
        use buck2_test_proto::external_runner_spec_value::*;
        Ok(match s.value.context("Missing `value`")? {
            Value::Verbatim(val) => ExternalRunnerSpecValue::Verbatim(val),
            Value::ArgHandle(val) => {
                ExternalRunnerSpecValue::ArgHandle(val.try_into().context("Invalid `arg_handle`")?)
            }
            Value::EnvHandle(val) => ExternalRunnerSpecValue::EnvHandle(val.into()),
        })
    }
}

impl TryInto<buck2_test_proto::ExternalRunnerSpecValue> for ExternalRunnerSpecValue {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExternalRunnerSpecValue, Self::Error> {
        use buck2_test_proto::external_runner_spec_value::*;

        let value = match self {
            Self::Verbatim(val) => Value::Verbatim(val),
            Self::ArgHandle(ArgHandle(val)) => {
                Value::ArgHandle(val.try_into().context("Invalid `arg_handle`")?)
            }
            Self::EnvHandle(EnvHandle(val)) => Value::EnvHandle(val),
        };

        Ok(buck2_test_proto::ExternalRunnerSpecValue { value: Some(value) })
    }
}

impl From<DeclaredOutput> for buck2_test_proto::DeclaredOutput {
    fn from(o: DeclaredOutput) -> Self {
        Self {
            name: o.name.as_str().to_owned(),
        }
    }
}

impl TryFrom<buck2_test_proto::DeclaredOutput> for DeclaredOutput {
    type Error = anyhow::Error;

    fn try_from(o: buck2_test_proto::DeclaredOutput) -> Result<Self, Self::Error> {
        let name = ForwardRelativePathBuf::try_from(o.name)?;
        Ok(Self { name })
    }
}

impl From<ExecutorConfigOverride> for buck2_test_proto::ExecutorConfigOverride {
    fn from(o: ExecutorConfigOverride) -> Self {
        Self {
            name: o.name.as_str().to_owned(),
        }
    }
}

impl From<buck2_test_proto::ExecutorConfigOverride> for ExecutorConfigOverride {
    fn from(o: buck2_test_proto::ExecutorConfigOverride) -> Self {
        Self { name: o.name }
    }
}

impl From<LocalResourceType> for buck2_test_proto::LocalResourceType {
    fn from(r: LocalResourceType) -> Self {
        Self {
            name: r.name.as_str().to_owned(),
        }
    }
}

impl From<buck2_test_proto::LocalResourceType> for LocalResourceType {
    fn from(o: buck2_test_proto::LocalResourceType) -> Self {
        Self { name: o.name }
    }
}

impl TryInto<buck2_test_proto::ArgValue> for ArgValue {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ArgValue, Self::Error> {
        Ok(buck2_test_proto::ArgValue {
            content: Some(self.content.try_into().context("Invalid `content`")?),
            format: self
                .format
                .map(|f| buck2_test_proto::ArgFormat { format: f }),
        })
    }
}

impl TryFrom<buck2_test_proto::ArgValue> for ArgValue {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ArgValue) -> Result<Self, Self::Error> {
        let content = s
            .content
            .context("Missing `content`")?
            .try_into()
            .context("Invalid `content`")?;
        let format = s.format.map(|f| f.format);

        Ok(Self { content, format })
    }
}

impl TryInto<buck2_test_proto::ArgValueContent> for ArgValueContent {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ArgValueContent, Self::Error> {
        use buck2_test_proto::arg_value_content::*;

        let value = match self {
            Self::ExternalRunnerSpecValue(value) => Value::SpecValue(
                value
                    .try_into()
                    .context("Invalid external runner spec value")?,
            ),
            Self::DeclaredOutput(value) => Value::DeclaredOutput(value.into()),
        };

        Ok(buck2_test_proto::ArgValueContent { value: Some(value) })
    }
}

impl TryFrom<buck2_test_proto::ArgValueContent> for ArgValueContent {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ArgValueContent) -> Result<Self, Self::Error> {
        use buck2_test_proto::arg_value_content::*;

        Ok(match s.value.context("Missing `value`")? {
            Value::SpecValue(value) => Self::ExternalRunnerSpecValue(
                value
                    .try_into()
                    .context("Invalid external runner spec value")?,
            ),
            Value::DeclaredOutput(value) => {
                Self::DeclaredOutput(value.try_into().context("Invalid `value`")?)
            }
        })
    }
}

impl TryFrom<buck2_test_proto::ExecuteRequest2> for ExecuteRequest2 {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExecuteRequest2) -> Result<Self, Self::Error> {
        let buck2_test_proto::ExecuteRequest2 {
            test_executable,
            timeout,
            host_sharing_requirements,
            executor_override,
            required_local_resources,
        } = s;

        let test_executable = test_executable
            .context("Missing `test_executable`")?
            .try_into()
            .context("Invalid `test_executable`")?;

        let timeout = convert::to_std_duration(timeout.context("Missing `timeout`")?)
            .context("Invalid `timeout`")?;

        let host_sharing_requirements =
            host_sharing_requirements.context("Missing `host_sharing_requirements`")?;
        let host_sharing_requirements =
            host_sharing_requirements_from_grpc(host_sharing_requirements)
                .context("Invalid `host_sharing_requirements`")?;

        let executor_override = executor_override.map(|o| o.into());

        let required_local_resources = RequiredLocalResources {
            resources: required_local_resources.into_map(|r| r.into()),
        };

        Ok(ExecuteRequest2 {
            test_executable,
            timeout,
            host_sharing_requirements,
            executor_override,
            required_local_resources,
        })
    }
}

impl TryInto<buck2_test_proto::ExecuteRequest2> for ExecuteRequest2 {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExecuteRequest2, Self::Error> {
        let test_executable = Some(
            self.test_executable
                .try_into()
                .context("Invalid `test_executable`")?,
        );

        Ok(buck2_test_proto::ExecuteRequest2 {
            test_executable,
            timeout: Some(self.timeout.try_into()?),
            host_sharing_requirements: Some(
                host_sharing_requirements_to_grpc(self.host_sharing_requirements)
                    .context("Invalid `host_sharing_requirements`")?,
            ),
            executor_override: self.executor_override.map(|o| o.into()),
            required_local_resources: self
                .required_local_resources
                .resources
                .into_map(|r| r.into()),
        })
    }
}

impl TryInto<buck2_test_proto::Output> for Output {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::Output, Self::Error> {
        use buck2_test_proto::output::*;

        let value = match self {
            Self::LocalPath(value) => {
                Value::LocalPath(value.to_str().context("Invalid local path")?.to_owned())
            }
        };

        Ok(buck2_test_proto::Output { value: Some(value) })
    }
}

impl TryFrom<buck2_test_proto::Output> for Output {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::Output) -> Result<Self, Self::Error> {
        use buck2_test_proto::output::*;

        Ok(match s.value.context("Missing `value`")? {
            Value::LocalPath(value) => {
                Self::LocalPath(value.try_into().context("Invalid local path value.")?)
            }
        })
    }
}

impl TryInto<buck2_test_proto::ExecutionResult2> for ExecutionResult2 {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::ExecutionResult2, Self::Error> {
        Ok(buck2_test_proto::ExecutionResult2 {
            status: Some(self.status.try_into().context("Invalid `status`")?),
            stdout: Some(self.stdout.try_into().context("Invalid `stdout`")?),
            stderr: Some(self.stderr.try_into().context("Invalid `stderr`")?),
            outputs: self
                .outputs
                .into_iter()
                .map(|(k, v)| {
                    Ok(buck2_test_proto::OutputEntry {
                        declared_output: Some(k.try_into().context("Invalid `declared_output`")?),
                        output: Some(v.try_into().context("Invalid `output`")?),
                    })
                })
                .collect::<Result<_, Self::Error>>()?,
            start_time: Some(
                self.start_time
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or_default()
                    .try_into()?,
            ),
            execution_time: Some(self.execution_time.try_into()?),
        })
    }
}

impl TryFrom<buck2_test_proto::ExecutionResult2> for ExecutionResult2 {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::ExecutionResult2) -> Result<Self, Self::Error> {
        let buck2_test_proto::ExecutionResult2 {
            status,
            stdout,
            stderr,
            outputs,
            start_time,
            execution_time,
        } = s;
        let status = status
            .context("Missing `status`")?
            .try_into()
            .context("Invalid `status`")?;
        let stdout = stdout
            .context("Missing `stdout`")?
            .try_into()
            .context("Invalid `stdout`")?;
        let stderr = stderr
            .context("Missing `stderr`")?
            .try_into()
            .context("Invalid `stderr`")?;

        let outputs = outputs
            .into_iter()
            .map(|entry| {
                let buck2_test_proto::OutputEntry {
                    declared_output,
                    output,
                } = entry;
                let declared_output = declared_output
                    .context("Missing `declared_output`")?
                    .try_into()
                    .context("Invalid `declared_output`")?;
                let output = output
                    .context("Missing `output`")?
                    .try_into()
                    .context("Invalid `output`")?;
                Ok((declared_output, output))
            })
            .collect::<Result<_, Self::Error>>()?;

        let start_time = SystemTime::UNIX_EPOCH
            + convert::to_std_duration(start_time.context("Missing `start_time`")?)
                .context("Invalid `start_time`")?;

        let execution_time =
            convert::to_std_duration(execution_time.context("Missing `execution_time`")?)
                .context("Invalid `execution_time`")?;

        Ok(ExecutionResult2 {
            status,
            stdout,
            stderr,
            outputs,
            start_time,
            execution_time,
        })
    }
}

impl TryFrom<buck2_test_proto::TestExecutable> for TestExecutable {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::TestExecutable) -> Result<Self, Self::Error> {
        let buck2_test_proto::TestExecutable {
            ui_prints,
            target,
            cmd,
            pre_create_dirs,
            env,
        } = s;
        let ui_prints = ui_prints
            .context("Missing `ui_prints`")?
            .try_into()
            .context("Invalid `ui_prints`")?;

        let target = target
            .context("Missing `target`")?
            .try_into()
            .context("Invalid `target`")?;

        let cmd = cmd
            .into_try_map(|c| c.try_into())
            .context("Invalid `cmd`")?;

        let env = env
            .into_iter()
            .map(|env_var| {
                let buck2_test_proto::EnvironmentVariable { key, value } = env_var;
                value
                    .context("Missing `value`")?
                    .try_into()
                    .context("Invalid `env`")
                    .map(|v: ArgValue| (key, v))
            })
            .collect::<anyhow::Result<_>>()?;

        let pre_create_dirs = pre_create_dirs
            .into_try_map(|c| c.try_into())
            .context("Invalid `pre_create_dirs`")?;

        Ok(TestExecutable {
            ui_prints,
            target,
            cmd,
            env,
            pre_create_dirs,
        })
    }
}

impl TryInto<buck2_test_proto::TestExecutable> for TestExecutable {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::TestExecutable, Self::Error> {
        let ui_prints = Some(self.ui_prints.try_into().context("Invalid `ui_prints`")?);
        let target = Some(self.target.try_into().context("Invalid `target`")?);
        let cmd = self
            .cmd
            .into_try_map(|i| i.try_into())
            .context("Invalid `cmd`")?;

        let env = self
            .env
            .into_iter()
            .map(|(k, v)| {
                v.try_into()
                    .context("Invalid `env`")
                    .map(
                        |v: buck2_test_proto::ArgValue| buck2_test_proto::EnvironmentVariable {
                            key: k,
                            value: Some(v),
                        },
                    )
            })
            .collect::<anyhow::Result<_>>()?;

        let pre_create_dirs = self
            .pre_create_dirs
            .into_try_map(|i| i.try_into())
            .context("Invalid `pre_create_dirs`")?;

        Ok(buck2_test_proto::TestExecutable {
            ui_prints,
            target,
            cmd,
            pre_create_dirs,
            env,
        })
    }
}

impl TryInto<buck2_test_proto::PrepareForLocalExecutionResult> for PrepareForLocalExecutionResult {
    type Error = anyhow::Error;

    fn try_into(self) -> Result<buck2_test_proto::PrepareForLocalExecutionResult, Self::Error> {
        let cwd = self.cwd.to_str().context("Invalid cwd path")?.to_owned();

        Ok(buck2_test_proto::PrepareForLocalExecutionResult {
            cmd: self.cmd,
            cwd,
            env: self
                .env
                .into_iter()
                .map(|(key, value)| buck2_test_proto::VerbatimEnvironmentVariable { key, value })
                .collect(),
        })
    }
}

impl TryFrom<buck2_test_proto::PrepareForLocalExecutionResult> for PrepareForLocalExecutionResult {
    type Error = anyhow::Error;

    fn try_from(s: buck2_test_proto::PrepareForLocalExecutionResult) -> Result<Self, Self::Error> {
        let buck2_test_proto::PrepareForLocalExecutionResult { cmd, cwd, env } = s;
        let cwd = cwd.try_into().context("Invalid cwd value.")?;

        let env = env
            .into_iter()
            .map(|env_var| (env_var.key, env_var.value))
            .collect();

        Ok(PrepareForLocalExecutionResult { cmd, env, cwd })
    }
}

#[cfg(test)]
mod tests {
    use std::fmt::Debug;
    use std::time::Duration;

    use host_sharing::HostSharingRequirements;
    use sorted_vector_map::sorted_vector_map;

    use super::*;

    fn assert_roundtrips<P, S>(s: &S)
    where
        S: Clone + PartialEq + Debug,
        P: TryInto<S, Error = anyhow::Error>,
        S: TryInto<P, Error = anyhow::Error>,
    {
        let proto: P = s.clone().try_into().unwrap();
        let roundtrip: S = proto.try_into().unwrap();
        assert_eq!(*s, roundtrip);
    }

    #[test]
    fn external_runner_spec_roundtrip() {
        let test_spec = ExternalRunnerSpec {
            target: ConfiguredTarget {
                handle: ConfiguredTargetHandle(1),
                cell: "qux".into(),
                package: "foo".into(),
                target: "bar".into(),
                configuration: "xxx".into(),
                package_project_relative_path: ForwardRelativePathBuf::unchecked_new(
                    "qux/foo".to_owned(),
                ),
            },
            test_type: "some_type".to_owned(),
            command: vec![
                ExternalRunnerSpecValue::Verbatim("arg".to_owned()),
                ExternalRunnerSpecValue::ArgHandle(ArgHandle(42)),
            ],
            env: [
                (
                    "FOO".to_owned(),
                    ExternalRunnerSpecValue::EnvHandle(EnvHandle("FOO".to_owned())),
                ),
                (
                    "BAR".to_owned(),
                    ExternalRunnerSpecValue::Verbatim("BAR".to_owned()),
                ),
            ]
            .into_iter()
            .collect(),
            labels: vec!["label1".to_owned(), "label2".to_owned()],
            contacts: vec!["contact1".to_owned(), "contact2".to_owned()],
            oncall: Some("contact1".to_owned()),
            working_dir_cell: CellName::testing_new("qux"),
        };
        assert_roundtrips::<buck2_test_proto::ExternalRunnerSpec, ExternalRunnerSpec>(&test_spec);
    }

    #[test]
    fn execute_request2_roundtrip() {
        let declared_output = DeclaredOutput {
            name: ForwardRelativePathBuf::unchecked_new("name".to_owned()),
        };

        let test_executable = TestExecutable {
            ui_prints: DisplayMetadata::Listing("name".to_owned()),
            target: ConfiguredTargetHandle(42),
            cmd: vec![
                ArgValue {
                    content: ArgValueContent::ExternalRunnerSpecValue(
                        ExternalRunnerSpecValue::Verbatim("arg".to_owned()),
                    ),
                    format: None,
                },
                ArgValue {
                    content: ArgValueContent::DeclaredOutput(declared_output.clone()),
                    format: Some("--output={}".to_owned()),
                },
            ],
            env: [(
                "FOO".to_owned(),
                ArgValue {
                    content: ArgValueContent::ExternalRunnerSpecValue(
                        ExternalRunnerSpecValue::EnvHandle(EnvHandle("FOO".to_owned())),
                    ),
                    format: None,
                },
            )]
            .into_iter()
            .collect(),
            pre_create_dirs: vec![declared_output],
        };
        let request = ExecuteRequest2 {
            test_executable,
            timeout: Duration::from_millis(42),
            host_sharing_requirements: HostSharingRequirements::ExclusiveAccess,
            executor_override: Some(ExecutorConfigOverride {
                name: "foo".to_owned(),
            }),
            required_local_resources: RequiredLocalResources { resources: vec![] },
        };
        assert_roundtrips::<buck2_test_proto::ExecuteRequest2, ExecuteRequest2>(&request);
    }

    #[test]
    fn execution_result2_roundtrips() {
        let local_path = if cfg!(not(windows)) {
            "/some/path"
        } else {
            "c:/some/path"
        };

        let result = ExecutionResult2 {
            status: ExecutionStatus::Finished { exitcode: 42 },
            stdout: ExecutionStream::Inline(vec![97, 115, 109]),
            stderr: ExecutionStream::Inline(vec![118, 105, 109]),
            outputs: [(
                DeclaredOutput {
                    name: ForwardRelativePathBuf::unchecked_new("name".to_owned()),
                },
                Output::LocalPath(String::from(local_path).try_into().expect("valid abs path")),
            )]
            .into_iter()
            .collect(),
            start_time: SystemTime::UNIX_EPOCH + Duration::from_secs(123),
            execution_time: Duration::from_secs(456),
        };
        assert_roundtrips::<buck2_test_proto::ExecutionResult2, ExecutionResult2>(&result);
    }

    #[test]
    fn prepare_for_local_execution_result_roundtrip() {
        let cmd = vec![
            "my_cmd".to_owned(),
            "--some-arg".to_owned(),
            "some_value".to_owned(),
        ];
        let local_path = if cfg!(not(windows)) {
            "/some/path"
        } else {
            "c:/some/path"
        };
        let cwd = String::from(local_path).try_into().expect("valid abs path");
        let env = sorted_vector_map! { "some_env".to_owned() => "some_env_val".to_owned() };

        let result = PrepareForLocalExecutionResult { cmd, env, cwd };

        assert_roundtrips::<
            buck2_test_proto::PrepareForLocalExecutionResult,
            PrepareForLocalExecutionResult,
        >(&result);
    }

    #[test]
    fn test_executable_roundtrip() {
        let declared_output = DeclaredOutput {
            name: ForwardRelativePathBuf::unchecked_new("name".to_owned()),
        };

        let test_executable = TestExecutable {
            ui_prints: DisplayMetadata::Listing("name".to_owned()),
            target: ConfiguredTargetHandle(42),
            cmd: vec![
                ArgValue {
                    content: ArgValueContent::ExternalRunnerSpecValue(
                        ExternalRunnerSpecValue::Verbatim("arg".to_owned()),
                    ),
                    format: None,
                },
                ArgValue {
                    content: ArgValueContent::DeclaredOutput(declared_output.clone()),
                    format: Some("--output={}".to_owned()),
                },
            ],
            env: [(
                "FOO".to_owned(),
                ArgValue {
                    content: ArgValueContent::ExternalRunnerSpecValue(
                        ExternalRunnerSpecValue::EnvHandle(EnvHandle("FOO".to_owned())),
                    ),
                    format: None,
                },
            )]
            .into_iter()
            .collect(),
            pre_create_dirs: vec![declared_output],
        };

        assert_roundtrips::<buck2_test_proto::TestExecutable, TestExecutable>(&test_executable);
    }
}


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

//! Core data objects used in the protocol

mod convert;

use std::collections::HashMap;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::time::Duration;
use std::time::SystemTime;

use buck2_core::cells::name::CellName;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use derive_more::From;
use dupe::Dupe;
use host_sharing::HostSharingRequirements;
use sorted_vector_map::SortedVectorMap;

/// A handle generated by the TestOrchestrator. It can be used by the TestExecutor to access this
/// target.
#[derive(Debug, Copy, Clone, Dupe, From, Hash, Eq, PartialEq)]
pub struct ConfiguredTargetHandle(u64);

/// The Target of a test rule
#[derive(Debug, Clone, PartialEq)]
pub struct ConfiguredTarget {
    pub handle: ConfiguredTargetHandle,
    /// Structured data
    pub cell: String,
    pub package: String,
    pub target: String,
    pub configuration: String,
    pub package_project_relative_path: ForwardRelativePathBuf,
}

/// Metadata about the execution to display
#[derive(Debug, Clone, PartialEq)]
pub enum DisplayMetadata {
    // Listing the test binary to discover tests. The String is the name of the suite at the binary
    Listing(String),
    // the name of the test(s) that we are running for the suite of a target
    Testing {
        suite: String,
        testcases: Vec<String>,
    },
}

#[derive(Clone, PartialEq)]
pub enum ExecutionStream {
    Inline(Vec<u8>),
}

impl Debug for ExecutionStream {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::Inline(d) => {
                write!(f, "{}", String::from_utf8_lossy(d))
            }
        }
    }
}

#[derive(Clone, Debug, Dupe, PartialEq)]
pub enum ExecutionStatus {
    Finished { exitcode: i32 },
    TimedOut { duration: Duration },
}

/// The result of running a test
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct TestResult {
    // the target the test came from
    pub target: ConfiguredTargetHandle,
    // the name of the test
    pub name: String,
    // the status of running the test
    pub status: TestStatus,
    // additional optional messages
    pub msg: Option<String>,
    // the duration of the test run
    // TODO(skcd) should this be optional? why doesn't everything have duration
    pub duration: Option<Duration>,
    // the output of the test execution (combining stdout and stderr)
    pub details: String,
}

/// different possible test results
#[derive(PartialEq, Eq, Debug, Clone, Dupe)]
#[allow(non_camel_case_types)]
pub enum TestStatus {
    PASS,
    FAIL,
    SKIP,
    OMITTED,
    FATAL,
    TIMEOUT,
    // There is something called unknown, adding it here for now,
    // we can change it later on.
    UNKNOWN,
    // We also have re-runs
    RERUN,
    LISTING_SUCCESS,
    LISTING_FAILED,
}

/// The set of information about a test rule that is passed to the test executor
#[derive(Clone, Debug, PartialEq)]
pub struct ExternalRunnerSpec {
    /// Target the spec belongs to
    pub target: ConfiguredTarget,
    /// Type of test spec
    pub test_type: String,
    /// Base command used for further processing. A mix of verbatim arguments and
    /// opaque handles for more complex arguments.
    pub command: Vec<ExternalRunnerSpecValue>,
    /// Environment variables a specified by the rule. A mapping from keys to
    /// verbatim values or opaque handles for more complex values.
    pub env: HashMap<String, ExternalRunnerSpecValue>,
    /// Labels defined on the rule.
    pub labels: Vec<String>,
    /// Contacts defined on the rule.
    pub contacts: Vec<String>,
    /// Oncall for the test
    pub oncall: Option<String>,
    /// Cell of current working directory for test command.
    pub working_dir_cell: CellName,
}

/// Command line argument or environment variable value
///
/// It is either a verbatim string, or a reference to a more complex value that's opaque to the
/// test run coordinator.
#[derive(Clone, Debug, PartialEq)]
pub enum ExternalRunnerSpecValue {
    Verbatim(String),
    ArgHandle(ArgHandle),
    EnvHandle(EnvHandle),
}

/// Handle referring to a complex argument defined on the test rule
#[derive(Clone, Debug, Dupe, PartialEq, From)]
pub struct ArgHandle(pub usize);

impl TryFrom<i64> for ArgHandle {
    type Error = anyhow::Error;

    fn try_from(i: i64) -> Result<Self, Self::Error> {
        Ok(ArgHandle(i.try_into()?))
    }
}

/// Handle referring to a complex environment value defined on the test rule
#[derive(Clone, Debug, PartialEq, From)]
pub struct EnvHandle(pub String);

#[derive(Clone, Debug, PartialEq)]
pub struct ArgValue {
    pub content: ArgValueContent,
    pub format: Option<String>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum ArgValueContent {
    ExternalRunnerSpecValue(ExternalRunnerSpecValue),
    DeclaredOutput(DeclaredOutput),
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct DeclaredOutput {
    pub name: ForwardRelativePathBuf,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct ExecutorConfigOverride {
    pub name: String,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct LocalResourceType {
    pub name: String,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct RequiredLocalResources {
    pub resources: Vec<LocalResourceType>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct ExecuteRequest2 {
    pub test_executable: TestExecutable,
    pub timeout: Duration,
    pub host_sharing_requirements: HostSharingRequirements,
    pub executor_override: Option<ExecutorConfigOverride>,
    pub required_local_resources: RequiredLocalResources,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Output {
    LocalPath(AbsNormPathBuf),
}

#[derive(Clone, Debug, PartialEq)]
pub struct ExecutionResult2 {
    pub status: ExecutionStatus,
    pub stdout: ExecutionStream,
    pub stderr: ExecutionStream,
    pub outputs: HashMap<DeclaredOutput, Output>,
    pub start_time: SystemTime,
    pub execution_time: Duration,
}

pub enum ExecuteResponse {
    /// A result is available.
    Result(ExecutionResult2),

    /// The test run is being cancelled.
    Cancelled,
}

#[derive(Clone, Debug, PartialEq)]
pub struct TestExecutable {
    pub ui_prints: DisplayMetadata,
    pub target: ConfiguredTargetHandle,
    pub cmd: Vec<ArgValue>,
    pub env: SortedVectorMap<String, ArgValue>,
    pub pre_create_dirs: Vec<DeclaredOutput>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct PrepareForLocalExecutionResult {
    pub cmd: Vec<String>,
    pub env: SortedVectorMap<String, String>,
    pub cwd: AbsNormPathBuf,
}

pub mod testing {
    use crate::data::ConfiguredTargetHandle;

    pub trait ConfiguredTargetHandleExt {
        fn testing_new(id: u64) -> Self;
    }

    impl ConfiguredTargetHandleExt for ConfiguredTargetHandle {
        fn testing_new(id: u64) -> Self {
            Self(id)
        }
    }
}


/*
 * 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 as _;
use buck2_grpc::make_channel;
use buck2_grpc::spawn_oneshot;
use buck2_grpc::to_tonic;
use buck2_grpc::ServerHandle;
use buck2_test_proto::test_executor_client;
use buck2_test_proto::test_executor_server;
use buck2_test_proto::Empty;
use buck2_test_proto::ExternalRunnerSpecRequest;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tonic::transport::Channel;

use crate::data::ExternalRunnerSpec;
use crate::protocol::TestExecutor;

pub struct TestExecutorClient {
    client: test_executor_client::TestExecutorClient<Channel>,
}

impl TestExecutorClient {
    pub async fn new<T>(io: T) -> anyhow::Result<Self>
    where
        T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
    {
        let channel = make_channel(io, "executor").await?;

        Ok(Self {
            client: test_executor_client::TestExecutorClient::new(channel)
                .max_encoding_message_size(usize::MAX)
                .max_decoding_message_size(usize::MAX),
        })
    }
}

#[async_trait::async_trait]
impl TestExecutor for TestExecutorClient {
    async fn external_runner_spec(&self, s: ExternalRunnerSpec) -> anyhow::Result<()> {
        self.client
            .clone()
            .external_runner_spec(ExternalRunnerSpecRequest {
                test_spec: Some(s.try_into().context("Invalid `test_spec`")?),
            })
            .await?;

        Ok(())
    }

    async fn end_of_test_requests(&self) -> anyhow::Result<()> {
        self.client.clone().end_of_test_requests(Empty {}).await?;
        Ok(())
    }
}

pub struct Service<T> {
    inner: T,
}

#[async_trait::async_trait]
impl<T> test_executor_server::TestExecutor for Service<T>
where
    T: TestExecutor + Send + Sync + 'static,
{
    async fn external_runner_spec(
        &self,
        request: tonic::Request<ExternalRunnerSpecRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let ExternalRunnerSpecRequest { test_spec } = request.into_inner();

            let test_spec = test_spec
                .context("Missing `test_spec`")?
                .try_into()
                .context("Invalid `test_spec`")?;

            self.inner
                .external_runner_spec(test_spec)
                .await
                .context("Failed to dispatch test_spec")?;

            Ok(Empty {})
        })
        .await
    }

    async fn end_of_test_requests(
        &self,
        _: tonic::Request<Empty>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            self.inner
                .end_of_test_requests()
                .await
                .context("Failed to report end-of-tests")?;

            Ok(Empty {})
        })
        .await
    }
}

pub fn spawn_executor_server<I, E>(io: I, executor: E) -> ServerHandle
where
    I: AsyncRead + AsyncWrite + Send + Unpin + 'static + tonic::transport::server::Connected,
    E: TestExecutor + Send + Sync + 'static,
{
    let router = tonic::transport::Server::builder().add_service(
        test_executor_server::TestExecutorServer::new(Service { inner: executor })
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX),
    );

    spawn_oneshot(io, router)
}


/*
 * 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 executor;
mod orchestrator;

pub use executor::spawn_executor_server;
pub use executor::TestExecutorClient;
pub use orchestrator::spawn_orchestrator_server;
pub use orchestrator::TestOrchestratorClient;

#[cfg(test)]
mod test;


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

use std::collections::HashMap;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;

use anyhow::Context as _;
use buck2_downward_api::DownwardApi;
use buck2_downward_api_proto::downward_api_client;
use buck2_downward_api_proto::downward_api_server;
use buck2_downward_api_proto::ConsoleRequest;
use buck2_downward_api_proto::ExternalEventRequest;
use buck2_downward_api_proto::LogRequest;
use buck2_events::dispatch::with_dispatcher_async;
use buck2_events::dispatch::EventDispatcher;
use buck2_grpc::make_channel;
use buck2_grpc::spawn_oneshot;
use buck2_grpc::to_tonic;
use buck2_grpc::ServerHandle;
use buck2_test_proto::test_orchestrator_client;
use buck2_test_proto::test_orchestrator_server;
use buck2_test_proto::AttachInfoMessageRequest;
use buck2_test_proto::Empty;
use buck2_test_proto::EndOfTestResultsRequest;
use buck2_test_proto::ExecuteResponse2;
use buck2_test_proto::PrepareForLocalExecutionResponse;
use buck2_test_proto::ReportTestResultRequest;
use buck2_test_proto::ReportTestSessionRequest;
use buck2_test_proto::ReportTestsDiscoveredRequest;
use buck2_test_proto::Testing;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use host_sharing::HostSharingRequirements;
use sorted_vector_map::SortedVectorMap;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tonic::transport::Channel;
use tower_layer::Layer;
use tracing::Level;

use crate::data::ArgValue;
use crate::data::ConfiguredTargetHandle;
use crate::data::DeclaredOutput;
use crate::data::DisplayMetadata;
use crate::data::ExecuteRequest2;
use crate::data::ExecuteResponse;
use crate::data::ExecutorConfigOverride;
use crate::data::PrepareForLocalExecutionResult;
use crate::data::RequiredLocalResources;
use crate::data::TestExecutable;
use crate::data::TestResult;
use crate::protocol::TestOrchestrator;

/// Test runner client to buck2 test orchestrator.
///
/// When running `buck2 test`, buck2 starts a gRPC server
/// and spawns a test runner which connects to buck2 process using this client.
pub struct TestOrchestratorClient {
    test_orchestrator_client: test_orchestrator_client::TestOrchestratorClient<Channel>,
    downward_api_client: downward_api_client::DownwardApiClient<Channel>,
}

impl TestOrchestratorClient {
    pub async fn new<T>(io: T) -> anyhow::Result<Self>
    where
        T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
    {
        let channel = make_channel(io, "orchestrator").await?;

        Ok(Self {
            test_orchestrator_client: test_orchestrator_client::TestOrchestratorClient::new(
                channel.clone(),
            )
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX),
            downward_api_client: downward_api_client::DownwardApiClient::new(channel)
                .max_encoding_message_size(usize::MAX)
                .max_decoding_message_size(usize::MAX),
        })
    }
}

#[async_trait::async_trait]
impl DownwardApi for TestOrchestratorClient {
    async fn console(&self, level: Level, message: String) -> anyhow::Result<()> {
        let level = level.try_into().context("Invalid `level`")?;

        self.downward_api_client
            .clone()
            .console(ConsoleRequest {
                level: Some(level),
                message,
            })
            .await?;

        Ok(())
    }

    async fn log(&self, level: Level, message: String) -> anyhow::Result<()> {
        let level = level.try_into().context("Invalid `level`")?;

        self.downward_api_client
            .clone()
            .log(LogRequest {
                level: Some(level),
                message,
            })
            .await?;

        Ok(())
    }

    async fn external(&self, data: HashMap<String, String>) -> anyhow::Result<()> {
        let event = data.into();

        self.downward_api_client
            .clone()
            .external_event(ExternalEventRequest { event: Some(event) })
            .await?;

        Ok(())
    }
}

impl TestOrchestratorClient {
    pub async fn execute2(
        &self,
        ui_prints: DisplayMetadata,
        target: ConfiguredTargetHandle,
        cmd: Vec<ArgValue>,
        env: SortedVectorMap<String, ArgValue>,
        timeout: Duration,
        host_sharing_requirements: HostSharingRequirements,
        pre_create_dirs: Vec<DeclaredOutput>,
        executor_override: Option<ExecutorConfigOverride>,
        required_local_resources: RequiredLocalResources,
    ) -> anyhow::Result<ExecuteResponse> {
        let test_executable = TestExecutable {
            ui_prints,
            target,
            cmd,
            env,
            pre_create_dirs,
        };

        let req = ExecuteRequest2 {
            test_executable,
            timeout,
            host_sharing_requirements,
            executor_override,
            required_local_resources,
        };

        let req: buck2_test_proto::ExecuteRequest2 =
            req.try_into().context("Invalid execute request")?;

        let ExecuteResponse2 { response } = self
            .test_orchestrator_client
            .clone()
            .execute2(req)
            .await?
            .into_inner();

        let response = match response.context("Missing `response`")? {
            buck2_test_proto::execute_response2::Response::Result(res) => {
                ExecuteResponse::Result(res.try_into().context("Invalid `result`")?)
            }
            buck2_test_proto::execute_response2::Response::Cancelled(
                buck2_test_proto::Cancelled {},
            ) => ExecuteResponse::Cancelled,
        };

        Ok(response)
    }

    pub async fn report_test_result(&self, result: TestResult) -> anyhow::Result<()> {
        let result = result.try_into().context("Invalid `result`")?;

        self.test_orchestrator_client
            .clone()
            .report_test_result(ReportTestResultRequest {
                result: Some(result),
            })
            .await?;

        Ok(())
    }

    pub async fn report_tests_discovered(
        &self,
        target: ConfiguredTargetHandle,
        suite: String,
        tests: Vec<String>,
    ) -> anyhow::Result<()> {
        let target = target.try_into().context("Invalid `target`")?;

        self.test_orchestrator_client
            .clone()
            .report_tests_discovered(ReportTestsDiscoveredRequest {
                target: Some(target),
                testing: Some(Testing {
                    suite,
                    testcases: tests,
                }),
            })
            .await?;

        Ok(())
    }

    pub async fn report_test_session(&self, session_info: String) -> anyhow::Result<()> {
        self.test_orchestrator_client
            .clone()
            .report_test_session(ReportTestSessionRequest { session_info })
            .await?;

        Ok(())
    }

    pub async fn end_of_test_results(&self, exit_code: i32) -> anyhow::Result<()> {
        self.test_orchestrator_client
            .clone()
            .end_of_test_results(EndOfTestResultsRequest { exit_code })
            .await?;

        Ok(())
    }

    pub async fn prepare_for_local_execution(
        &self,
        ui_prints: DisplayMetadata,
        target: ConfiguredTargetHandle,
        cmd: Vec<ArgValue>,
        env: SortedVectorMap<String, ArgValue>,
        pre_create_dirs: Vec<DeclaredOutput>,
    ) -> anyhow::Result<PrepareForLocalExecutionResult> {
        let executable = TestExecutable {
            ui_prints,
            target,
            cmd,
            env,
            pre_create_dirs,
        };

        let executable: buck2_test_proto::TestExecutable = executable
            .try_into()
            .context("Invalid prepare_for_local_execution request")?;

        let request = buck2_test_proto::PrepareForLocalExecutionRequest {
            test_executable: Some(executable),
        };
        let PrepareForLocalExecutionResponse { result } = self
            .test_orchestrator_client
            .clone()
            .prepare_for_local_execution(request)
            .await?
            .into_inner();

        let result = result
            .context("Missing `result`")?
            .try_into()
            .context("Invalid `result`")?;

        Ok(result)
    }

    pub async fn attach_info_message(&self, message: String) -> anyhow::Result<()> {
        self.test_orchestrator_client
            .clone()
            .attach_info_message(AttachInfoMessageRequest { message })
            .await?;
        Ok(())
    }
}

struct TestOrchestratorService<T: TestOrchestrator> {
    inner: T,
}

#[async_trait::async_trait]
impl<T> test_orchestrator_server::TestOrchestrator for TestOrchestratorService<T>
where
    T: TestOrchestrator + Send + Sync + 'static,
{
    async fn execute2(
        &self,
        request: tonic::Request<buck2_test_proto::ExecuteRequest2>,
    ) -> Result<tonic::Response<ExecuteResponse2>, tonic::Status> {
        to_tonic(async move {
            let ExecuteRequest2 {
                test_executable,
                timeout,
                host_sharing_requirements,
                executor_override,
                required_local_resources,
            } = request
                .into_inner()
                .try_into()
                .context("Invalid execute2 request")?;

            let TestExecutable {
                ui_prints,
                target,
                cmd,
                env,
                pre_create_dirs,
            } = test_executable;

            let response = self
                .inner
                .execute2(
                    ui_prints,
                    target,
                    cmd,
                    env,
                    timeout,
                    host_sharing_requirements,
                    pre_create_dirs,
                    executor_override,
                    required_local_resources,
                )
                .await
                .context("Execution failed")?;

            let response = match response {
                ExecuteResponse::Result(r) => {
                    buck2_test_proto::execute_response2::Response::Result(
                        r.try_into().context("Failed to serialize result")?,
                    )
                }
                ExecuteResponse::Cancelled => {
                    buck2_test_proto::execute_response2::Response::Cancelled(
                        buck2_test_proto::Cancelled {},
                    )
                }
            };

            Ok(ExecuteResponse2 {
                response: Some(response),
            })
        })
        .await
    }

    async fn end_of_test_results(
        &self,
        req: tonic::Request<EndOfTestResultsRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let EndOfTestResultsRequest { exit_code } = req.into_inner();

            self.inner
                .end_of_test_results(exit_code)
                .await
                .context("Failed to report end-of-tests")?;

            Ok(Empty {})
        })
        .await
    }

    async fn report_test_result(
        &self,
        request: tonic::Request<ReportTestResultRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let ReportTestResultRequest { result } = request.into_inner();

            let result = result
                .context("Missing `result`")?
                .try_into()
                .context("Invalid `result`")?;

            self.inner
                .report_test_result(result)
                .await
                .context("Failed to report end-of-tests")?;

            Ok(Empty {})
        })
        .await
    }

    async fn report_tests_discovered(
        &self,
        request: tonic::Request<ReportTestsDiscoveredRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let ReportTestsDiscoveredRequest { target, testing } = request.into_inner();

            let target = target
                .context("Missing `target`")?
                .try_into()
                .context("Invalid `target`")?;

            let Testing { suite, testcases } = testing.context("Missing `testing`")?;

            self.inner
                .report_tests_discovered(target, suite, testcases)
                .await
                .context("Failed to report end-of-tests")?;

            Ok(Empty {})
        })
        .await
    }

    async fn report_test_session(
        &self,
        request: tonic::Request<ReportTestSessionRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let ReportTestSessionRequest { session_info } = request.into_inner();

            self.inner
                .report_test_session(session_info)
                .await
                .context("Failed to report test session summary")?;

            Ok(Empty {})
        })
        .await
    }

    async fn prepare_for_local_execution(
        &self,
        request: tonic::Request<buck2_test_proto::PrepareForLocalExecutionRequest>,
    ) -> Result<tonic::Response<PrepareForLocalExecutionResponse>, tonic::Status> {
        to_tonic(async move {
            let buck2_test_proto::PrepareForLocalExecutionRequest { test_executable } =
                request.into_inner();

            let TestExecutable {
                ui_prints,
                target,
                cmd,
                env,
                pre_create_dirs,
            } = test_executable
                .context("Missing `test_executable`")?
                .try_into()
                .context("Invalid `test_executable`")
                .context("Invalid prepare_for_local_execution request")?;

            let result = self
                .inner
                .prepare_for_local_execution(ui_prints, target, cmd, env, pre_create_dirs)
                .await
                .context("Prepare for local execution failed")?;

            let result = result.try_into().context("Failed to serialize result")?;

            Ok(PrepareForLocalExecutionResponse {
                result: Some(result),
            })
        })
        .await
    }

    async fn attach_info_message(
        &self,
        request: tonic::Request<AttachInfoMessageRequest>,
    ) -> Result<tonic::Response<Empty>, tonic::Status> {
        to_tonic(async move {
            let AttachInfoMessageRequest { message } = request.into_inner();

            self.inner
                .attach_info_message(message)
                .await
                .context("Failed to attach info messages")?;

            Ok(Empty {})
        })
        .await
    }
}

struct DownwardApiService<T: DownwardApi> {
    inner: T,
}

#[async_trait::async_trait]
impl<T> downward_api_server::DownwardApi for DownwardApiService<T>
where
    T: DownwardApi + Send + Sync + 'static,
{
    async fn console(
        &self,
        request: tonic::Request<ConsoleRequest>,
    ) -> Result<tonic::Response<buck2_downward_api_proto::Empty>, tonic::Status> {
        to_tonic(async move {
            let ConsoleRequest { level, message } = request.into_inner();

            let level = level
                .context("Missing `level`")?
                .try_into()
                .context("Invalid `level`")?;

            self.inner
                .console(level, message)
                .await
                .context("Failed to console")?;

            Ok(buck2_downward_api_proto::Empty {})
        })
        .await
    }

    async fn log(
        &self,
        request: tonic::Request<LogRequest>,
    ) -> Result<tonic::Response<buck2_downward_api_proto::Empty>, tonic::Status> {
        to_tonic(async move {
            let LogRequest { level, message } = request.into_inner();

            let level = level
                .context("Missing `level`")?
                .try_into()
                .context("Invalid `level`")?;

            self.inner
                .log(level, message)
                .await
                .context("Failed to log")?;

            Ok(buck2_downward_api_proto::Empty {})
        })
        .await
    }

    async fn external_event(
        &self,
        request: tonic::Request<ExternalEventRequest>,
    ) -> Result<tonic::Response<buck2_downward_api_proto::Empty>, tonic::Status> {
        to_tonic(async move {
            let ExternalEventRequest { event } = request.into_inner();

            let event = event
                .context("Missing `event`")?
                .try_into()
                .context("Invalid `event`")?;

            self.inner
                .external(event)
                .await
                .context("Failed to deliver event")?;

            Ok(buck2_downward_api_proto::Empty {})
        })
        .await
    }
}

pub fn spawn_orchestrator_server<I, O, D>(
    io: I,
    orchestrator: O,
    downward_api: D,
    dispatcher: EventDispatcher,
) -> ServerHandle
where
    I: AsyncRead + AsyncWrite + Send + Unpin + 'static + tonic::transport::server::Connected,
    O: TestOrchestrator + Send + Sync + 'static,
    D: DownwardApi + Send + Sync + 'static,
{
    let router = tonic::transport::Server::builder()
        .layer(EventDispatcherLayer::new(dispatcher))
        .add_service(
            test_orchestrator_server::TestOrchestratorServer::new(TestOrchestratorService {
                inner: orchestrator,
            })
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX),
        )
        .add_service(
            downward_api_server::DownwardApiServer::new(DownwardApiService {
                inner: downward_api,
            })
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX),
        );

    spawn_oneshot(io, router)
}

/// Used to wrap the Tonic server so that the spawned server has its EventDispatcher set.
#[derive(Clone, Dupe)]
pub struct EventDispatcherLayer {
    dispatcher: EventDispatcher,
}

impl EventDispatcherLayer {
    pub fn new(dispatcher: EventDispatcher) -> Self {
        EventDispatcherLayer { dispatcher }
    }
}

impl<S> Layer<S> for EventDispatcherLayer {
    type Service = DispatcherService<S>;

    fn layer(&self, service: S) -> Self::Service {
        DispatcherService {
            service,
            dispatcher: self.dispatcher.dupe(),
        }
    }
}

#[derive(Clone)]
pub struct DispatcherService<S> {
    service: S,
    dispatcher: EventDispatcher,
}

impl<T, Request> tower_service::Service<Request> for DispatcherService<T>
where
    T: tower_service::Service<Request>,
    T::Future: 'static + Send,
    T::Error: 'static,
    T::Response: 'static,
{
    type Response = T::Response;
    type Error = T::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, req: Request) -> Self::Future {
        // Call the inner service and get a future that resolves to the response
        let fut = self.service.call(req);

        with_dispatcher_async(self.dispatcher.dupe(), fut).boxed()
    }
}


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

use std::time::Duration;

use anyhow::Context as _;
use assert_matches::assert_matches;
use buck2_grpc::DuplexChannel;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;

use crate::data::ExternalRunnerSpec;
use crate::grpc::spawn_executor_server;
use crate::grpc::TestExecutorClient;
use crate::protocol::TestExecutor;

struct MockExecutor;

#[async_trait::async_trait]
impl TestExecutor for MockExecutor {
    async fn external_runner_spec(&self, _: ExternalRunnerSpec) -> anyhow::Result<()> {
        Ok(())
    }

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

#[tokio::test]
async fn test_basic() -> anyhow::Result<()> {
    let (client_io, server_io) = tokio::io::duplex(64);
    let server = spawn_executor_server(to_duplex_channel(server_io), MockExecutor);
    let client = TestExecutorClient::new(client_io)
        .await
        .context("Failed to create client")?;

    client.end_of_test_requests().await.context("Call 1")?;
    client.end_of_test_requests().await.context("Call 2")?;
    client.end_of_test_requests().await.context("Call 3")?;

    server.shutdown().await.context("Failed to shutdown")?;

    assert_matches!(client.end_of_test_requests().await, Err(..));

    Ok(())
}

#[tokio::test]
async fn test_client_disconnect() -> anyhow::Result<()> {
    let (client_io, server_io) = tokio::io::duplex(64);
    let mut server =
        spawn_executor_server(to_duplex_channel(server_io), MockExecutor).into_join_handle();

    assert_matches!(
        tokio::time::timeout(Duration::from_millis(10), &mut server).await,
        Err(..)
    );

    drop(client_io);

    assert_matches!(
        tokio::time::timeout(Duration::from_millis(10), &mut server).await,
        Ok(..)
    );

    Ok(())
}

fn to_duplex_channel(
    io: impl AsyncRead + AsyncWrite + Send + Unpin + 'static,
) -> impl AsyncRead + AsyncWrite + Send + Unpin + 'static + tonic::transport::server::Connected {
    let (read, write) = tokio::io::split(io);
    DuplexChannel::new(read, write)
}


/*
 * 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 as _;
use host_sharing::HostSharingRequirements;
use host_sharing::WeightClass;
use host_sharing::WeightPercentage;

pub fn weight_class_from_grpc(input: buck2_test_proto::WeightClass) -> anyhow::Result<WeightClass> {
    use buck2_test_proto::weight_class::*;

    Ok(match input.value.context("Missing `value`")? {
        Value::Permits(p) => WeightClass::Permits(p.try_into().context("Invalid `permits`")?),
        Value::Percentage(p) => {
            WeightClass::Percentage(WeightPercentage::try_new(p).context("Invalid `percentage`")?)
        }
    })
}

pub fn host_sharing_requirements_from_grpc(
    input: buck2_test_proto::HostSharingRequirements,
) -> anyhow::Result<HostSharingRequirements> {
    use buck2_test_proto::host_sharing_requirements::*;

    let requirements = match input.requirements.context("Missing `requirements`")? {
        Requirements::Shared(Shared { weight_class }) => {
            host_sharing::HostSharingRequirements::Shared(weight_class_from_grpc(
                weight_class.context("Missing `weight_class`")?,
            )?)
        }
        Requirements::ExclusiveAccess(ExclusiveAccess {}) => {
            host_sharing::HostSharingRequirements::ExclusiveAccess
        }
        Requirements::OnePerToken(OnePerToken {
            identifier,
            weight_class,
        }) => host_sharing::HostSharingRequirements::OnePerToken(
            identifier,
            weight_class_from_grpc(weight_class.context("Missing `weight_class`")?)?,
        ),
    };

    Ok(requirements)
}

pub fn weight_class_to_grpc(input: WeightClass) -> anyhow::Result<buck2_test_proto::WeightClass> {
    use buck2_test_proto::weight_class::*;

    let value = match input {
        host_sharing::WeightClass::Permits(p) => {
            Value::Permits(p.try_into().context("Invalid `permits`")?)
        }
        host_sharing::WeightClass::Percentage(p) => Value::Percentage(p.into_value().into()),
    };

    Ok(buck2_test_proto::WeightClass { value: Some(value) })
}

pub fn host_sharing_requirements_to_grpc(
    input: HostSharingRequirements,
) -> anyhow::Result<buck2_test_proto::HostSharingRequirements> {
    use buck2_test_proto::host_sharing_requirements::*;

    let requirements = match input {
        host_sharing::HostSharingRequirements::Shared(weight) => Requirements::Shared(Shared {
            weight_class: Some(weight_class_to_grpc(weight)?),
        }),
        host_sharing::HostSharingRequirements::ExclusiveAccess => {
            Requirements::ExclusiveAccess(ExclusiveAccess {})
        }
        host_sharing::HostSharingRequirements::OnePerToken(identifier, weight) => {
            Requirements::OnePerToken(OnePerToken {
                identifier,
                weight_class: Some(weight_class_to_grpc(weight)?),
            })
        }
    };

    Ok(buck2_test_proto::HostSharingRequirements {
        requirements: Some(requirements),
    })
}


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

//! The traits that defines the protocol between buck and a test executor.
//!
//! Test executors are expected to implement the trait `TestExecutor`. Buck will need to implement
//! the trait `Buck` for the test executor to interact against.

pub mod convert;

use std::time::Duration;

use dupe::Dupe;
use host_sharing::HostSharingRequirements;
use sorted_vector_map::SortedVectorMap;

use crate::data::ArgValue;
use crate::data::ConfiguredTargetHandle;
use crate::data::DeclaredOutput;
use crate::data::DisplayMetadata;
use crate::data::ExecuteResponse;
use crate::data::ExecutorConfigOverride;
use crate::data::ExternalRunnerSpec;
use crate::data::PrepareForLocalExecutionResult;
use crate::data::RequiredLocalResources;
use crate::data::TestResult;

/// available to buck to interact with the test executor
#[async_trait::async_trait]
pub trait TestExecutor: Send + Sync {
    /// sends an external runner spec to the test executor
    async fn external_runner_spec(&self, s: ExternalRunnerSpec) -> anyhow::Result<()>;

    // report that there are no more test specs to send
    async fn end_of_test_requests(&self) -> anyhow::Result<()>;
}

// available to the test executor to interact with the orchestrator
#[async_trait::async_trait]
pub trait TestOrchestrator: Send + Sync {
    /// executes the given command and returns the result of the execution
    async fn execute2(
        &self,
        // information about this execute request for Buck's UX
        ui_prints: DisplayMetadata,
        // the label of the rule being tested
        target: ConfiguredTargetHandle,
        // the command to run
        cmd: Vec<ArgValue>,
        // environment variables to set at runtime
        env: SortedVectorMap<String, ArgValue>,
        // timeout for command
        timeout: Duration,
        // parameters used to effectively share the executor host for this command.
        host_sharing_requirements: HostSharingRequirements,
        // outputs that need to be pre created as directories
        pre_create_dirs: Vec<DeclaredOutput>,
        // A specific executor to use for this. It must be declared on the underlying
        // ExternalRunnerTestInfo to work.
        executor_override: Option<ExecutorConfigOverride>,
        required_local_resources: RequiredLocalResources,
    ) -> anyhow::Result<ExecuteResponse>;

    /// reports a test is done
    async fn report_test_result(&self, r: TestResult) -> anyhow::Result<()>;

    async fn report_tests_discovered(
        &self,
        target: ConfiguredTargetHandle,
        suite: String,
        name: Vec<String>,
    ) -> anyhow::Result<()>;

    /// report a summary about the current test executor
    async fn report_test_session(&self, session_info: String) -> anyhow::Result<()>;

    /// report that all tests are done and provide the exit code that this test executor wants to
    /// return for the test command, no more executions
    async fn end_of_test_results(&self, exit_code: i32) -> anyhow::Result<()>;

    /// prepare the given test executable to be available for local execution.
    /// Return the actual command with all the args, env and cwd to be executed locally.
    async fn prepare_for_local_execution(
        &self,
        ui_prints: DisplayMetadata,
        target: ConfiguredTargetHandle,
        cmd: Vec<ArgValue>,
        env: SortedVectorMap<String, ArgValue>,
        pre_create_dirs: Vec<DeclaredOutput>,
    ) -> anyhow::Result<PrepareForLocalExecutionResult>;

    /// attach a message containing information that the executor wants to be surfaced
    /// to the user
    async fn attach_info_message(&self, message: String) -> anyhow::Result<()>;
}

// TODO need to figure out what this is. we can go without it for now
#[derive(Debug, Default, Clone, Dupe)]
pub struct ExecPlatformRefinement;

/// the template name for external test executor to fill with the test execution location
pub const OUTPUT_DIR: &str = "output_dir";

/// the template name for buck to fill with the commands outputs
pub const OUTPUTS_TEMPLATE: &str = "outputs";

/// the template name for buck to fill with the run cmd based on the test's run information
/// this should be multi-arity
pub const TEST_RUN_CMD: &str = "test_run_cmd";

/// the template name for buck to fill with the run environment based on the test's run information
/// this should be multi-arity
pub const TEST_RUN_ENV_KEYS: &str = "test_run_env_keys";

/// the template name for buck to fill with the run environment based on the test's run information
/// this should be multi-arity
pub const TEST_RUN_ENV_VALUES: &str = "test_run_env_values";

/// the template name for the external test runner to fill with test filters
pub const TEST_FILTER_TEMPLATE: &str = "test_filter";


/*
 * 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 = &["test.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!("buck.test");


/*
 * 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::str::FromStr;
use std::time::Duration;

use anyhow::Context;
use clap::Parser;

#[derive(Debug, Parser)]
pub struct Config {
    /// add a list of environment variables using format: --env VAR1=Value1 VAR2='Value 2'
    #[clap(long)]
    pub env: Vec<EnvValue>,

    /// Max number of seconds allowed to run a test.
    #[clap(long, default_value = "600", parse(try_from_str=try_parse_timeout_from_str))]
    pub timeout: Duration,

    #[clap(flatten)]
    ignored_args: IgnoredArgs,
}

/// Ignored args included for backwards compatibility.
#[derive(Debug, Parser)]
struct IgnoredArgs {
    #[clap(long, hidden = true)]
    buck_test_info: String,
}

/// Uiltity that can be used to parse Env values from CLI arguments.
#[derive(Debug, PartialEq)]
pub struct EnvValue {
    pub name: String,
    pub value: String,
}

impl EnvValue {
    pub fn new(name: &str, value: &str) -> EnvValue {
        EnvValue {
            name: name.to_owned(),
            value: value.to_owned(),
        }
    }
}

impl FromStr for EnvValue {
    type Err = EnvValueParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        match input.split_once('=') {
            Some((key, value)) => Ok(EnvValue::new(key, value)),
            None => Err(EnvValueParseError::IncorrectSyntax(input.to_owned())),
        }
    }
}

#[derive(Debug, thiserror::Error, PartialEq)]
pub enum EnvValueParseError {
    #[error("Incorrect syntax for env value. Please use name=value. Input: `{0}`")]
    IncorrectSyntax(String),
}

fn try_parse_timeout_from_str(input: &str) -> anyhow::Result<Duration> {
    let seconds = input.parse().context("Could not parse provided timeout")?;
    Ok(Duration::from_secs(seconds))
}


/*
 * 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_test_api::data::ExternalRunnerSpec;
use buck2_test_api::protocol::TestExecutor;
use futures::channel::mpsc::UnboundedSender;

pub type SpecSender = UnboundedSender<ExternalRunnerSpec>;

pub struct Buck2TestExecutor {
    pub sender: SpecSender,
}

impl Buck2TestExecutor {
    pub fn new(sender: SpecSender) -> Self {
        Self { sender }
    }
}

#[async_trait::async_trait]
impl TestExecutor for Buck2TestExecutor {
    async fn external_runner_spec(&self, spec: ExternalRunnerSpec) -> anyhow::Result<()> {
        self.sender
            .clone()
            .start_send(spec)
            .expect("Sending to not fail if all core invariants are held.");
        Ok(())
    }

    async fn end_of_test_requests(&self) -> anyhow::Result<()> {
        // This ensures that all senders are dropped so the receiver can terminate
        self.sender.close_channel();
        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.
 */

#![feature(async_closure)]

mod config;
mod executor;
mod runner;
mod service;
pub mod tcp;

#[cfg(unix)]
pub mod unix;


/*
 * 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 buck2_test_api::data::ArgValue;
use buck2_test_api::data::ArgValueContent;
use buck2_test_api::data::ConfiguredTargetHandle;
use buck2_test_api::data::DisplayMetadata;
use buck2_test_api::data::ExecuteResponse;
use buck2_test_api::data::ExecutionResult2;
use buck2_test_api::data::ExecutionStatus;
use buck2_test_api::data::ExternalRunnerSpec;
use buck2_test_api::data::ExternalRunnerSpecValue;
use buck2_test_api::data::RequiredLocalResources;
use buck2_test_api::data::TestResult;
use buck2_test_api::data::TestStatus;
use buck2_test_api::grpc::TestOrchestratorClient;
use clap::Parser;
use futures::channel::mpsc::UnboundedReceiver;
use futures::StreamExt;
use host_sharing::HostSharingRequirements;
use parking_lot::Mutex;

use crate::config::Config;
use crate::config::EnvValue;

pub type SpecReceiver = UnboundedReceiver<ExternalRunnerSpec>;

/// Internal test runner implementation for Buck2.
///
/// This is a basic test runner intended to be used by the open-source Buck2 build
/// if no external test runner is provided. This ensures that `buck2 test` works
/// out-of-the-box for open-source users.
///
/// **This is intended for open-source use only.**
pub struct Buck2TestRunner {
    orchestrator_client: TestOrchestratorClient,
    spec_receiver: Mutex<Option<SpecReceiver>>,
    config: Config,
}

impl Buck2TestRunner {
    pub fn new(
        orchestrator_client: TestOrchestratorClient,
        spec_receiver: SpecReceiver,
        args: Vec<String>,
    ) -> anyhow::Result<Self> {
        let config = Config::try_parse_from(args).context("Error parsing test runner arguments")?;
        Ok(Self {
            orchestrator_client,
            spec_receiver: Mutex::new(Some(spec_receiver)),
            config,
        })
    }

    pub async fn run_all_tests(&self) -> anyhow::Result<()> {
        let receiver;
        {
            let mut maybe_receiver = self.spec_receiver.lock();
            receiver = maybe_receiver
                .take()
                .context("Spec channel has already been consumed")?;
            drop(maybe_receiver);
        }
        let run_verdict = receiver
            .map(async move |spec| {
                let name = format!(
                    "{}//{}:{}",
                    spec.target.cell, spec.target.package, spec.target.target
                );
                let target_handle = spec.target.handle.to_owned();

                let execution_response = self
                    .execute_test_from_spec(spec)
                    .await
                    .expect("Test execution request failed");

                let execution_result = match execution_response {
                    ExecuteResponse::Result(r) => r,
                    ExecuteResponse::Cancelled => return TestStatus::OMITTED,
                };

                let test_result = get_test_result(name, target_handle, execution_result);
                let test_status = test_result.status.clone();

                self.report_test_result(test_result)
                    .await
                    .expect("Test result reporting failed");

                test_status
            })
            // Use an arbitrarily large buffer -- execution throttling will be handled by the Buck2
            // executor, so no need to hold back on requests here.
            .buffer_unordered(10000)
            // If any individual test failed, consider the entire run to have failed.
            .fold(
                RunVerdict::Pass,
                async move |mut run_verdict, test_status| {
                    if test_status != TestStatus::PASS {
                        run_verdict = RunVerdict::Fail;
                    }
                    run_verdict
                },
            )
            .await;

        self.orchestrator_client
            .end_of_test_results(run_verdict.exit_code())
            .await
    }

    async fn execute_test_from_spec(
        &self,
        spec: ExternalRunnerSpec,
    ) -> anyhow::Result<ExecuteResponse> {
        let display_metadata = DisplayMetadata::Testing {
            suite: spec.target.target,
            testcases: Vec::new(),
        };

        let command = spec
            .command
            .into_iter()
            .map(|spec_value| ArgValue {
                content: ArgValueContent::ExternalRunnerSpecValue(spec_value),
                format: None,
            })
            .collect();

        let config_env = self.config.env.iter().map(|EnvValue { name, value }| {
            (
                name.to_owned(),
                ArgValue {
                    content: ArgValueContent::ExternalRunnerSpecValue(
                        ExternalRunnerSpecValue::Verbatim(value.to_owned()),
                    ),
                    format: None,
                },
            )
        });

        let env = spec
            .env
            .into_iter()
            .map(|(key, value)| {
                (
                    key,
                    ArgValue {
                        content: ArgValueContent::ExternalRunnerSpecValue(value),
                        format: None,
                    },
                )
            })
            .chain(config_env)
            .collect();

        let target_handle = spec.target.handle;
        let host_sharing_requirements = HostSharingRequirements::default();
        let pre_create_dirs = Vec::new();
        let executor_override = None;

        self.orchestrator_client
            .execute2(
                display_metadata,
                target_handle,
                command,
                env,
                self.config.timeout,
                host_sharing_requirements,
                pre_create_dirs,
                executor_override,
                RequiredLocalResources { resources: vec![] },
            )
            .await
    }

    async fn report_test_result(&self, test_result: TestResult) -> anyhow::Result<()> {
        self.orchestrator_client
            .report_test_result(test_result)
            .await
    }
}

fn get_test_result(
    name: String,
    target: ConfiguredTargetHandle,
    execution_result: ExecutionResult2,
) -> TestResult {
    let status = match execution_result.status {
        ExecutionStatus::Finished { exitcode } => match exitcode {
            0 => TestStatus::PASS,
            _ => TestStatus::FAIL,
        },
        ExecutionStatus::TimedOut { .. } => TestStatus::TIMEOUT,
    };
    TestResult {
        target,
        name,
        status,
        msg: None,
        duration: Some(execution_result.execution_time),
        details: format!(
            "---- STDOUT ----\n{:?}\n---- STDERR ----\n{:?}\n",
            execution_result.stdout, execution_result.stderr
        ),
    }
}

#[derive(Debug)]
enum RunVerdict {
    Pass,
    Fail,
}

impl RunVerdict {
    fn exit_code(&self) -> i32 {
        match self {
            RunVerdict::Pass => 0,
            RunVerdict::Fail => 32,
        }
    }
}


/*
 * 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 buck2_grpc::DuplexChannel;
use buck2_test_api::grpc::spawn_executor_server;
use buck2_test_api::grpc::TestOrchestratorClient;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;

use crate::executor::Buck2TestExecutor;
use crate::runner::Buck2TestRunner;

pub async fn run<OC, ER, EW>(
    orchestrator_channel: OC,
    executor_channel: DuplexChannel<ER, EW>,
    args: Vec<String>,
) -> anyhow::Result<()>
where
    OC: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
    ER: AsyncRead + Send + Unpin + 'static,
    EW: AsyncWrite + Send + Unpin + 'static,
{
    let (spec_sender, spec_receiver) = futures::channel::mpsc::unbounded();

    let executor_server =
        spawn_executor_server(executor_channel, Buck2TestExecutor::new(spec_sender));

    let orchestrator_client = TestOrchestratorClient::new(orchestrator_channel)
        .await
        .context("Failed to TestOrchestratorClient")?;

    let runner = Buck2TestRunner::new(orchestrator_client, spec_receiver, args)?;

    runner.run_all_tests().await?;

    executor_server
        .shutdown()
        .await
        .context("Failed to shutdown server")?;

    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::net::SocketAddr;

use buck2_grpc::DuplexChannel;
use clap::Parser;
use tokio::net::TcpStream;

#[derive(Debug, Parser)]
pub struct Buck2TestRunnerTcp {
    #[clap(long)]
    executor_addr: String,

    #[clap(long)]
    orchestrator_addr: String,

    args: Vec<String>,
}

impl Buck2TestRunnerTcp {
    pub async fn run(self) -> anyhow::Result<()> {
        // NOTE: We can remove this code once Tokio supports UNIX domain socket or Named pipe
        // https://github.com/tokio-rs/tokio/issues/2201

        // NOTE: We assume the parameters we received from the caller are correct here. If
        // they're not, things are probably going to go wrong but that's on our caller.
        //
        // There are a few ways in which the params could be incorrect:
        // - The TCP socket addresses aren't valid
        // - The TCP socket don't exist
        // - The TCP socket aren't alive
        let orchestrator_addr: SocketAddr = self
            .orchestrator_addr
            .parse()
            .expect("Invalid orchestrator address");
        let executor_addr: SocketAddr = self
            .executor_addr
            .parse()
            .expect("Invalid executor address");

        let orchestrator_io = TcpStream::connect(&orchestrator_addr)
            .await
            .expect("Failed to create orchestrator_io");

        let executor_io = TcpStream::connect(&executor_addr)
            .await
            .expect("Failed to create executor_io");

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

        crate::service::run(orchestrator_io, executor_io, self.args).await
    }
}


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

use std::os::unix::io::FromRawFd;
use std::os::unix::io::RawFd;
use std::os::unix::net::UnixStream as StdUnixStream;

use buck2_grpc::DuplexChannel;
use clap::Parser;
use tokio::net::UnixStream;

#[derive(Debug, Parser)]
pub struct Buck2TestRunnerUnix {
    #[clap(long)]
    executor_fd: RawFd,

    #[clap(long)]
    orchestrator_fd: RawFd,

    args: Vec<String>,
}

impl Buck2TestRunnerUnix {
    pub async fn run(self) -> anyhow::Result<()> {
        // NOTE: We assume the parameters we received from the caller are correct here. If
        // they're not, things are probably going to go wrong but that's on our caller.
        //
        // There are a few ways in which the params could be incorrect:
        // - The FDs don't exist
        // - The FDs aren't streams
        // - The FDs are the same
        //
        // In all those cases though, the unsafety below is going to result in bad file
        // descriptors at worse, which is basically the best we can do anyway.
        let orchestrator_io =
            UnixStream::from_std(unsafe { StdUnixStream::from_raw_fd(self.orchestrator_fd) })
                .expect("Failed to create orchestrator_io");

        let executor_io =
            UnixStream::from_std(unsafe { StdUnixStream::from_raw_fd(self.executor_fd) })
                .expect("Failed to create executor_io");

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

        crate::service::run(orchestrator_io, executor_io, self.args).await
    }
}


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

use buck2_node::attrs::coerced_attr::CoercedAttr;
use buck2_node::attrs::display::AttrDisplayWithContextExt;
use starlark::values::dict::Dict;
use starlark::values::list::AllocList;
use starlark::values::tuple::AllocTuple;
use starlark::values::Heap;
use starlark::values::Value;
use starlark_map::small_map::SmallMap;

#[derive(Debug, thiserror::Error)]
enum CoercedAttrResolveError {
    #[error("Attribute cannot be converted to Starlark value: `{0}`")]
    AttrCannotBeConvertedToValue(String),
}

pub(crate) trait CoercedAttrResolveExt {
    fn to_value<'v>(&self, heap: &'v Heap) -> anyhow::Result<Value<'v>>;
}

impl CoercedAttrResolveExt for CoercedAttr {
    fn to_value<'v>(&self, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        match self {
            x @ (CoercedAttr::Concat(..) | CoercedAttr::Selector(..)) => {
                // It is possible to convert selects back to Starlark objects,
                // but there's no need to do it for now (and probably never will).
                Err(CoercedAttrResolveError::AttrCannotBeConvertedToValue(
                    x.as_display_no_ctx().to_string(),
                )
                .into())
            }
            CoercedAttr::None => Ok(Value::new_none()),
            CoercedAttr::Bool(b) => Ok(Value::new_bool(b.0)),
            CoercedAttr::Int(i) => Ok(heap.alloc(*i)),
            CoercedAttr::String(s) | CoercedAttr::EnumVariant(s) => {
                Ok(heap.alloc_str(s).to_value())
            }
            CoercedAttr::List(l) => {
                let mut v = Vec::with_capacity(l.len());
                for e in l.iter() {
                    v.push(e.to_value(heap)?);
                }
                Ok(heap.alloc(AllocList(v)))
            }
            CoercedAttr::Tuple(l) => {
                let mut v = Vec::with_capacity(l.len());
                for e in l.iter() {
                    v.push(e.to_value(heap)?);
                }
                Ok(heap.alloc(AllocTuple(v)))
            }
            CoercedAttr::Dict(d) => {
                let mut m = SmallMap::with_capacity(d.len());
                for (k, v) in d.iter() {
                    m.insert_hashed(k.to_value(heap)?.get_hashed()?, v.to_value(heap)?);
                }
                Ok(heap.alloc(Dict::new(m)))
            }
            x => {
                // For now this function is used to convert attributes to Starlark values
                // for transition rules which access attributes.
                //
                // For regular deps this function should fail.
                //
                // For configuration deps, this function should resolve attributes to providers,
                // but it is not implemented yet.
                Err(CoercedAttrResolveError::AttrCannotBeConvertedToValue(
                    x.as_display_no_ctx().to_string(),
                )
                .into())
            }
        }
    }
}


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

#![feature(try_blocks)]

pub(crate) mod coerced_attr;
pub(crate) mod transition;

pub fn init_late_bindings() {
    transition::calculation_apply_transition::init_transition_calculation();
    transition::starlark::init_register_transition();
}