//! 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.
///
/// The default output selection is
///
/// ```json
///   {
///    "*": {
///      "*": [
///        "abi",
///        "evm.bytecode",
///        "evm.deployedBytecode",
///        "evm.methodIdentifiers"
///      ],
///      "": [
///        "ast"
///      ]
///    }
///  }
/// ```
#[derive(Debug, Clone, Eq, PartialEq, Default, 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()]),
                ("".to_string(), vec!["*".to_string()]),
            ]),
        )])
        .into()
    }

    /// Default output selection for compiler output:
    ///
    /// `{ "*": { "*": [ "*" ], "": [
    /// "abi","evm.bytecode","evm.deployedBytecode","evm.methodIdentifiers"] } }`
    ///
    /// Which enables it for all files and all their contracts ("*" wildcard)
    pub fn default_output_selection() -> Self {
        BTreeMap::from([("*".to_string(), Self::default_file_output_selection())]).into()
    }

    /// Default output selection for a single file:
    ///
    /// `{ "*": [ "*" ], "": [
    /// "abi","evm.bytecode","evm.deployedBytecode","evm.methodIdentifiers"] }`
    ///
    /// Which enables it for all the contracts in the file ("*" wildcard)
    pub fn default_file_output_selection() -> FileOutputSelection {
        BTreeMap::from([(
            "*".to_string(),
            vec![
                "abi".to_string(),
                "evm.bytecode".to_string(),
                "evm.deployedBytecode".to_string(),
                "evm.methodIdentifiers".to_string(),
            ],
        )])
    }

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

// 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 {
        OutputSelection(s)
    }
}

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

