action_messages.rs

// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::time::{Duration, SystemTime};

use nativelink_error::{error_if, make_input_err, Error, ResultExt};
use nativelink_proto::build::bazel::remote::execution::v2::{
    execution_stage, Action, ActionResult as ProtoActionResult, ExecuteOperationMetadata,
    ExecuteRequest, ExecuteResponse, ExecutedActionMetadata, FileNode, LogFile, OutputDirectory,
    OutputFile, OutputSymlink, SymlinkNode,
};
use nativelink_proto::google::longrunning::operation::Result as LongRunningResult;
use nativelink_proto::google::longrunning::Operation;
use nativelink_proto::google::rpc::Status;
use prost::bytes::Bytes;
use prost::Message;
use prost_types::Any;
use serde::ser::Error as SerdeError;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::common::{DigestInfo, HashMapExt, VecExt};
use crate::digest_hasher::DigestHasherFunc;
use crate::metrics_utils::{CollectorState, MetricsComponent};
use crate::platform_properties::PlatformProperties;

/// Default priority remote execution jobs will get when not provided.
pub const DEFAULT_EXECUTION_PRIORITY: i32 = 0;

/// Exit code sent if there is an internal error.
pub const INTERNAL_ERROR_EXIT_CODE: i32 = -178;

/// Holds an id that is unique to the client for a requested operation.
/// Each client should be issued a unique id even if they are attached
/// to the same underlying operation.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ClientOperationId(String);

impl ClientOperationId {
    pub fn new(unique_qualifier: ActionUniqueQualifier) -> Self {
        Self(OperationId::new(unique_qualifier).to_string())
    }

    pub fn from_raw_string(name: String) -> Self {
        Self(name)
    }

    pub fn into_string(self) -> String {
        self.0
    }
}

impl std::fmt::Display for ClientOperationId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.0.clone()))
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct OperationId {
    pub unique_qualifier: ActionUniqueQualifier,
    pub id: Uuid,
}

impl PartialEq for OperationId {
    fn eq(&self, other: &Self) -> bool {
        self.id.eq(&other.id)
    }
}

impl Eq for OperationId {}

impl PartialOrd for OperationId {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for OperationId {
    fn cmp(&self, other: &Self) -> Ordering {
        self.id.cmp(&other.id)
    }
}

impl Hash for OperationId {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.id.hash(state)
    }
}

impl OperationId {
    pub fn new(unique_qualifier: ActionUniqueQualifier) -> Self {
        Self {
            id: Uuid::new_v4(),
            unique_qualifier,
        }
    }
}

impl TryFrom<&str> for OperationId {
    type Error = Error;

    /// Attempts to convert a string slice into an `OperationId`.
    ///
    /// The input string `value` is expected to be in the format:
    /// `<instance_name>/<digest_function>/<digest_hash>-<digest_size>/<cached>/<id>`.
    ///
    /// # Parameters
    ///
    /// - `value`: A `&str` representing the `OperationId` to be converted.
    ///
    /// # Returns
    ///
    /// - `Result<OperationId, Error>`: Returns `Ok(OperationId)` if the conversion is
    /// successful, or an `Err(Error)` if it fails.
    ///
    /// # Errors
    ///
    /// This function will return an error if the input string is not in the expected format or if
    /// any of the components cannot be parsed correctly. The detailed error messages provide
    /// insight into which part of the input string caused the failure.
    ///
    /// ## Example Usage
    ///
    /// ```no_run
    /// use nativelink_util::action_messages::OperationId;
    /// let operation_id_str = "main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/u/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52";
    /// let operation_id = OperationId::try_from(operation_id_str);
    /// ```
    ///
    /// In this example, `operation_id_str` is a string representing an `OperationId`.
    /// The `try_from` method is used to convert it into an `OperationId` instance.
    /// If any part of the string is incorrectly formatted, an error will be returned with a
    /// descriptive message.
    fn try_from(value: &str) -> Result<Self, Error> {
        let (unique_qualifier, id) = value
            .rsplit_once('/')
            .err_tip(|| format!("Invalid OperationId unique_qualifier / id fragment - {value}"))?;
        let (instance_name, rest) = unique_qualifier
            .split_once('/')
            .err_tip(|| format!("Invalid UniqueQualifier instance name fragment - {value}"))?;
        let (digest_function, rest) = rest
            .split_once('/')
            .err_tip(|| format!("Invalid UniqueQualifier digest function fragment - {value}"))?;
        let (digest_hash, rest) = rest
            .split_once('-')
            .err_tip(|| format!("Invalid UniqueQualifier digest hash fragment - {value}"))?;
        let (digest_size, cachable) = rest
            .split_once('/')
            .err_tip(|| format!("Invalid UniqueQualifier digest size fragment - {value}"))?;
        let digest = DigestInfo::try_new(
            digest_hash,
            digest_size
                .parse::<u64>()
                .err_tip(|| format!("Invalid UniqueQualifier size value fragment - {value}"))?,
        )
        .err_tip(|| format!("Invalid DigestInfo digest hash - {value}"))?;
        let cachable = match cachable {
            "u" => false,
            "c" => true,
            _ => {
                return Err(make_input_err!(
                    "Invalid UniqueQualifier cachable value fragment - {value}"
                ));
            }
        };
        let unique_key = ActionUniqueKey {
            instance_name: instance_name.to_string(),
            digest_function: digest_function.try_into()?,
            digest,
        };
        let unique_qualifier = if cachable {
            ActionUniqueQualifier::Cachable(unique_key)
        } else {
            ActionUniqueQualifier::Uncachable(unique_key)
        };
        let id = Uuid::parse_str(id).map_err(|e| make_input_err!("Failed to parse {e} as uuid"))?;
        Ok(Self {
            unique_qualifier,
            id,
        })
    }
}

