//! Bindings for standard json output selection.

use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer};
use std::{collections::BTreeMap, fmt, str::FromStr};

/// Represents the desired outputs based on a File `(file -> (contract -> [outputs]))`
pub type FileOutputSelection = BTreeMap<String, Vec<String>>;

/// Represents the selected output of files and contracts.
///
/// The first level key is the file name and the second level key is the
/// contract name. An empty contract name is used for outputs that are
/// not tied to a contract but to the whole source file like the AST.
/// A star as contract name refers to all contracts in the file.
/// Similarly, a star as a file name matches all files.
/// To select all outputs the compiler can possibly generate, use
/// "outputSelection: { "*": { "*": [ "*" ], "": [ "*" ] } }"
/// but note that this might slow down the compilation process needlessly.
///
/// The available output types are as follows:
///
/// File level (needs empty string as contract name):
///   ast - AST of all source files
///
/// Contract level (needs the contract name or "*"):
///   abi - ABI
///   devdoc - Developer documentation (natspec)
///   userdoc - User documentation (natspec)
///   metadata - Metadata
///   ir - Yul intermediate representation of the code before optimization
///   irOptimized - Intermediate representation after optimization
///   storageLayout - Slots, offsets and types of the contract's state
///     variables.
///   evm.assembly - New assembly format
///   evm.legacyAssembly - Old-style assembly format in JSON
///   evm.bytecode.functionDebugData - Debugging information at function level
///   evm.bytecode.object - Bytecode object
///   evm.bytecode.opcodes - Opcodes list
///   evm.bytecode.sourceMap - Source mapping (useful for debugging)
///   evm.bytecode.linkReferences - Link references (if unlinked object)
///   evm.bytecode.generatedSources - Sources generated by the compiler
///   evm.deployedBytecode* - Deployed bytecode (has all the options that
///     evm.bytecode has)
///   evm.deployedBytecode.immutableReferences - Map from AST ids to
///     bytecode ranges that reference immutables
///   evm.methodIdentifiers - The list of function hashes
///   evm.gasEstimates - Function gas estimates
///   ewasm.wast - Ewasm in WebAssembly S-expressions format
///   ewasm.wasm - Ewasm in WebAssembly binary format
///
/// Note that using a using `evm`, `evm.bytecode`, `ewasm`, etc. will select
/// every target part of that output. Additionally, `*` can be used as a
/// wildcard to request everything.
#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
#[serde(transparent)]
pub struct OutputSelection(pub BTreeMap<String, FileOutputSelection>);

impl OutputSelection {
    /// select all outputs the compiler can possibly generate, use
    /// `{ "*": { "*": [ "*" ], "": [ "*" ] } }`
    /// but note that this might slow down the compilation process needlessly.
    pub fn complete_output_selection() -> Self {
        BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([
                ("*".to_string(), vec!["*".to_string()]),
                (String::new(), vec!["*".to_string()]),
            ]),
        )])
        .into()
    }

    /// Default output selection.
    pub fn default_output_selection() -> Self {
        BTreeMap::from([("*".to_string(), Self::default_file_output_selection())]).into()
    }

    /// Default file output selection.
    ///
    /// Uses [`ContractOutputSelection::basic`].
    pub fn default_file_output_selection() -> FileOutputSelection {
        BTreeMap::from([(
            "*".to_string(),
            ContractOutputSelection::basic().iter().map(ToString::to_string).collect(),
        )])
    }

    /// Returns output selection configuration which enables the same provided outputs for each
    /// contract of each source.
    pub fn common_output_selection(outputs: impl IntoIterator<Item = String>) -> Self {
        BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([("*".to_string(), outputs.into_iter().collect())]),
        )])
        .into()
    }

    /// Returns an empty output selection which corresponds to an empty map `{}`
    pub fn empty_file_output_select() -> FileOutputSelection {
        Default::default()
    }

    /// Returns output selection which requests only AST for all sources.
    pub fn ast_output_selection() -> Self {
        BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([
                // Do not request any output for separate contracts
                ("*".to_string(), vec![]),
                // Request AST for all sources.
                (String::new(), vec!["ast".to_string()]),
            ]),
        )])
        .into()
    }

    /// Returns true if this output selection is a subset of the other output selection.
    // TODO: correctly process wildcard keys to reduce false negatives
    pub fn is_subset_of(&self, other: &Self) -> bool {
        self.0.iter().all(|(file, selection)| {
            other.0.get(file).is_some_and(|other_selection| {
                selection.iter().all(|(contract, outputs)| {
                    other_selection.get(contract).is_some_and(|other_outputs| {
                        outputs.iter().all(|output| other_outputs.contains(output))
                    })
                })
            })
        })
    }
}

