use heck::ToUpperCamelCase;
use std::{io::BufWriter, path::PathBuf};

static MESSAGE_TEMPLATE: &str = include_str!("./templates/message.handlebars");

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, serde::Serialize)]
struct ParsedMessageKind {
    name: String,
    id: u64,
}

fn main() -> anyhow::Result<()> {
    let proto_files = ["protos/api.proto", "protos/api_options.proto"];
    let includes = ["protos"];

    prost_build::Config::new()
        .btree_map(["."])
        .message_attribute(
            ".",
            "#[cfg_attr(feature = \"defmt\", derive(defmt::Format))]",
        )
        .default_package_filename("api")
        .compile_protos(&proto_files, &includes)?;

    let valid_sources = 0..=2;

    let mut message_kinds: Vec<ParsedMessageKind> = protobuf_parse::Parser::new()
        .inputs(proto_files)
        .includes(includes)
        .parse_and_typecheck()?
        .file_descriptors
        .into_iter()
        .flat_map(|fd| fd.message_type)
        .filter_map(|message| {
            if let Some(protobuf::UnknownValueRef::Varint(id)) = message
                .options
                .get_or_default()
                .special_fields
                .unknown_fields()
                .get(1036)
                && let Some(protobuf::UnknownValueRef::Varint(source)) = message
                    .options
                    .get_or_default()
                    .special_fields
                    .unknown_fields()
                    .get(1037)
                && valid_sources.contains(&source)
            {
                Some(ParsedMessageKind {
                    name: sanitize_identifier(message.name().to_upper_camel_case()),
                    id,
                })
            } else {
                None
            }
        })
        .collect();

    message_kinds.sort();

    let out_dir = PathBuf::from(std::env::var("OUT_DIR")?).join("api.rs");

    let api_constants = BufWriter::new(std::fs::OpenOptions::new().append(true).open(out_dir)?);

    handlebars::Handlebars::new().render_template_to_write(
        MESSAGE_TEMPLATE,
        &message_kinds,
        api_constants,
    )?;

    Ok(())
}

/// From prost-build, as it is in a private module.
fn sanitize_identifier(ident: String) -> String {
    // Use a raw identifier if the identifier matches a Rust keyword:
    // https://doc.rust-lang.org/reference/keywords.html.
    match ident.as_str() {
        // 2015 strict keywords.
        | "as" | "break" | "const" | "continue" | "else" | "enum" | "false"
        | "fn" | "for" | "if" | "impl" | "in" | "let" | "loop" | "match" | "mod" | "move" | "mut"
        | "pub" | "ref" | "return" | "static" | "struct" | "trait" | "true"
        | "type" | "unsafe" | "use" | "where" | "while"
        // 2018 strict keywords.
        | "dyn"
        // 2015 reserved keywords.
        | "abstract" | "become" | "box" | "do" | "final" | "macro" | "override" | "priv" | "typeof"
        | "unsized" | "virtual" | "yield"
        // 2018 reserved keywords.
        | "async" | "await" | "try"
        // 2024 reserved keywords.
        | "gen" => format!("r#{ident}"),
        // the following keywords are not supported as raw identifiers and are therefore suffixed with an underscore.
        "_" | "super" | "self" | "Self" | "extern" | "crate" => format!("{ident}_"),
        // the following keywords begin with a number and are therefore prefixed with an underscore.
        s if s.starts_with(char::is_numeric) => format!("_{ident}"),
        _ => ident,
    }
}