// Copyright 2026 Rigetti Computing
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Build-time utilities for generating, post-processing, and committing protobuf
//! code produced by [`prost-build`](super::prost_build) and [`pbjson-build`](super::pbjson_build),
//! with optional gRPC service stub generation via [`tonic-prost-build`](super::tonic_prost_build).
//!
//! # Typical `build.rs` pattern
//!
//! ```rust,no_run
//! use qcs_dependencies_client::codegen_proto;
//!
//! fn main() {
//!     let mut builder = codegen_proto::Builder::new("proto", "src/gen");
//!     builder.prost_config.type_attribute("myservice.Foo.bar", "#[derive(serde::Serialize)]");
//!     builder
//!         .pbjson_prefixes([".myservice"])
//!         .tonic(
//!             codegen_grpc::tonic_prost_build::configure()
//!                 .server_mod_attribute("myservice", r#"#[cfg(feature = "server")]"#),
//!         )
//!         .license_header(include_str!("../LICENSE_HEADER"))
//!         .build();
//! }
//! ```
//!
//! Omit `.tonic(...)` to use plain `prost_build` (message types only, no service stubs).
//!
//! For advanced use cases (custom `prost_build::Config` pipeline), the lower-level
//! helpers [`qcs_prost_config`], [`copy_generated_to_src`], and [`fix_generated_paths`]
//! are also available.
//!
//! # Path-rewriting strategy
//!
//! Because `prost`, `tonic`, `tonic_prost`, and `pbjson` are not direct
//! dependencies of consuming crates (they come through `qcs_dependencies_client`),
//! generated code needs two adjustments:
//!
//! - **prost** — handled at generation time via
//!   [`prost_build::Config::prost_path`](super::prost_build::Config::prost_path);
//!   [`qcs_prost_config`] sets this automatically.
//! - **tonic / tonic_prost / pbjson** — handled at copy time by
//!   [`fix_generated_paths`], which is called inside [`copy_generated_to_src`].

use std::{
    env,
    fs::OpenOptions,
    io::Write,
    path::{Path, PathBuf},
};

#[cfg(feature = "grpc")]
pub use super::tonic_prost_build;

/// Builder for a protobuf codegen pipeline.
///
/// Orchestrates proto compilation via `prost_build` (or `tonic_prost_build` when
/// [`.tonic()`](Self::tonic) is called), optional pbjson JSON stub generation, and
/// copying the resulting `.rs` files into the `gen_dir`.
///
/// See the [module-level documentation](self) for a usage example.
pub struct Builder {
    /// A handle to the Prost build-config, to permit applying additional configuration
    pub prost_config: super::prost_build::Config,

    proto_root: PathBuf,
    gen_dir: PathBuf,
    out_dir: PathBuf,
    pbjson_prefixes: Vec<String>,
    #[cfg(feature = "grpc")]
    tonic_builder: Option<tonic_prost_build::Builder>,
    license_header: Option<String>,
}

impl Builder {
    /// Create a new builder.
    ///
    /// `proto_root` and `gen_dir` are resolved relative to `CARGO_MANIFEST_DIR`
    /// (the directory containing the `Cargo.toml` of the crate being built).
    /// Absolute paths are used as-is.
    ///
    /// - `proto_root`: directory containing `.proto` files; used as the include path.
    /// - `gen_dir`: directory where generated `.rs` files will be written.
    pub fn new(proto_root: impl Into<PathBuf>, gen_dir: impl Into<PathBuf>) -> Self {
        let manifest_dir = PathBuf::from(
            env::var("CARGO_MANIFEST_DIR")
                .expect("CARGO_MANIFEST_DIR should be set by Cargo in build scripts"),
        );
        let out_dir = env::var("OUT_DIR").expect("OUT_DIR should be set by Cargo in build scripts");

        let resolve = |p: PathBuf| {
            if p.is_absolute() {
                p
            } else {
                manifest_dir.join(p)
            }
        };

        let prost_config = qcs_prost_config(out_dir.as_str());
        Self {
            proto_root: resolve(proto_root.into()),
            gen_dir: resolve(gen_dir.into()),
            pbjson_prefixes: Vec::new(),
            #[cfg(feature = "grpc")]
            tonic_builder: None,
            license_header: None,
            prost_config,
            out_dir: PathBuf::from(out_dir),
        }
    }

    /// Set the protobuf package prefixes passed to `pbjson_build`.
    ///
    /// If not called, the pbjson step is skipped.
    pub fn pbjson_prefixes(
        mut self,
        prefixes: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        self.pbjson_prefixes = prefixes.into_iter().map(Into::into).collect();
        self
    }