// this will make sure that if the `FileOutputSelection` for a certain file is empty will be
// serializes as `"*" : []` because
// > Contract level (needs the contract name or "*") <https://docs.soliditylang.org/en/v0.8.13/using-the-compiler.html>
impl Serialize for OutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        struct EmptyFileOutput;

        impl Serialize for EmptyFileOutput {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: Serializer,
            {
                let mut map = serializer.serialize_map(Some(1))?;
                map.serialize_entry("*", &[] as &[String])?;
                map.end()
            }
        }

        let mut map = serializer.serialize_map(Some(self.0.len()))?;
        for (file, selection) in self.0.iter() {
            if selection.is_empty() {
                map.serialize_entry(file, &EmptyFileOutput {})?;
            } else {
                map.serialize_entry(file, selection)?;
            }
        }
        map.end()
    }
}

impl AsRef<BTreeMap<String, FileOutputSelection>> for OutputSelection {
    fn as_ref(&self) -> &BTreeMap<String, FileOutputSelection> {
        &self.0
    }
}

impl AsMut<BTreeMap<String, FileOutputSelection>> for OutputSelection {
    fn as_mut(&mut self) -> &mut BTreeMap<String, FileOutputSelection> {
        &mut self.0
    }
}

impl From<BTreeMap<String, FileOutputSelection>> for OutputSelection {
    fn from(s: BTreeMap<String, FileOutputSelection>) -> Self {
        Self(s)
    }
}

/// Contract level output selection
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum ContractOutputSelection {
    Abi,
    DevDoc,
    UserDoc,
    Metadata,
    Ir,
    IrOptimized,
    IrOptimizedAst,
    StorageLayout,
    TransientStorageLayout,
    Evm(EvmOutputSelection),
    Ewasm(EwasmOutputSelection),
}

impl ContractOutputSelection {
    /// Returns the basic set of contract level settings that should be included in the `Contract`
    /// that solc emits.
    ///
    /// These correspond to the fields in `CompactBytecode`, `CompactDeployedBytecode`, ABI, and
    /// method identifiers.
    pub fn basic() -> Vec<Self> {
        // We don't include all the `bytecode` fields because `generatedSources` is a massive JSON
        // object and is not used by Foundry.
        vec![
            Self::Abi,
            // The fields in `CompactBytecode`.
            BytecodeOutputSelection::Object.into(),
            BytecodeOutputSelection::SourceMap.into(),
            BytecodeOutputSelection::LinkReferences.into(),
            // The fields in `CompactDeployedBytecode`.
            DeployedBytecodeOutputSelection::Object.into(),
            DeployedBytecodeOutputSelection::SourceMap.into(),
            DeployedBytecodeOutputSelection::LinkReferences.into(),
            DeployedBytecodeOutputSelection::ImmutableReferences.into(),
            EvmOutputSelection::MethodIdentifiers.into(),
        ]
    }
}

impl Serialize for ContractOutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for ContractOutputSelection {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?.parse().map_err(serde::de::Error::custom)
    }
}

impl fmt::Display for ContractOutputSelection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Abi => f.write_str("abi"),
            Self::DevDoc => f.write_str("devdoc"),
            Self::UserDoc => f.write_str("userdoc"),
            Self::Metadata => f.write_str("metadata"),
            Self::Ir => f.write_str("ir"),
            Self::IrOptimized => f.write_str("irOptimized"),
            Self::IrOptimizedAst => f.write_str("irOptimizedAst"),
            Self::StorageLayout => f.write_str("storageLayout"),
            Self::TransientStorageLayout => f.write_str("transientStorageLayout"),
            Self::Evm(e) => e.fmt(f),
            Self::Ewasm(e) => e.fmt(f),
        }
    }
}

impl FromStr for ContractOutputSelection {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "abi" => Ok(Self::Abi),
            "devdoc" => Ok(Self::DevDoc),
            "userdoc" => Ok(Self::UserDoc),
            "metadata" => Ok(Self::Metadata),
            "ir" => Ok(Self::Ir),
            "ir-optimized" | "irOptimized" | "iroptimized" => Ok(Self::IrOptimized),
            "irOptimizedAst" | "ir-optimized-ast" | "iroptimizedast" => Ok(Self::IrOptimizedAst),
            "storage-layout" | "storagelayout" | "storageLayout" => Ok(Self::StorageLayout),
            "transient-storage-layout" | "transientstoragelayout" | "transientStorageLayout" => {
                Ok(Self::TransientStorageLayout)
            }
            s => EvmOutputSelection::from_str(s)
                .map(ContractOutputSelection::Evm)
                .or_else(|_| EwasmOutputSelection::from_str(s).map(ContractOutputSelection::Ewasm))
                .map_err(|_| format!("Invalid contract output selection: {s}")),
        }
    }
}

impl<T: Into<EvmOutputSelection>> From<T> for ContractOutputSelection {
    fn from(evm: T) -> Self {
        Self::Evm(evm.into())
    }
}