impl std::fmt::Display for OperationId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}/{}", self.unique_qualifier, self.id))
    }
}

impl std::fmt::Debug for OperationId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self, f)
    }
}

/// Unique id of worker.
#[derive(Default, Eq, PartialEq, Hash, Copy, Clone, Serialize, Deserialize)]
pub struct WorkerId(pub Uuid);

impl std::fmt::Display for WorkerId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut buf = Uuid::encode_buffer();
        let worker_id_str = self.0.hyphenated().encode_lower(&mut buf);
        f.write_fmt(format_args!("{worker_id_str}"))
    }
}

impl std::fmt::Debug for WorkerId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self, f)
    }
}

impl TryFrom<String> for WorkerId {
    type Error = Error;
    fn try_from(s: String) -> Result<Self, Self::Error> {
        match Uuid::parse_str(&s) {
            Err(e) => Err(make_input_err!(
                "Failed to convert string to WorkerId : {s} : {e:?}",
            )),
            Ok(my_uuid) => Ok(WorkerId(my_uuid)),
        }
    }
}

/// Holds the information needed to uniquely identify an action
/// and if it is cachable or not.
#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
pub enum ActionUniqueQualifier {
    /// The action is cachable.
    Cachable(ActionUniqueKey),
    /// The action is uncachable.
    Uncachable(ActionUniqueKey),
}

impl ActionUniqueQualifier {
    /// Get the instance_name of the action.
    pub const fn instance_name(&self) -> &String {
        match self {
            Self::Cachable(action) => &action.instance_name,
            Self::Uncachable(action) => &action.instance_name,
        }
    }

    /// Get the digest function of the action.
    pub const fn digest_function(&self) -> DigestHasherFunc {
        match self {
            Self::Cachable(action) => action.digest_function,
            Self::Uncachable(action) => action.digest_function,
        }
    }

    /// Get the digest of the action.
    pub const fn digest(&self) -> DigestInfo {
        match self {
            Self::Cachable(action) => action.digest,
            Self::Uncachable(action) => action.digest,
        }
    }
}

impl std::fmt::Display for ActionUniqueQualifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let (cachable, unique_key) = match self {
            Self::Cachable(action) => (true, action),
            Self::Uncachable(action) => (false, action),
        };
        f.write_fmt(format_args!(
            "{}/{}/{}-{}/{}",
            unique_key.instance_name,
            unique_key.digest_function,
            unique_key.digest.hash_str(),
            unique_key.digest.size_bytes,
            if cachable { 'c' } else { 'u' },
        ))
    }
}

/// This is a utility struct used to make it easier to match `ActionInfos` in a
/// `HashMap` without needing to construct an entire `ActionInfo`.
#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
pub struct ActionUniqueKey {
    /// Name of instance group this action belongs to.
    pub instance_name: String,
    /// The digest function this action expects.
    pub digest_function: DigestHasherFunc,
    /// Digest of the underlying `Action`.
    pub digest: DigestInfo,
}

/// Information needed to execute an action. This struct is used over bazel's proto `Action`
/// for simplicity and offers a `salt`, which is useful to ensure during hashing (for dicts)
/// to ensure we never match against another `ActionInfo` (when a task should never be cached).
/// This struct must be 100% compatible with `ExecuteRequest` struct in `remote_execution.proto`
/// except for the salt field.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ActionInfo {
    /// Digest of the underlying `Command`.
    pub command_digest: DigestInfo,
    /// Digest of the underlying `Directory`.
    pub input_root_digest: DigestInfo,
    /// Timeout of the action.
    pub timeout: Duration,
    /// The properties rules that must be applied when finding a worker that can run this action.
    pub platform_properties: PlatformProperties,
    /// The priority of the action. Higher value means it should execute faster.
    pub priority: i32,
    /// When this action started to be loaded from the CAS.
    pub load_timestamp: SystemTime,
    /// When this action was created.
    pub insert_timestamp: SystemTime,
    /// Info used to uniquely identify this ActionInfo and if it is cachable.
    /// This is primarily used to join actions/operations together using this key.
    pub unique_qualifier: ActionUniqueQualifier,
}

