oas3-gen 0.26.2

use std::{
  collections::{BTreeSet, HashMap},
  rc::Rc,
};

use clap::ValueEnum;
use proc_macro2::TokenStream;
use quote::{ToTokens, quote};

use self::{client::ClientFragment, mod_file::ModFileFragment, server::ServerGenerator, types::TypesFragment};
use super::ast::{ClientRootNode, GlobalLintsNode, OperationInfo, RustType, ServerRequestTraitDef};
use crate::generator::{
  ast::{Documentation, FileHeaderNode, constants::HttpHeaderRef},
  converter::CodegenConfig,
};

pub mod attributes;
pub mod client;
pub mod coercion;
pub mod constants;
pub mod enums;
pub(crate) mod headers;
pub mod http;
pub mod methods;
pub mod mod_file;
pub mod server;
pub mod structs;
pub mod type_aliases;
pub mod types;

#[cfg(test)]
mod tests;

#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Visibility {
  #[default]
  Public,
  Crate,
  File,
}

impl ToTokens for Visibility {
  fn to_tokens(&self, tokens: &mut TokenStream) {
    let ts = match self {
      Visibility::Public => quote! { pub },
      Visibility::Crate => quote! { pub(crate) },
      Visibility::File => quote! {},
    };
    tokens.extend(ts);
  }
}

impl Visibility {
  pub(crate) fn to_tokens(self) -> TokenStream {
    match self {
      Visibility::Public => quote! { pub },
      Visibility::Crate => quote! { pub(crate) },
      Visibility::File => quote! {},
    }
  }
}

pub fn format(code: &TokenStream) -> anyhow::Result<String> {
  let syntax_tree = syn::parse2(code.clone())?;
  Ok(prettyplease::unparse(&syntax_tree))
}

impl ToTokens for FileHeaderNode {
  fn to_tokens(&self, tokens: &mut TokenStream) {
    self.lints.to_tokens(tokens);

    let title_line = format!(" {}", self.title);
    let source_line = format!(" Source: {}", self.source_path);
    let version_line = format!(" Version: {}", self.version);
    let generator_line = format!(" Generated by `oas3-gen v{}`", self.generator_version);

    let header = quote! {
      #![doc = ""]
      #![doc = " AUTO-GENERATED CODE - DO NOT EDIT!"]
      #![doc = ""]
      #![doc = #title_line]
      #![doc = #source_line]
      #![doc = #version_line]
      #![doc = #generator_line]
      #![doc = ""]
    };

    tokens.extend(header);

    let description_text = self.description.clone().map(|d| d.with_top_level(true));

    if let Some(description_text) = description_text {
      tokens.extend(description_text.to_token_stream());
    }
  }
}

