use std::fs;
use std::path::{Path, PathBuf};

const JSON_PACKAGES: [&str; 5] = [
    "v1",
    "common.v1",
    "admin.v1",
    "cluster.v1",
    "cluster.internal.v1",
];

const SERDE_HELPERS: &str = "crate::serde_helpers";

const DESCRIPTOR_TMP: &str = ".descriptor.bin";

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let manifest_dir = Path::new(env!("CARGO_MANIFEST_DIR"));
    let out_dir = PathBuf::from(std::env::var("OUT_DIR")?);
    build_protos(manifest_dir, &out_dir)
}

fn build_protos(manifest_dir: &Path, out_dir: &Path) -> Result<(), Box<dyn std::error::Error>> {
    let proto_dir = PathBuf::from("proto");
    let proto_root = manifest_dir.join(&proto_dir);

    let mut files = Vec::new();
    collect_protos(&proto_root, &mut files)?;
    files.sort();
    if files.is_empty() {
        return Err(format!("no .proto found under {}", proto_root.display()).into());
    }
    let relative: Vec<PathBuf> = files
        .iter()
        .map(|file| proto_dir.join(file.strip_prefix(&proto_root).expect("scanned under root")))
        .collect();

    println!("cargo:rerun-if-changed={}", proto_root.display());
    for file in &files {
        println!("cargo:rerun-if-changed={}", file.display());
    }

    let mut builder = tonic_prost_build::configure()
        .build_client(true)
        .build_server(true)
        .file_descriptor_set_path(DESCRIPTOR_TMP)
        .enum_attribute(".", "#[allow(clippy::large_enum_variant)]")
        .include_file("_include.rs");

    for package in JSON_PACKAGES {
        let prefix = format!(".{package}");
        builder = builder
            .type_attribute(&prefix, "#[derive(serde::Serialize, serde::Deserialize)]")
            .message_attribute(&prefix, "#[serde(default)]")
            .enum_attribute(&prefix, "#[serde(rename_all = \"SCREAMING_SNAKE_CASE\")]");
    }

    let bytes_helper = format!(
        "#[serde(with = \"{SERDE_HELPERS}::bytes\")] #[schema(value_type = String)] #[doc = \"Hex encoded bytes (two hex digits per byte).\"]"
    );
    let bytes_list_helper = format!(
        "#[serde(with = \"{SERDE_HELPERS}::bytes_list\")] #[schema(value_type = Vec<String>)]"
    );
    let timestamp_helper = format!(
        "#[serde(with = \"{SERDE_HELPERS}::timestamp_opt\")] #[schema(value_type = Option<String>, default = json!(null))]"
    );

    for file in &files {
        let text = fs::read_to_string(file)?;
        for decl in scan_proto(&text) {
            builder = match decl {
                Decl::Schema(path) => builder.type_attribute(path, "#[derive(utoipa::ToSchema)]"),
                Decl::Oneof(holder) => builder
                    .field_attribute(holder, "#[schema(value_type = Option<serde_json::Value>)]"),
                Decl::Field(path, FieldKind::Bytes) => builder.field_attribute(path, &bytes_helper),
                Decl::Field(path, FieldKind::BytesList) => {
                    builder.field_attribute(path, &bytes_list_helper)
                }
                Decl::Field(path, FieldKind::Timestamp) => {
                    builder.field_attribute(path, &timestamp_helper)
                }
            };
        }
    }

    builder.compile_protos(&relative, std::slice::from_ref(&proto_dir))?;

    fs::copy(DESCRIPTOR_TMP, out_dir.join("descriptor.bin"))?;
    let _ = fs::remove_file(DESCRIPTOR_TMP);
    Ok(())
}

fn collect_protos(dir: &Path, out: &mut Vec<PathBuf>) -> std::io::Result<()> {
    for entry in fs::read_dir(dir)? {
        let path = entry?.path();
        if path.is_dir() {
            collect_protos(&path, out)?;
        } else if path.extension().is_some_and(|ext| ext == "proto") {
            out.push(path);
        }
    }
    Ok(())
}

enum Decl {
    Schema(String),
    Oneof(String),
    Field(String, FieldKind),
}