impl ContractOutputSelection {
    /// Returns the basic set of contract level settings that should be included in the `Contract`
    /// that solc emits:
    ///    - "abi"
    ///    - "evm.bytecode"
    ///    - "evm.deployedBytecode"
    ///    - "evm.methodIdentifiers"
    pub fn basic() -> Vec<ContractOutputSelection> {
        vec![
            ContractOutputSelection::Abi,
            BytecodeOutputSelection::All.into(),
            DeployedBytecodeOutputSelection::All.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 {
            ContractOutputSelection::Abi => f.write_str("abi"),
            ContractOutputSelection::DevDoc => f.write_str("devdoc"),
            ContractOutputSelection::UserDoc => f.write_str("userdoc"),
            ContractOutputSelection::Metadata => f.write_str("metadata"),
            ContractOutputSelection::Ir => f.write_str("ir"),
            ContractOutputSelection::IrOptimized => f.write_str("irOptimized"),
            ContractOutputSelection::StorageLayout => f.write_str("storageLayout"),
            ContractOutputSelection::Evm(e) => e.fmt(f),
            ContractOutputSelection::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(ContractOutputSelection::Abi),
            "devdoc" => Ok(ContractOutputSelection::DevDoc),
            "userdoc" => Ok(ContractOutputSelection::UserDoc),
            "metadata" => Ok(ContractOutputSelection::Metadata),
            "ir" => Ok(ContractOutputSelection::Ir),
            "ir-optimized" | "irOptimized" | "iroptimized" => {
                Ok(ContractOutputSelection::IrOptimized)
            }
            "storage-layout" | "storagelayout" | "storageLayout" => {
                Ok(ContractOutputSelection::StorageLayout)
            }
            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 {
        ContractOutputSelection::Evm(evm.into())
    }
}

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

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

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

impl From<DeployedBytecodeOutputSelection> for EvmOutputSelection {
    fn from(b: DeployedBytecodeOutputSelection) -> Self {
        EvmOutputSelection::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 {
            EvmOutputSelection::All => f.write_str("evm"),
            EvmOutputSelection::Assembly => f.write_str("evm.assembly"),
            EvmOutputSelection::LegacyAssembly => f.write_str("evm.legacyAssembly"),
            EvmOutputSelection::MethodIdentifiers => f.write_str("evm.methodIdentifiers"),
            EvmOutputSelection::GasEstimates => f.write_str("evm.gasEstimates"),
            EvmOutputSelection::ByteCode(b) => b.fmt(f),
            EvmOutputSelection::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(EvmOutputSelection::All),
            "asm" | "evm.assembly" => Ok(EvmOutputSelection::Assembly),
            "evm.legacyAssembly" => Ok(EvmOutputSelection::LegacyAssembly),
            "methodidentifiers" | "evm.methodIdentifiers" | "evm.methodidentifiers" => {
                Ok(EvmOutputSelection::MethodIdentifiers)
            }
            "gas" | "evm.gasEstimates" | "evm.gasestimates" => Ok(EvmOutputSelection::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, Hash, PartialOrd, Ord)]
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 {
            BytecodeOutputSelection::All => f.write_str("evm.bytecode"),
            BytecodeOutputSelection::FunctionDebugData => {
                f.write_str("evm.bytecode.functionDebugData")
            }
            BytecodeOutputSelection::Object => f.write_str("evm.bytecode.object"),
            BytecodeOutputSelection::Opcodes => f.write_str("evm.bytecode.opcodes"),
            BytecodeOutputSelection::SourceMap => f.write_str("evm.bytecode.sourceMap"),
            BytecodeOutputSelection::LinkReferences => f.write_str("evm.bytecode.linkReferences"),
            BytecodeOutputSelection::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(BytecodeOutputSelection::All),
            "evm.bytecode.functionDebugData" => Ok(BytecodeOutputSelection::FunctionDebugData),
            "evm.bytecode.object" => Ok(BytecodeOutputSelection::Object),
            "evm.bytecode.opcodes" => Ok(BytecodeOutputSelection::Opcodes),
            "evm.bytecode.sourceMap" => Ok(BytecodeOutputSelection::SourceMap),
            "evm.bytecode.linkReferences" => Ok(BytecodeOutputSelection::LinkReferences),
            "evm.bytecode.generatedSources" => Ok(BytecodeOutputSelection::GeneratedSources),
            s => Err(format!("Invalid bytecode selection: {s}")),
        }
    }
}

/// Contract level output selection for `evm.deployedBytecode`
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
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 {
            DeployedBytecodeOutputSelection::All => f.write_str("evm.deployedBytecode"),
            DeployedBytecodeOutputSelection::FunctionDebugData => {
                f.write_str("evm.deployedBytecode.functionDebugData")
            }
            DeployedBytecodeOutputSelection::Object => f.write_str("evm.deployedBytecode.object"),
            DeployedBytecodeOutputSelection::Opcodes => f.write_str("evm.deployedBytecode.opcodes"),
            DeployedBytecodeOutputSelection::SourceMap => {
                f.write_str("evm.deployedBytecode.sourceMap")
            }
            DeployedBytecodeOutputSelection::LinkReferences => {
                f.write_str("evm.deployedBytecode.linkReferences")
            }
            DeployedBytecodeOutputSelection::GeneratedSources => {
                f.write_str("evm.deployedBytecode.generatedSources")
            }
            DeployedBytecodeOutputSelection::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(DeployedBytecodeOutputSelection::All),
            "evm.deployedBytecode.functionDebugData" => {
                Ok(DeployedBytecodeOutputSelection::FunctionDebugData)
            }
            "evm.deployedBytecode.object" => Ok(DeployedBytecodeOutputSelection::Object),
            "evm.deployedBytecode.opcodes" => Ok(DeployedBytecodeOutputSelection::Opcodes),
            "evm.deployedBytecode.sourceMap" => Ok(DeployedBytecodeOutputSelection::SourceMap),
            "evm.deployedBytecode.linkReferences" => {
                Ok(DeployedBytecodeOutputSelection::LinkReferences)
            }
            "evm.deployedBytecode.generatedSources" => {
                Ok(DeployedBytecodeOutputSelection::GeneratedSources)
            }
            "evm.deployedBytecode.immutableReferences" => {
                Ok(DeployedBytecodeOutputSelection::ImmutableReferences)
            }
            s => Err(format!("Invalid deployedBytecode selection: {s}")),
        }
    }
}

/// Contract level output selection for `evm.ewasm`
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
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 {
            EwasmOutputSelection::All => f.write_str("ewasm"),
            EwasmOutputSelection::Wast => f.write_str("ewasm.wast"),
            EwasmOutputSelection::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(EwasmOutputSelection::All),
            "ewasm.wast" => Ok(EwasmOutputSelection::Wast),
            "ewasm.wasm" => Ok(EwasmOutputSelection::Wasm),
            s => Err(format!("Invalid ewasm selection: {s}")),
        }
    }
}

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

    #[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 deployed_bytecode_from_str() {
        assert_eq!(
            DeployedBytecodeOutputSelection::from_str("evm.deployedBytecode.immutableReferences")
                .unwrap(),
            DeployedBytecodeOutputSelection::ImmutableReferences
        )
    }
}