    /// Use `tonic_prost_build` to compile protos, generating service stubs in addition
    /// to message types.
    ///
    /// Pass a pre-configured [`tonic_prost_build::Builder`] (e.g., with
    /// `server_mod_attribute` or `build_server` calls already applied).
    /// When not called, compilation falls back to plain `prost_build`.
    #[cfg(feature = "grpc")]
    pub fn tonic(mut self, builder: tonic_prost_build::Builder) -> Self {
        self.tonic_builder = Some(builder);
        self
    }

    /// Prepend a license header to each generated file when copying to `gen_dir`.
    pub fn license_header(mut self, header: impl Into<String>) -> Self {
        self.license_header = Some(header.into());
        self
    }

    /// Run the full codegen pipeline:
    ///
    /// 1. Emit `cargo:rerun-if-changed` directives for `build.rs`, `gen_dir`, and each proto file.
    /// 2. Compile protos
    /// 3. Generate pbjson JSON stubs (if [`pbjson_prefixes`](Self::pbjson_prefixes) was called).
    /// 4. Copy generated `.rs` files from `OUT_DIR` to `gen_dir` (with optional license header
    ///    and path rewrites via [`fix_generated_paths`]).
    ///
    /// # Panics
    ///
    /// Panics on any I/O or compilation error.
    pub fn build(mut self) {
        println!("cargo:rerun-if-changed=build.rs");
        println!("cargo:rerun-if-changed={}", self.gen_dir.display());

        let proto_files = find_proto_files(&self.proto_root);
        for proto_file in &proto_files {
            println!("cargo:rerun-if-changed={}", proto_file.display());
        }

        #[cfg(feature = "grpc")]
        let tonic_builder = self.tonic_builder;
        #[cfg(not(feature = "grpc"))]
        let tonic_builder: Option<std::convert::Infallible> = None;

        match tonic_builder {
            #[cfg(feature = "grpc")]
            Some(tb) => tb
                .compile_with_config(self.prost_config, &proto_files, &[self.proto_root])
                .expect("tonic_prost_build failed to compile protos"),
            None => self
                .prost_config
                .compile_protos(&proto_files, &[self.proto_root])
                .expect("prost_build failed to compile protos"),
        }

        if !self.pbjson_prefixes.is_empty() {
            let descriptor_path = PathBuf::from(&self.out_dir).join("proto_descriptor.bin");
            let descriptor_set =
                std::fs::read(&descriptor_path).expect("failed to read proto descriptor set");
            super::pbjson_build::Builder::new()
                .register_descriptors(&descriptor_set)
                .expect("failed to register descriptors with pbjson_build")
                .build(&self.pbjson_prefixes)
                .expect("pbjson_build failed to generate JSON stubs");
        }

        copy_generated_to_src(&self.out_dir, &self.gen_dir, self.license_header.as_deref());
    }
}

/// Recursively walks `root` and returns absolute paths to every `.proto` file found,
/// sorted for a deterministic compilation order.
///
/// # Panics
///
/// Panics on any I/O error while traversing the directory tree.
pub fn find_proto_files(root: &Path) -> Vec<PathBuf> {
    let mut files = Vec::new();
    collect_proto_files(root, &mut files);
    files.sort();
    files
}

fn collect_proto_files(dir: &Path, out: &mut Vec<PathBuf>) {
    for entry in std::fs::read_dir(dir)
        .unwrap_or_else(|e| panic!("cannot read directory {}: {e}", dir.display()))
    {
        let path = entry
            .unwrap_or_else(|e| panic!("bad dir entry in {}: {e}", dir.display()))
            .path();
        if path.is_dir() {
            collect_proto_files(&path, out);
        } else if path.extension().is_some_and(|ext| ext == "proto") {
            out.push(path);
        }
    }
}

/// Creates a [`prost_build::Config`] pre-configured to route all generated paths
/// through `qcs_dependencies_client`:
///
/// - `compile_well_known_types()` — replaces `prost-types` well-known types.
/// - `extern_path(".google.protobuf", "qcs_dependencies_client::pbjson_types")` —
///   maps protobuf well-known types to `pbjson-types`.
/// - `prost_path("::qcs_dependencies_client::prost")` — routes all prost
///   runtime paths in generated `impl` bodies through `qcs_dependencies_client`,
///   so consuming crates don't need a direct `prost` dependency.
///
/// Callers should add `file_descriptor_set_path`, `type_attribute`, etc. as needed.
pub fn qcs_prost_config<O>(out_dir: O) -> super::prost_build::Config
where
    PathBuf: From<O>,
{
    let mut config = super::prost_build::Config::new();
    let descriptor_path = PathBuf::from(out_dir).join("proto_descriptor.bin");
    config
        .file_descriptor_set_path(&descriptor_path)
        .compile_well_known_types()
        .extern_path(".google.protobuf", "qcs_dependencies_client::pbjson_types")
        // Route prost runtime paths through qcs_dependencies_client so grpc crates don't need a direct prost dependency.
        .prost_path("::qcs_dependencies_client::prost");
    config
}

