use anyhow::{Result, bail};
use wasm_encoder::{IndirectNameMap, NameMap, NameSection, RawSection};
use wasmparser::{KnownCustom, Name, NameSectionReader, Parser, Payload::*};

/// Demangle Rust and C++ symbol names in the `name` section.
///
/// This command will detect a `name` section in a wasm executable and demangle
/// any Rust and C++ symbol names found within it. Tooling for debugging a wasm
/// module which otherwise uses the `name` section but doesn't run a demangler
/// will use the demangled names since the `name` section will be replaced.
#[derive(clap::Parser)]
#[clap(after_help = "\
Examples:

Suppose foo.wasm has the following textual representation:

(module
  (func $do_not_demangle_me)
  (func $_ZN4rustE)
  (func $_Z3food)
)

The second two functions are mangled Rust symbol names.

    # Demangle symbol names in foo.wasm and print the textual form of
    # the output to stdout.
    $ wasm-tools demangle -t foo.wasm
(module
  (type (;0;) (func))
  (func $do_not_demangle_me (;0;) (type 0))
  (func $rust (;1;) (type 0))
  (func $\"foo(double)\" (;2;) (type 0))
)

    # Demangle symbol names in foo.wasm and save the binary output
    # to foo-demangled.wasm.
    $ wasm-tools demangle foo.wasm -o foo-demangled.wasm

Exit status:
    0 on success,
    nonzero if the input file fails to parse.
")]
pub struct Opts {
    #[clap(flatten)]
    io: wasm_tools::InputOutput,

    /// Output the text format of WebAssembly instead of the binary format.
    #[clap(short = 't', long)]
    wat: bool,
}

impl Opts {
    pub fn general_opts(&self) -> &wasm_tools::GeneralOpts {
        self.io.general_opts()
    }

    pub fn run(&self) -> Result<()> {
        let input = self.io.get_input_wasm(None)?;
        let mut module = wasm_encoder::Module::new();

        for payload in Parser::new(0).parse_all(&input) {
            let payload = payload?;
            match &payload {
                CustomSection(c) => {
                    if let KnownCustom::Name(s) = c.as_known() {
                        match self.demangle(s) {
                            Ok(new_section) => {
                                module.section(&new_section);
                                continue;
                            }
                            Err(e) => log::debug!("error parsing name section {e:?}"),
                        }
                    }
                }
                Version { encoding, .. } if *encoding == wasmparser::Encoding::Component => {
                    bail!("demangling components is not supported");
                }
                _ => {}
            }
            if let Some((id, range)) = payload.as_section() {
                module.section(&RawSection {
                    id,
                    data: &input[range],
                });
            }
        }

        self.io.output_wasm(module.as_slice(), self.wat)?;
        Ok(())
    }

    fn demangle(&self, section: NameSectionReader<'_>) -> Result<NameSection> {
        let mut new_section = NameSection::new();
        for section in section {
            match section? {
                Name::Module { name, .. } => new_section.module(name),
                Name::Memory(names) => new_section.memories(&self.name_map(names)?),
                Name::Global(names) => new_section.globals(&self.name_map(names)?),
                Name::Function(names) => new_section.functions(&self.name_map(names)?),
                Name::Type(names) => new_section.types(&self.name_map(names)?),
                Name::Table(names) => new_section.tables(&self.name_map(names)?),
                Name::Element(names) => new_section.elements(&self.name_map(names)?),
                Name::Data(names) => new_section.data(&self.name_map(names)?),
                Name::Local(names) => new_section.locals(&self.indirect_name_map(names)?),
                Name::Label(names) => new_section.labels(&self.indirect_name_map(names)?),
                Name::Field(names) => new_section.fields(&self.indirect_name_map(names)?),
                Name::Tag(names) => new_section.tags(&self.name_map(names)?),
                Name::Unknown { .. } => bail!("unknown name section"),
            }
        }
        Ok(new_section)
    }

    fn name_map(&self, names: wasmparser::NameMap<'_>) -> Result<NameMap> {
        let mut ret = NameMap::new();
        for naming in names {
            let naming = naming?;
            let name = match rustc_demangle::try_demangle(naming.name) {
                Ok(name) => name.to_string(),
                Err(_) => cpp_demangle::Symbol::new(naming.name)
                    .ok()
                    .and_then(|s| s.demangle().ok())
                    .unwrap_or_else(|| naming.name.to_string()),
            };
            ret.append(naming.index, &name);
        }
        Ok(ret)
    }

    fn indirect_name_map(&self, names: wasmparser::IndirectNameMap<'_>) -> Result<IndirectNameMap> {
        let mut ret = IndirectNameMap::new();
        for naming in names {
            let naming = naming?;
            let map = self.name_map(naming.names)?;
            ret.append(naming.index, &map);
        }
        Ok(ret)
    }
}