impl ActionInfo {
    #[inline]
    pub const fn instance_name(&self) -> &String {
        self.unique_qualifier.instance_name()
    }

    /// Returns the underlying digest of the `Action`.
    #[inline]
    pub const fn digest(&self) -> DigestInfo {
        self.unique_qualifier.digest()
    }

    pub fn try_from_action_and_execute_request(
        execute_request: ExecuteRequest,
        action: Action,
        load_timestamp: SystemTime,
        queued_timestamp: SystemTime,
    ) -> Result<Self, Error> {
        let unique_key = ActionUniqueKey {
            instance_name: execute_request.instance_name,
            digest_function: DigestHasherFunc::try_from(execute_request.digest_function)
                .err_tip(|| format!("Could not find digest_function in try_from_action_and_execute_request {:?}", execute_request.digest_function))?,
            digest: execute_request
                .action_digest
                .err_tip(|| "Expected action_digest to exist on ExecuteRequest")?
                .try_into()?,
        };
        let unique_qualifier = if execute_request.skip_cache_lookup {
            ActionUniqueQualifier::Uncachable(unique_key)
        } else {
            ActionUniqueQualifier::Cachable(unique_key)
        };
        Ok(Self {
            command_digest: action
                .command_digest
                .err_tip(|| "Expected command_digest to exist on Action")?
                .try_into()?,
            input_root_digest: action
                .input_root_digest
                .err_tip(|| "Expected input_root_digest to exist on Action")?
                .try_into()?,
            timeout: action
                .timeout
                .unwrap_or_default()
                .try_into()
                .map_err(|_| make_input_err!("Failed convert proto duration to system duration"))?,
            platform_properties: action.platform.unwrap_or_default().into(),
            priority: execute_request
                .execution_policy
                .unwrap_or_default()
                .priority,
            load_timestamp,
            insert_timestamp: queued_timestamp,
            unique_qualifier,
        })
    }
}

impl From<ActionInfo> for ExecuteRequest {
    fn from(val: ActionInfo) -> Self {
        let digest = val.digest().into();
        let (skip_cache_lookup, unique_qualifier) = match val.unique_qualifier {
            ActionUniqueQualifier::Cachable(unique_qualifier) => (false, unique_qualifier),
            ActionUniqueQualifier::Uncachable(unique_qualifier) => (true, unique_qualifier),
        };
        Self {
            instance_name: unique_qualifier.instance_name,
            action_digest: Some(digest),
            skip_cache_lookup,
            execution_policy: None,     // Not used in the worker.
            results_cache_policy: None, // Not used in the worker.
            digest_function: unique_qualifier.digest_function.proto_digest_func().into(),
        }
    }
}

/// Simple utility struct to determine if a string is representing a full path or
/// just the name of the file.
/// This is in order to be able to reuse the same struct instead of building different
/// structs when converting `FileInfo` -> {`OutputFile`, `FileNode`} and other similar
/// structs.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub enum NameOrPath {
    Name(String),
    Path(String),
}

impl PartialOrd for NameOrPath {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for NameOrPath {
    fn cmp(&self, other: &Self) -> Ordering {
        let self_lexical_name = match self {
            Self::Name(name) => name,
            Self::Path(path) => path,
        };
        let other_lexical_name = match other {
            Self::Name(name) => name,
            Self::Path(path) => path,
        };
        self_lexical_name.cmp(other_lexical_name)
    }
}

/// Represents an individual file and associated metadata.
/// This struct must be 100% compatible with `OutputFile` and `FileNode` structs
/// in `remote_execution.proto`.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct FileInfo {
    pub name_or_path: NameOrPath,
    pub digest: DigestInfo,
    pub is_executable: bool,
}

//TODO: Make this TryFrom.
impl From<FileInfo> for FileNode {
    fn from(val: FileInfo) -> Self {
        let NameOrPath::Name(name) = val.name_or_path else {
            panic!("Cannot return a FileInfo that uses a NameOrPath::Path(), it must be a NameOrPath::Name()");
        };
        Self {
            name,
            digest: Some((&val.digest).into()),
            is_executable: val.is_executable,
            node_properties: Option::default(), // Not supported.
        }
    }
}

impl TryFrom<OutputFile> for FileInfo {
    type Error = Error;

    fn try_from(output_file: OutputFile) -> Result<Self, Error> {
        Ok(Self {
            name_or_path: NameOrPath::Path(output_file.path),
            digest: output_file
                .digest
                .err_tip(|| "Expected digest to exist on OutputFile")?
                .try_into()?,
            is_executable: output_file.is_executable,
        })
    }
}

//TODO: Make this TryFrom.
impl From<FileInfo> for OutputFile {
    fn from(val: FileInfo) -> Self {
        let NameOrPath::Path(path) = val.name_or_path else {
            panic!("Cannot return a FileInfo that uses a NameOrPath::Name(), it must be a NameOrPath::Path()");
        };
        Self {
            path,
            digest: Some((&val.digest).into()),
            is_executable: val.is_executable,
            contents: Bytes::default(),
            node_properties: Option::default(), // Not supported.
        }
    }
}