impl From<EwasmOutputSelection> for ContractOutputSelection {
    fn from(ewasm: EwasmOutputSelection) -> Self {
        Self::Ewasm(ewasm)
    }
}

/// Contract level output selection for `evm`
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum EvmOutputSelection {
    All,
    Assembly,
    LegacyAssembly,
    MethodIdentifiers,
    GasEstimates,
    ByteCode(BytecodeOutputSelection),
    DeployedByteCode(DeployedBytecodeOutputSelection),
}

impl From<BytecodeOutputSelection> for EvmOutputSelection {
    fn from(b: BytecodeOutputSelection) -> Self {
        Self::ByteCode(b)
    }
}

impl From<DeployedBytecodeOutputSelection> for EvmOutputSelection {
    fn from(b: DeployedBytecodeOutputSelection) -> Self {
        Self::DeployedByteCode(b)
    }
}

impl Serialize for EvmOutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for EvmOutputSelection {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?.parse().map_err(serde::de::Error::custom)
    }
}

impl fmt::Display for EvmOutputSelection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::All => f.write_str("evm"),
            Self::Assembly => f.write_str("evm.assembly"),
            Self::LegacyAssembly => f.write_str("evm.legacyAssembly"),
            Self::MethodIdentifiers => f.write_str("evm.methodIdentifiers"),
            Self::GasEstimates => f.write_str("evm.gasEstimates"),
            Self::ByteCode(b) => b.fmt(f),
            Self::DeployedByteCode(b) => b.fmt(f),
        }
    }
}

impl FromStr for EvmOutputSelection {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "evm" => Ok(Self::All),
            "asm" | "evm.assembly" => Ok(Self::Assembly),
            "legacyAssembly" | "evm.legacyAssembly" => Ok(Self::LegacyAssembly),
            "methodidentifiers" | "evm.methodIdentifiers" | "evm.methodidentifiers" => {
                Ok(Self::MethodIdentifiers)
            }
            "gas" | "evm.gasEstimates" | "evm.gasestimates" => Ok(Self::GasEstimates),
            s => BytecodeOutputSelection::from_str(s)
                .map(EvmOutputSelection::ByteCode)
                .or_else(|_| {
                    DeployedBytecodeOutputSelection::from_str(s)
                        .map(EvmOutputSelection::DeployedByteCode)
                })
                .map_err(|_| format!("Invalid evm selection: {s}")),
        }
    }
}

/// Contract level output selection for `evm.bytecode`
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum BytecodeOutputSelection {
    All,
    FunctionDebugData,
    Object,
    Opcodes,
    SourceMap,
    LinkReferences,
    GeneratedSources,
}

impl Serialize for BytecodeOutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for BytecodeOutputSelection {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?.parse().map_err(serde::de::Error::custom)
    }
}

impl fmt::Display for BytecodeOutputSelection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::All => f.write_str("evm.bytecode"),
            Self::FunctionDebugData => f.write_str("evm.bytecode.functionDebugData"),
            Self::Object => f.write_str("evm.bytecode.object"),
            Self::Opcodes => f.write_str("evm.bytecode.opcodes"),
            Self::SourceMap => f.write_str("evm.bytecode.sourceMap"),
            Self::LinkReferences => f.write_str("evm.bytecode.linkReferences"),
            Self::GeneratedSources => f.write_str("evm.bytecode.generatedSources"),
        }
    }
}

impl FromStr for BytecodeOutputSelection {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "evm.bytecode" => Ok(Self::All),
            "evm.bytecode.functionDebugData" => Ok(Self::FunctionDebugData),
            "code" | "bin" | "evm.bytecode.object" => Ok(Self::Object),
            "evm.bytecode.opcodes" => Ok(Self::Opcodes),
            "evm.bytecode.sourceMap" => Ok(Self::SourceMap),
            "evm.bytecode.linkReferences" => Ok(Self::LinkReferences),
            "evm.bytecode.generatedSources" => Ok(Self::GeneratedSources),
            s => Err(format!("Invalid bytecode selection: {s}")),
        }
    }
}

/// Contract level output selection for `evm.deployedBytecode`
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum DeployedBytecodeOutputSelection {
    All,
    FunctionDebugData,
    Object,
    Opcodes,
    SourceMap,
    LinkReferences,
    GeneratedSources,
    ImmutableReferences,
}

impl Serialize for DeployedBytecodeOutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for DeployedBytecodeOutputSelection {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?.parse().map_err(serde::de::Error::custom)
    }
}