fn generate_source(
  code: &TokenStream,
  metadata: &ClientRootNode,
  lint_config: Option<&GlobalLintsNode>,
  source_path: &str,
  gen_version: &str,
) -> anyhow::Result<String> {
  let lints = lint_config.map_or_else(|| GlobalLintsNode::builder().build(), std::clone::Clone::clone);
  let header_node = FileHeaderNode::builder()
    .title(metadata.title.clone())
    .version(metadata.version.clone())
    .source_path(source_path.to_string())
    .generator_version(gen_version.to_string())
    .description(Documentation::from_optional(metadata.description.as_ref()))
    .lints(lints)
    .build();

  let header_formatted = format(&header_node.to_token_stream())?;
  let code_formatted = format(code)?;
  Ok(format!("{header_formatted}\n{code_formatted}\n"))
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum GeneratedFileType {
  Client,
  Module,
  Server,
  Types,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GeneratedResult {
  pub code: HashMap<GeneratedFileType, String>,
}

impl GeneratedResult {
  pub fn client(code: String) -> Self {
    Self {
      code: HashMap::from([(GeneratedFileType::Client, code)]),
    }
  }

  pub fn types(code: String) -> Self {
    Self {
      code: HashMap::from([(GeneratedFileType::Types, code)]),
    }
  }

  pub fn full_client(mod_file: String, client: String, types: String) -> Self {
    Self {
      code: HashMap::from([
        (GeneratedFileType::Module, mod_file),
        (GeneratedFileType::Client, client),
        (GeneratedFileType::Types, types),
      ]),
    }
  }

  pub fn full_server(mod_file: String, server: String, types: String) -> Self {
    Self {
      code: HashMap::from([
        (GeneratedFileType::Module, mod_file),
        (GeneratedFileType::Server, server),
        (GeneratedFileType::Types, types),
      ]),
    }
  }

  pub fn code(&self, file_type: &GeneratedFileType) -> Option<&String> {
    self.code.get(file_type)
  }
}

#[derive(Debug, Clone)]
pub struct SchemaCodeGenerator {
  config: CodegenConfig,
  rust_types: Rc<Vec<RustType>>,
  operations: Rc<Vec<OperationInfo>>,
  header_refs: Rc<Vec<HttpHeaderRef>>,
  uses: Rc<BTreeSet<String>>,
  client: Rc<ClientRootNode>,
  server_trait: Option<ServerRequestTraitDef>,
  visibility: Visibility,
  source_path: String,
  gen_version: String,
}

#[bon::bon]
impl SchemaCodeGenerator {
  #[builder]
  pub fn new(
    config: CodegenConfig,
    rust_types: Vec<RustType>,
    operations: Vec<OperationInfo>,
    header_refs: Vec<HttpHeaderRef>,
    uses: BTreeSet<String>,
    client: ClientRootNode,
    server_trait: Option<ServerRequestTraitDef>,
    visibility: Visibility,
    source_path: String,
    gen_version: String,
  ) -> Self {
    Self {
      config,
      rust_types: Rc::new(rust_types),
      operations: Rc::new(operations),
      header_refs: Rc::new(header_refs),
      uses: Rc::new(uses),
      client: Rc::new(client),
      server_trait,
      visibility,
      source_path,
      gen_version,
    }
  }
}

impl SchemaCodeGenerator {
  /// Generates a standalone `types.rs` file containing all type definitions.
  ///
  /// Includes structs, enums, type aliases, and their serde/validation derives.
  pub fn generate_types(&self) -> anyhow::Result<GeneratedResult> {
    let code = self.format_tokens_with_lints(&self.types_fragment())?;
    Ok(GeneratedResult::types(code))
  }

  /// Generates a standalone `client.rs` file with types and HTTP client combined.
  ///
  /// The client struct includes methods for each API operation.
  pub fn generate_client(&self) -> anyhow::Result<GeneratedResult> {
    let code = self.format_tokens_with_lints(&self.client_fragment(false))?;
    Ok(GeneratedResult::client(code))
  }

  /// Generates a modular client with separate `mod.rs`, `client.rs`, and `types.rs` files.
  ///
  /// The client imports types from the sibling `types` module.
  pub fn generate_client_mod(&self) -> anyhow::Result<GeneratedResult> {
    let types_code = self.format_tokens(&self.types_fragment())?;
    let client_code = self.format_tokens(&self.client_fragment(true))?;
    let mod_fragment = ModFileFragment::for_client(
      (*self.client).clone(),
      self.visibility,
      self.source_path.clone(),
      self.gen_version.clone(),
    );
    let mod_code = mod_fragment.generate()?;

    Ok(GeneratedResult::full_client(mod_code, client_code, types_code))
  }

  /// Generates a modular server with separate `mod.rs`, `server.rs`, and `types.rs` files.
  ///
  /// The server trait imports types from the sibling `types` module.
  pub fn generate_server_mod(&self) -> anyhow::Result<GeneratedResult> {
    let types_code = self.format_tokens(&self.types_fragment())?;
    let server_code = self.format_tokens(&self.server_fragment())?;
    let mod_fragment = ModFileFragment::for_server(
      (*self.client).clone(),
      self.visibility,
      self.source_path.clone(),
      self.gen_version.clone(),
    );
    let mod_code = mod_fragment.generate()?;

    Ok(GeneratedResult::full_server(mod_code, server_code, types_code))
  }

  /// Creates a types generator fragment for all Rust type definitions.
  fn types_fragment(&self) -> TypesFragment {
    TypesFragment::new(
      self.rust_types.clone(),
      self.header_refs.clone(),
      (*self.uses).clone(),
      self.visibility,
      self.config.target,
    )
  }

  /// Creates a client fragment for HTTP client code generation.
  fn client_fragment(&self, with_types_import: bool) -> ClientFragment {
    let fragment = ClientFragment::new(&self.client, &self.operations, self.visibility);
    if with_types_import {
      fragment.with_types_import()
    } else {
      fragment
    }
  }

  /// Creates a server fragment for axum server trait generation.
  fn server_fragment(&self) -> ServerGenerator {
    ServerGenerator::new(self.server_trait.clone(), self.visibility).with_types_import()
  }

  /// Formats tokens into source code with a file header (no lint attributes).
  fn format_tokens(&self, fragment: &impl ToTokens) -> anyhow::Result<String> {
    generate_source(
      &fragment.to_token_stream(),
      &self.client,
      None,
      &self.source_path,
      &self.gen_version,
    )
  }

  /// Formats tokens into source code with a file header and default lint configuration.
  fn format_tokens_with_lints(&self, fragment: &impl ToTokens) -> anyhow::Result<String> {
    let lints = GlobalLintsNode::default();
    generate_source(
      &fragment.to_token_stream(),
      &self.client,
      Some(&lints),
      &self.source_path,
      &self.gen_version,
    )
  }
}