/// Represents an individual symlink file and associated metadata.
/// This struct must be 100% compatible with `SymlinkNode` and `OutputSymlink`.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct SymlinkInfo {
    pub name_or_path: NameOrPath,
    pub target: String,
}

impl TryFrom<SymlinkNode> for SymlinkInfo {
    type Error = Error;

    fn try_from(symlink_node: SymlinkNode) -> Result<Self, Error> {
        Ok(Self {
            name_or_path: NameOrPath::Name(symlink_node.name),
            target: symlink_node.target,
        })
    }
}

// TODO: Make this TryFrom.
impl From<SymlinkInfo> for SymlinkNode {
    fn from(val: SymlinkInfo) -> Self {
        let NameOrPath::Name(name) = val.name_or_path else {
            panic!("Cannot return a SymlinkInfo that uses a NameOrPath::Path(), it must be a NameOrPath::Name()");
        };
        Self {
            name,
            target: val.target,
            node_properties: Option::default(), // Not supported.
        }
    }
}

impl TryFrom<OutputSymlink> for SymlinkInfo {
    type Error = Error;

    fn try_from(output_symlink: OutputSymlink) -> Result<Self, Error> {
        Ok(Self {
            name_or_path: NameOrPath::Path(output_symlink.path),
            target: output_symlink.target,
        })
    }
}

// TODO: Make this TryFrom.
impl From<SymlinkInfo> for OutputSymlink {
    fn from(val: SymlinkInfo) -> Self {
        let NameOrPath::Path(path) = val.name_or_path else {
            panic!("Cannot return a SymlinkInfo that uses a NameOrPath::Path(), it must be a NameOrPath::Name()");
        };
        Self {
            path,
            target: val.target,
            node_properties: Option::default(), // Not supported.
        }
    }
}

/// Represents an individual directory file and associated metadata.
/// This struct must be 100% compatible with `SymlinkNode` and `OutputSymlink`.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct DirectoryInfo {
    pub path: String,
    pub tree_digest: DigestInfo,
}

impl TryFrom<OutputDirectory> for DirectoryInfo {
    type Error = Error;

    fn try_from(output_directory: OutputDirectory) -> Result<Self, Error> {
        Ok(Self {
            path: output_directory.path,
            tree_digest: output_directory
                .tree_digest
                .err_tip(|| "Expected tree_digest to exist in OutputDirectory")?
                .try_into()?,
        })
    }
}

impl From<DirectoryInfo> for OutputDirectory {
    fn from(val: DirectoryInfo) -> Self {
        Self {
            path: val.path,
            tree_digest: Some(val.tree_digest.into()),
            is_topologically_sorted: false,
        }
    }
}

/// Represents the metadata associated with the execution result.
/// This struct must be 100% compatible with `ExecutedActionMetadata`.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionMetadata {
    pub worker: String,
    pub queued_timestamp: SystemTime,
    pub worker_start_timestamp: SystemTime,
    pub worker_completed_timestamp: SystemTime,
    pub input_fetch_start_timestamp: SystemTime,
    pub input_fetch_completed_timestamp: SystemTime,
    pub execution_start_timestamp: SystemTime,
    pub execution_completed_timestamp: SystemTime,
    pub output_upload_start_timestamp: SystemTime,
    pub output_upload_completed_timestamp: SystemTime,
}

impl Default for ExecutionMetadata {
    fn default() -> Self {
        Self {
            worker: "".to_string(),
            queued_timestamp: SystemTime::UNIX_EPOCH,
            worker_start_timestamp: SystemTime::UNIX_EPOCH,
            worker_completed_timestamp: SystemTime::UNIX_EPOCH,
            input_fetch_start_timestamp: SystemTime::UNIX_EPOCH,
            input_fetch_completed_timestamp: SystemTime::UNIX_EPOCH,
            execution_start_timestamp: SystemTime::UNIX_EPOCH,
            execution_completed_timestamp: SystemTime::UNIX_EPOCH,
            output_upload_start_timestamp: SystemTime::UNIX_EPOCH,
            output_upload_completed_timestamp: SystemTime::UNIX_EPOCH,
        }
    }
}

impl From<ExecutionMetadata> for ExecutedActionMetadata {
    fn from(val: ExecutionMetadata) -> Self {
        Self {
            worker: val.worker,
            queued_timestamp: Some(val.queued_timestamp.into()),
            worker_start_timestamp: Some(val.worker_start_timestamp.into()),
            worker_completed_timestamp: Some(val.worker_completed_timestamp.into()),
            input_fetch_start_timestamp: Some(val.input_fetch_start_timestamp.into()),
            input_fetch_completed_timestamp: Some(val.input_fetch_completed_timestamp.into()),
            execution_start_timestamp: Some(val.execution_start_timestamp.into()),
            execution_completed_timestamp: Some(val.execution_completed_timestamp.into()),
            output_upload_start_timestamp: Some(val.output_upload_start_timestamp.into()),
            output_upload_completed_timestamp: Some(val.output_upload_completed_timestamp.into()),
            virtual_execution_duration: val
                .execution_completed_timestamp
                .duration_since(val.execution_start_timestamp)
                .ok()
                .and_then(|duration| prost_types::Duration::try_from(duration).ok()),
            auxiliary_metadata: Vec::default(),
        }
    }
}