impl fmt::Display for DeployedBytecodeOutputSelection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::All => f.write_str("evm.deployedBytecode"),
            Self::FunctionDebugData => f.write_str("evm.deployedBytecode.functionDebugData"),
            Self::Object => f.write_str("evm.deployedBytecode.object"),
            Self::Opcodes => f.write_str("evm.deployedBytecode.opcodes"),
            Self::SourceMap => f.write_str("evm.deployedBytecode.sourceMap"),
            Self::LinkReferences => f.write_str("evm.deployedBytecode.linkReferences"),
            Self::GeneratedSources => f.write_str("evm.deployedBytecode.generatedSources"),
            Self::ImmutableReferences => f.write_str("evm.deployedBytecode.immutableReferences"),
        }
    }
}

impl FromStr for DeployedBytecodeOutputSelection {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "evm.deployedBytecode" => Ok(Self::All),
            "evm.deployedBytecode.functionDebugData" => Ok(Self::FunctionDebugData),
            "deployed-code"
            | "deployed-bin"
            | "runtime-code"
            | "runtime-bin"
            | "evm.deployedBytecode.object" => Ok(Self::Object),
            "evm.deployedBytecode.opcodes" => Ok(Self::Opcodes),
            "evm.deployedBytecode.sourceMap" => Ok(Self::SourceMap),
            "evm.deployedBytecode.linkReferences" => Ok(Self::LinkReferences),
            "evm.deployedBytecode.generatedSources" => Ok(Self::GeneratedSources),
            "evm.deployedBytecode.immutableReferences" => Ok(Self::ImmutableReferences),
            s => Err(format!("Invalid deployedBytecode selection: {s}")),
        }
    }
}

/// Contract level output selection for `evm.ewasm`
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum EwasmOutputSelection {
    All,
    Wast,
    Wasm,
}

impl Serialize for EwasmOutputSelection {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for EwasmOutputSelection {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?.parse().map_err(serde::de::Error::custom)
    }
}

impl fmt::Display for EwasmOutputSelection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::All => f.write_str("ewasm"),
            Self::Wast => f.write_str("ewasm.wast"),
            Self::Wasm => f.write_str("ewasm.wasm"),
        }
    }
}

impl FromStr for EwasmOutputSelection {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "ewasm" => Ok(Self::All),
            "ewasm.wast" => Ok(Self::Wast),
            "ewasm.wasm" => Ok(Self::Wasm),
            s => Err(format!("Invalid ewasm selection: {s}")),
        }
    }
}

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

    #[test]
    fn outputselection_serde_works() {
        let mut output = BTreeMap::default();
        output.insert(
            "*".to_string(),
            vec![
                "abi".to_string(),
                "evm.bytecode".to_string(),
                "evm.deployedBytecode".to_string(),
                "evm.methodIdentifiers".to_string(),
            ],
        );

        let json = serde_json::to_string(&output).unwrap();
        let deserde_selection: BTreeMap<String, Vec<ContractOutputSelection>> =
            serde_json::from_str(&json).unwrap();

        assert_eq!(json, serde_json::to_string(&deserde_selection).unwrap());
    }

    #[test]
    fn empty_outputselection_serde_works() {
        let mut empty = OutputSelection::default();
        empty.0.insert("contract.sol".to_string(), OutputSelection::empty_file_output_select());
        let s = serde_json::to_string(&empty).unwrap();
        assert_eq!(s, r#"{"contract.sol":{"*":[]}}"#);
    }

    #[test]
    fn outputselection_subset_of() {
        let output_selection = OutputSelection::from(BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([(
                "*".to_string(),
                vec!["abi".to_string(), "evm.bytecode".to_string()],
            )]),
        )]));

        let output_selection_abi = OutputSelection::from(BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([("*".to_string(), vec!["abi".to_string()])]),
        )]));

        assert!(output_selection_abi.is_subset_of(&output_selection));
        assert!(!output_selection.is_subset_of(&output_selection_abi));

        let output_selection_empty = OutputSelection::from(BTreeMap::from([(
            "*".to_string(),
            BTreeMap::from([("*".to_string(), vec![])]),
        )]));

        assert!(output_selection_empty.is_subset_of(&output_selection));
        assert!(output_selection_empty.is_subset_of(&output_selection_abi));
        assert!(!output_selection.is_subset_of(&output_selection_empty));
        assert!(!output_selection_abi.is_subset_of(&output_selection_empty));

        let output_selection_specific = OutputSelection::from(BTreeMap::from([(
            "Contract.sol".to_string(),
            BTreeMap::from([(
                "Contract".to_string(),
                vec![
                    "abi".to_string(),
                    "evm.bytecode".to_string(),
                    "evm.deployedBytecode".to_string(),
                ],
            )]),
        )]));

        assert!(!output_selection_specific.is_subset_of(&output_selection));
    }

    #[test]
    fn deployed_bytecode_from_str() {
        assert_eq!(
            DeployedBytecodeOutputSelection::from_str("evm.deployedBytecode.immutableReferences")
                .unwrap(),
            DeployedBytecodeOutputSelection::ImmutableReferences
        )
    }
}