enum FieldKind {
    Bytes,
    BytesList,
    Timestamp,
}

enum Block {
    Message(String),
    Other,
}

fn scan_proto(text: &str) -> Vec<Decl> {
    let text = strip_comments(text);
    let Some(package) = package_of(&text) else {
        return Vec::new();
    };
    if !JSON_PACKAGES.contains(&package.as_str()) {
        return Vec::new();
    }

    let mut decls = Vec::new();
    let mut stack: Vec<Block> = Vec::new();
    for raw in text.lines() {
        let line = raw.trim();
        if line.is_empty() {
            continue;
        }
        if let Some(head) = line.strip_suffix('{') {
            let head = head.trim();
            let mut tokens = head.split_whitespace();
            let keyword = tokens.next().unwrap_or_default();
            let name = tokens.next().unwrap_or_default().to_string();
            match keyword {
                "message" => {
                    decls.push(Decl::Schema(path_of(&package, &stack, &name)));
                    stack.push(Block::Message(name));
                }
                "enum" => {
                    decls.push(Decl::Schema(path_of(&package, &stack, &name)));
                    stack.push(Block::Other);
                }
                "oneof" => {
                    decls.push(Decl::Oneof(path_of(&package, &stack, &name)));
                    stack.push(Block::Other);
                }
                _ => stack.push(Block::Other),
            }
            continue;
        }
        if let Some(head) = line.strip_suffix("{}") {
            let head = head.trim();
            let mut tokens = head.split_whitespace();
            let keyword = tokens.next().unwrap_or_default();
            let name = tokens.next().unwrap_or_default();
            if matches!(keyword, "message" | "enum") && !name.is_empty() {
                decls.push(Decl::Schema(path_of(&package, &stack, name)));
            }
            continue;
        }
        if line == "}" {
            stack.pop();
            continue;
        }
        if !matches!(stack.last(), Some(Block::Message(_))) {
            continue;
        }
        let Some((ty, name, repeated)) = parse_field(line) else {
            continue;
        };
        let kind = match (ty.as_str(), repeated) {
            ("google.protobuf.Timestamp", _) => FieldKind::Timestamp,
            ("bytes", true) => FieldKind::BytesList,
            ("bytes", false) => FieldKind::Bytes,
            _ => continue,
        };
        decls.push(Decl::Field(path_of(&package, &stack, &name), kind));
    }
    decls
}

fn package_of(text: &str) -> Option<String> {
    text.lines()
        .find_map(|line| line.trim().strip_prefix("package "))
        .map(|rest| rest.trim().trim_end_matches(';').trim().to_string())
}

fn path_of(package: &str, stack: &[Block], name: &str) -> String {
    let mut path = format!(".{package}");
    for block in stack {
        if let Block::Message(outer) = block {
            path.push('.');
            path.push_str(outer);
        }
    }
    path.push('.');
    path.push_str(name);
    path
}

fn parse_field(line: &str) -> Option<(String, String, bool)> {
    let body = line.strip_suffix(';')?;
    let tokens: Vec<&str> = body.split_whitespace().collect();
    let eq = tokens.iter().position(|token| *token == "=")?;
    if eq < 2 {
        return None;
    }
    let name = tokens[eq - 1];
    let ty = tokens[eq - 2];
    if ty.contains('<') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
        return None;
    }
    let repeated = eq >= 3 && tokens[eq - 3] == "repeated";
    Some((ty.to_string(), name.to_string(), repeated))
}

fn strip_comments(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    let mut chars = text.chars().peekable();
    let mut line_comment = false;
    let mut block_comment = false;
    while let Some(current) = chars.next() {
        if line_comment {
            if current == '\n' {
                line_comment = false;
                out.push('\n');
            }
            continue;
        }
        if block_comment {
            if current == '*' && chars.peek() == Some(&'/') {
                chars.next();
                block_comment = false;
            }
            continue;
        }
        if current == '/' && chars.peek() == Some(&'/') {
            chars.next();
            line_comment = true;
            continue;
        }
        if current == '/' && chars.peek() == Some(&'*') {
            chars.next();
            block_comment = true;
            continue;
        }
        out.push(current);
    }
    out
}