impl TryFrom<ExecutedActionMetadata> for ExecutionMetadata {
    type Error = Error;

    fn try_from(eam: ExecutedActionMetadata) -> Result<Self, Error> {
        Ok(Self {
            worker: eam.worker,
            queued_timestamp: eam
                .queued_timestamp
                .err_tip(|| "Expected queued_timestamp to exist in ExecutedActionMetadata")?
                .try_into()?,
            worker_start_timestamp: eam
                .worker_start_timestamp
                .err_tip(|| "Expected worker_start_timestamp to exist in ExecutedActionMetadata")?
                .try_into()?,
            worker_completed_timestamp: eam
                .worker_completed_timestamp
                .err_tip(|| {
                    "Expected worker_completed_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            input_fetch_start_timestamp: eam
                .input_fetch_start_timestamp
                .err_tip(|| {
                    "Expected input_fetch_start_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            input_fetch_completed_timestamp: eam
                .input_fetch_completed_timestamp
                .err_tip(|| {
                    "Expected input_fetch_completed_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            execution_start_timestamp: eam
                .execution_start_timestamp
                .err_tip(|| {
                    "Expected execution_start_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            execution_completed_timestamp: eam
                .execution_completed_timestamp
                .err_tip(|| {
                    "Expected execution_completed_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            output_upload_start_timestamp: eam
                .output_upload_start_timestamp
                .err_tip(|| {
                    "Expected output_upload_start_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
            output_upload_completed_timestamp: eam
                .output_upload_completed_timestamp
                .err_tip(|| {
                    "Expected output_upload_completed_timestamp to exist in ExecutedActionMetadata"
                })?
                .try_into()?,
        })
    }
}

/// Represents the results of an execution.
/// This struct must be 100% compatible with `ActionResult` in `remote_execution.proto`.
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct ActionResult {
    pub output_files: Vec<FileInfo>,
    pub output_folders: Vec<DirectoryInfo>,
    pub output_directory_symlinks: Vec<SymlinkInfo>,
    pub output_file_symlinks: Vec<SymlinkInfo>,
    pub exit_code: i32,
    pub stdout_digest: DigestInfo,
    pub stderr_digest: DigestInfo,
    pub execution_metadata: ExecutionMetadata,
    pub server_logs: HashMap<String, DigestInfo>,
    pub error: Option<Error>,
    pub message: String,
}

impl Default for ActionResult {
    fn default() -> Self {
        ActionResult {
            output_files: Default::default(),
            output_folders: Default::default(),
            output_directory_symlinks: Default::default(),
            output_file_symlinks: Default::default(),
            exit_code: INTERNAL_ERROR_EXIT_CODE,
            stdout_digest: DigestInfo::new([0u8; 32], 0),
            stderr_digest: DigestInfo::new([0u8; 32], 0),
            execution_metadata: ExecutionMetadata {
                worker: "".to_string(),
                queued_timestamp: SystemTime::UNIX_EPOCH,
                worker_start_timestamp: SystemTime::UNIX_EPOCH,
                worker_completed_timestamp: SystemTime::UNIX_EPOCH,
                input_fetch_start_timestamp: SystemTime::UNIX_EPOCH,
                input_fetch_completed_timestamp: SystemTime::UNIX_EPOCH,
                execution_start_timestamp: SystemTime::UNIX_EPOCH,
                execution_completed_timestamp: SystemTime::UNIX_EPOCH,
                output_upload_start_timestamp: SystemTime::UNIX_EPOCH,
                output_upload_completed_timestamp: SystemTime::UNIX_EPOCH,
            },
            server_logs: Default::default(),
            error: None,
            message: String::new(),
        }
    }
}

// TODO(allada) Remove the need for clippy argument by making the ActionResult and ProtoActionResult
// a Box.
/// The execution status/stage. This should match `ExecutionStage::Value` in `remote_execution.proto`.
#[derive(PartialEq, Debug, Clone)]
#[allow(clippy::large_enum_variant)]
pub enum ActionStage {
    /// Stage is unknown.
    Unknown,
    /// Checking the cache to see if action exists.
    CacheCheck,
    /// Action has been accepted and waiting for worker to take it.
    Queued,
    // TODO(allada) We need a way to know if the job was sent to a worker, but hasn't begun
    // execution yet.
    /// Worker is executing the action.
    Executing,
    /// Worker completed the work with result.
    Completed(ActionResult),
    /// Result was found from cache, don't decode the proto just to re-encode it.
    CompletedFromCache(ProtoActionResult),
}

impl Serialize for ActionStage {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let s = SerdeActionStage::try_from(self.clone()).map_err(S::Error::custom)?;
        s.serialize(serializer)
    }
}

#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
#[allow(clippy::large_enum_variant)]
enum SerdeActionStage {
    /// Stage is unknown.
    Unknown,
    /// Checking the cache to see if action exists.
    CacheCheck,
    /// Action has been accepted and waiting for worker to take it.
    Queued,
    // TODO(allada) We need a way to know if the job was sent to a worker, but hasn't begun
    // execution yet.
    /// Worker is executing the action.
    Executing,
    /// Worker completed the work with result.
    Completed(ActionResult),
}

impl From<SerdeActionStage> for ActionStage {
    fn from(value: SerdeActionStage) -> ActionStage {
        match value {
            SerdeActionStage::Unknown => ActionStage::Unknown,
            SerdeActionStage::CacheCheck => ActionStage::CacheCheck,
            SerdeActionStage::Queued => ActionStage::Queued,
            SerdeActionStage::Executing => ActionStage::Executing,
            SerdeActionStage::Completed(res) => ActionStage::Completed(res),
        }
    }
}

impl TryFrom<ActionStage> for SerdeActionStage {
    type Error = Error;
    fn try_from(value: ActionStage) -> Result<SerdeActionStage, Error> {
        match value {
            ActionStage::CompletedFromCache(proto) => {
                let action_result = ActionResult::try_from(proto)?;
                Ok(SerdeActionStage::Completed(action_result))
            }
            ActionStage::Unknown => Ok(SerdeActionStage::Unknown),
            ActionStage::CacheCheck => Ok(SerdeActionStage::CacheCheck),
            ActionStage::Queued => Ok(SerdeActionStage::Queued),
            ActionStage::Executing => Ok(SerdeActionStage::Executing),
            ActionStage::Completed(res) => Ok(SerdeActionStage::Completed(res)),
        }
    }
}

impl<'de> Deserialize<'de> for ActionStage {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        SerdeActionStage::deserialize(deserializer).map(|v| v.into())
    }
}

impl ActionStage {
    pub const fn has_action_result(&self) -> bool {
        match self {
            Self::Unknown | Self::CacheCheck | Self::Queued | Self::Executing => false,
            Self::Completed(_) | Self::CompletedFromCache(_) => true,
        }
    }

    /// Returns true if the worker considers the action done and no longer needs to be tracked.
    // Note: This function is separate from `has_action_result()` to not mix the concept of
    //       "finished" with "has a result".
    pub const fn is_finished(&self) -> bool {
        self.has_action_result()
    }

    /// Returns if the stage enum is the same as the other stage enum, but
    /// does not compare the values of the enum.
    pub const fn is_same_stage(&self, other: &Self) -> bool {
        matches!(
            (self, other),
            (Self::Unknown, Self::Unknown)
                | (Self::CacheCheck, Self::CacheCheck)
                | (Self::Queued, Self::Queued)
                | (Self::Executing, Self::Executing)
                | (Self::Completed(_), Self::Completed(_))
                | (Self::CompletedFromCache(_), Self::CompletedFromCache(_))
        )
    }
}

impl MetricsComponent for ActionStage {
    fn gather_metrics(&self, c: &mut CollectorState) {
        let (stage, maybe_exit_code) = match self {
            ActionStage::Unknown => ("Unknown", None),
            ActionStage::CacheCheck => ("CacheCheck", None),
            ActionStage::Queued => ("Queued", None),
            ActionStage::Executing => ("Executing", None),
            ActionStage::Completed(action_result) => ("Completed", Some(action_result.exit_code)),
            ActionStage::CompletedFromCache(proto_action_result) => {
                ("CompletedFromCache", Some(proto_action_result.exit_code))
            }
        };
        c.publish("stage", &stage.to_string(), "The state of the action.");
        if let Some(exit_code) = maybe_exit_code {
            c.publish("exit_code", &exit_code, "The exit code of the action.");
        }
    }
}

impl From<&ActionStage> for execution_stage::Value {
    fn from(val: &ActionStage) -> Self {
        match val {
            ActionStage::Unknown => Self::Unknown,
            ActionStage::CacheCheck => Self::CacheCheck,
            ActionStage::Queued => Self::Queued,
            ActionStage::Executing => Self::Executing,
            ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => Self::Completed,
        }
    }
}

pub fn to_execute_response(action_result: ActionResult) -> ExecuteResponse {
    fn logs_from(server_logs: HashMap<String, DigestInfo>) -> HashMap<String, LogFile> {
        let mut logs = HashMap::with_capacity(server_logs.len());
        for (k, v) in server_logs {
            logs.insert(
                k.clone(),
                LogFile {
                    digest: Some(v.into()),
                    human_readable: false,
                },
            );
        }
        logs
    }

    let status = Some(
        action_result
            .error
            .clone()
            .map_or_else(Status::default, |v| v.into()),
    );
    let message = action_result.message.clone();
    ExecuteResponse {
        server_logs: logs_from(action_result.server_logs.clone()),
        result: Some(action_result.into()),
        cached_result: false,
        status,
        message,
    }
}

impl From<ActionStage> for ExecuteResponse {
    fn from(val: ActionStage) -> Self {
        match val {
            // We don't have an execute response if we don't have the results. It is defined
            // behavior to return an empty proto struct.
            ActionStage::Unknown
            | ActionStage::CacheCheck
            | ActionStage::Queued
            | ActionStage::Executing => Self::default(),
            ActionStage::Completed(action_result) => to_execute_response(action_result),
            // Handled separately as there are no server logs and the action
            // result is already in Proto format.
            ActionStage::CompletedFromCache(proto_action_result) => Self {
                server_logs: HashMap::new(),
                result: Some(proto_action_result),
                cached_result: true,
                status: Some(Status::default()),
                message: String::new(), // Will be populated later if applicable.
            },
        }
    }
}

impl From<ActionResult> for ProtoActionResult {
    fn from(val: ActionResult) -> Self {
        let mut output_symlinks = Vec::with_capacity(
            val.output_file_symlinks.len() + val.output_directory_symlinks.len(),
        );
        output_symlinks.extend_from_slice(val.output_file_symlinks.as_slice());
        output_symlinks.extend_from_slice(val.output_directory_symlinks.as_slice());

        Self {
            output_files: val.output_files.into_iter().map(Into::into).collect(),
            output_file_symlinks: val
                .output_file_symlinks
                .into_iter()
                .map(Into::into)
                .collect(),
            output_symlinks: output_symlinks.into_iter().map(Into::into).collect(),
            output_directories: val.output_folders.into_iter().map(Into::into).collect(),
            output_directory_symlinks: val
                .output_directory_symlinks
                .into_iter()
                .map(Into::into)
                .collect(),
            exit_code: val.exit_code,
            stdout_raw: Bytes::default(),
            stdout_digest: Some(val.stdout_digest.into()),
            stderr_raw: Bytes::default(),
            stderr_digest: Some(val.stderr_digest.into()),
            execution_metadata: Some(val.execution_metadata.into()),
        }
    }
}

impl TryFrom<ProtoActionResult> for ActionResult {
    type Error = Error;

    fn try_from(val: ProtoActionResult) -> Result<Self, Error> {
        let output_file_symlinks = val
            .output_file_symlinks
            .into_iter()
            .map(|output_symlink| {
                SymlinkInfo::try_from(output_symlink)
                    .err_tip(|| "Output File Symlinks could not be converted to SymlinkInfo")
            })
            .collect::<Result<Vec<_>, _>>()?;

        let output_directory_symlinks = val
            .output_directory_symlinks
            .into_iter()
            .map(|output_symlink| {
                SymlinkInfo::try_from(output_symlink)
                    .err_tip(|| "Output File Symlinks could not be converted to SymlinkInfo")
            })
            .collect::<Result<Vec<_>, _>>()?;

        let output_files = val
            .output_files
            .into_iter()
            .map(|output_file| {
                output_file
                    .try_into()
                    .err_tip(|| "Output File could not be converted")
            })
            .collect::<Result<Vec<_>, _>>()?;

        let output_folders = val
            .output_directories
            .into_iter()
            .map(|output_directory| {
                output_directory
                    .try_into()
                    .err_tip(|| "Output File could not be converted")
            })
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self {
            output_files,
            output_folders,
            output_file_symlinks,
            output_directory_symlinks,
            exit_code: val.exit_code,
            stdout_digest: val
                .stdout_digest
                .err_tip(|| "Expected stdout_digest to be set on ExecuteResponse msg")?
                .try_into()?,
            stderr_digest: val
                .stderr_digest
                .err_tip(|| "Expected stderr_digest to be set on ExecuteResponse msg")?
                .try_into()?,
            execution_metadata: val
                .execution_metadata
                .err_tip(|| "Expected execution_metadata to be set on ExecuteResponse msg")?
                .try_into()?,
            server_logs: Default::default(),
            error: None,
            message: String::new(),
        })
    }
}

impl TryFrom<ExecuteResponse> for ActionStage {
    type Error = Error;

    fn try_from(execute_response: ExecuteResponse) -> Result<Self, Error> {
        let proto_action_result = execute_response
            .result
            .err_tip(|| "Expected result to be set on ExecuteResponse msg")?;
        let action_result = ActionResult {
            output_files: proto_action_result
                .output_files
                .try_map(TryInto::try_into)?,
            output_directory_symlinks: proto_action_result
                .output_directory_symlinks
                .try_map(TryInto::try_into)?,
            output_file_symlinks: proto_action_result
                .output_file_symlinks
                .try_map(TryInto::try_into)?,
            output_folders: proto_action_result
                .output_directories
                .try_map(TryInto::try_into)?,
            exit_code: proto_action_result.exit_code,

            stdout_digest: proto_action_result
                .stdout_digest
                .err_tip(|| "Expected stdout_digest to be set on ExecuteResponse msg")?
                .try_into()?,
            stderr_digest: proto_action_result
                .stderr_digest
                .err_tip(|| "Expected stderr_digest to be set on ExecuteResponse msg")?
                .try_into()?,
            execution_metadata: proto_action_result
                .execution_metadata
                .err_tip(|| "Expected execution_metadata to be set on ExecuteResponse msg")?
                .try_into()?,
            server_logs: execute_response.server_logs.try_map(|v| {
                v.digest
                    .err_tip(|| "Expected digest to be set on LogFile msg")?
                    .try_into()
            })?,
            error: execute_response.status.clone().and_then(|v| {
                if v.code == 0 {
                    None
                } else {
                    Some(v.into())
                }
            }),
            message: execute_response.message,
        };

        if execute_response.cached_result {
            return Ok(Self::CompletedFromCache(action_result.into()));
        }
        Ok(Self::Completed(action_result))
    }
}

// TODO: Should be able to remove this after tokio-rs/prost#299
trait TypeUrl: Message {
    const TYPE_URL: &'static str;
}

impl TypeUrl for ExecuteResponse {
    const TYPE_URL: &'static str =
        "type.googleapis.com/build.bazel.remote.execution.v2.ExecuteResponse";
}

impl TypeUrl for ExecuteOperationMetadata {
    const TYPE_URL: &'static str =
        "type.googleapis.com/build.bazel.remote.execution.v2.ExecuteOperationMetadata";
}

fn from_any<T>(message: &Any) -> Result<T, Error>
where
    T: TypeUrl + Default,
{
    error_if!(
        message.type_url != T::TYPE_URL,
        "Incorrect type when decoding Any. {} != {}",
        message.type_url,
        T::TYPE_URL.to_string()
    );
    Ok(T::decode(message.value.as_slice())?)
}

fn to_any<T>(message: &T) -> Any
where
    T: TypeUrl,
{
    Any {
        type_url: T::TYPE_URL.to_string(),
        value: message.encode_to_vec(),
    }
}

/// Current state of the action.
/// This must be 100% compatible with `Operation` in `google/longrunning/operations.proto`.
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct ActionState {
    pub stage: ActionStage,
    pub id: OperationId,
}

impl ActionState {
    pub fn try_from_operation(
        operation: Operation,
        operation_id: OperationId,
    ) -> Result<Self, Error> {
        let metadata = from_any::<ExecuteOperationMetadata>(
            &operation
                .metadata
                .err_tip(|| "No metadata in upstream operation")?,
        )
        .err_tip(|| "Could not decode metadata in upstream operation")?;

        let stage = match execution_stage::Value::try_from(metadata.stage).err_tip(|| {
            format!(
                "Could not convert {} to execution_stage::Value",
                metadata.stage
            )
        })? {
            execution_stage::Value::Unknown => ActionStage::Unknown,
            execution_stage::Value::CacheCheck => ActionStage::CacheCheck,
            execution_stage::Value::Queued => ActionStage::Queued,
            execution_stage::Value::Executing => ActionStage::Executing,
            execution_stage::Value::Completed => {
                let execute_response = operation
                    .result
                    .err_tip(|| "No result data for completed upstream action")?;
                match execute_response {
                    LongRunningResult::Error(error) => ActionStage::Completed(ActionResult {
                        error: Some(error.into()),
                        ..ActionResult::default()
                    }),
                    LongRunningResult::Response(response) => {
                        // Could be Completed, CompletedFromCache or Error.
                        from_any::<ExecuteResponse>(&response)
                            .err_tip(|| {
                                "Could not decode result structure for completed upstream action"
                            })?
                            .try_into()?
                    }
                }
            }
        };

        Ok(Self {
            id: operation_id,
            stage,
        })
    }

    pub fn as_operation(&self, client_operation_id: ClientOperationId) -> Operation {
        let stage = Into::<execution_stage::Value>::into(&self.stage) as i32;
        let name = client_operation_id.into_string();

        let result = if self.stage.has_action_result() {
            let execute_response: ExecuteResponse = self.stage.clone().into();
            Some(LongRunningResult::Response(to_any(&execute_response)))
        } else {
            None
        };
        let digest = Some(self.id.unique_qualifier.digest().into());

        let metadata = ExecuteOperationMetadata {
            stage,
            action_digest: digest,
            // TODO(blaise.bruer) We should support stderr/stdout streaming.
            stdout_stream_name: String::default(),
            stderr_stream_name: String::default(),
            partial_execution_metadata: None,
        };

        Operation {
            name,
            metadata: Some(to_any(&metadata)),
            done: result.is_some(),
            result,
        }
    }
}

impl MetricsComponent for ActionState {
    fn gather_metrics(&self, c: &mut CollectorState) {
        c.publish("stage", &self.stage, "");
    }
}