/// Copies every `.rs` file from `out_dir` to `dest_dir`.
///
/// Each file is written as follows:
/// 1. If `header` is `Some(text)`, `text` is written first (with a trailing
///    newline).
/// 2. The file content is processed by [`fix_generated_paths`] and written.
///
/// This is the standard way to commit generated protobuf code from `OUT_DIR`
/// into a `src/gen/` directory in source control.
///
/// # Panics
///
/// Panics on any I/O or compilation error.
pub fn copy_generated_to_src(
    out_dir: impl AsRef<Path>,
    dest_dir: impl AsRef<Path>,
    header: Option<&str>,
) {
    let out_dir = out_dir.as_ref();
    let dest_dir = dest_dir.as_ref();

    std::fs::create_dir_all(dest_dir)
        .unwrap_or_else(|e| panic!("cannot create dest_dir {}: {e}", dest_dir.display()));

    for entry in std::fs::read_dir(out_dir)
        .unwrap_or_else(|e| panic!("cannot read OUT_DIR {}: {e}", out_dir.display()))
    {
        let src_path = entry
            .unwrap_or_else(|e| panic!("bad entry in OUT_DIR: {e}"))
            .path();

        if !src_path.to_string_lossy().ends_with(".rs") {
            continue;
        }

        let dest = dest_dir.join(
            src_path
                .file_name()
                .expect("OUT_DIR entry has no file name"),
        );

        let mut file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&dest)
            .unwrap_or_else(|e| panic!("cannot open {} for writing: {e}", dest.display()));

        if let Some(h) = header {
            writeln!(file, "{h}").expect("failed to write header");
        }

        let content = std::fs::read_to_string(&src_path)
            .unwrap_or_else(|e| panic!("cannot read {}: {e}", src_path.display()));
        writeln!(file, "{}", fix_generated_paths(content))
            .expect("failed to write generated content");
    }
}

/// Post-processes a generated `.rs` file so it compiles when `tonic`,
/// `tonic_prost`, and `pbjson` are accessed through `qcs_dependencies_client`.
///
/// Instead of rewriting every path inline, this adds `use` imports so that the
/// generated code's short paths (`tonic::`, `tonic_prost::`, `pbjson::`) resolve
/// correctly through `qcs_dependencies_client`. Imports are added both at the
/// top of the file and inside each `pub mod` block (since child modules don't
/// inherit `use` imports from their parent).
///
/// **prost paths are not rewritten here** — they are handled at generation time
/// via [`prost_build::Config::prost_path`](super::prost_build::Config::prost_path)
/// (set by [`qcs_prost_config`]).
pub fn fix_generated_paths(content: String) -> String {
    let has_tonic = content.contains("tonic::");
    let has_tonic_prost = content.contains("tonic_prost::");
    let has_pbjson = content.contains("pbjson::");

    if !has_tonic && !has_tonic_prost && !has_pbjson {
        return content;
    }

    // pbjson-build emits absolute paths like `::pbjson::private::` which bypass
    // `use` imports entirely.  Rewrite those to rely on the `use` statement.
    let content = content.replace("::pbjson::", "pbjson::");

    // Build the result: processed content only.
    // `use tonic::` statements inside nested `mod` blocks are rewritten to go
    // through `qcs_dependencies_client`, and we inject a local `tonic` (and
    // optionally `tonic_prost`) binding at each site. File-level pbjson import
    // is needed for non-absolute pbjson paths (if any exist at file scope).
    let mut result = String::with_capacity(content.len() + 256);
    if has_pbjson {
        result.push_str("use qcs_dependencies_client::pbjson;\n\n");
    }

    let mut last_was_tonic_use = false;
    for line in content.lines() {
        // Rewrite `use tonic::` inside nested modules to fully-qualified paths.
        // Also inject `use qcs_dependencies_client::{tonic, tonic_prost};` (once)
        // so that bare `tonic::` and `tonic_prost::` references in expression
        // position resolve correctly within the nested module.
        let trimmed = line.trim_start();
        if let Some(stripped) = trimmed.strip_prefix("use tonic::") {
            let indent = &line[..line.len() - trimmed.len()];
            if !last_was_tonic_use {
                // First `use tonic::` in this block: add local bindings.
                result.push_str(indent);
                result.push_str("use qcs_dependencies_client::tonic;\n");
                if has_tonic_prost {
                    result.push_str(indent);
                    result.push_str("use qcs_dependencies_client::tonic_prost;\n");
                }
                last_was_tonic_use = true;
            }
            // Rewrite the glob/path import through qcs_dependencies_client.
            result.push_str(indent);
            result.push_str("use qcs_dependencies_client::tonic::");
            result.push_str(stripped);
            result.push('\n');
        } else {
            last_was_tonic_use = false;
            result.push_str(line);
            result.push('\n');
        }
    }

    result
}
