// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

mod bg;
mod fields;

use {
  super::parse::Ast,
  crate::{args::Args, ast::parse::FieldType},
  bg::*,
  fields::*,
  heck::{ToSnakeCase, ToUpperCamelCase},
  proc_macro2::{self, TokenStream},
  quote::quote,
};

pub fn codegen(
  args: Args,
  ast: Ast,
) -> Result<TokenStream, crate::error::Error> {
  let Ast { name, fields } = ast;

  let field_data = &fields
    .iter()
    .map(|f| -> Result<TokenStream, crate::error::Error> {
      let (name, ty) = match f.ty {
        | FieldType::Literal(ty_name) => (&f.name, ty_name),
        | FieldType::Error(ref _val) => {
          return Err(crate::error::Error::UnknownFieldType(
            proc_macro2::Span::call_site(),
            "Error field type not supported".to_string(),
          ));
        },
        | FieldType::Path(ref val) => (&f.name, val.as_str()),
        | FieldType::SingleNode(ref _val) => (&f.name, "crate::NodeId"),
        | FieldType::EnumNodeId(ref _val) => (&f.name, "crate::NodeId"),
        | FieldType::MultipleNode(ref _types) => (&f.name, "crate::NodeId"),
        | FieldType::String => (&f.name, "laburnum::Span"),
        | FieldType::Span => (&f.name, "laburnum::Span"),
        | FieldType::Ident => (&f.name, "laburnum::Ident"),
        | FieldType::SpannedIdent => {
          (&f.name, "laburnum::Spanned<laburnum::Ident>")
        },
        | FieldType::Enum(ref ty_name) => (&f.name, ty_name.as_str()),
        | _ => {
          return Err(crate::error::Error::UnknownFieldType(
            proc_macro2::Span::call_site(),
            format!("{:?}", f.ty),
          ));
        },
      };

      let ty = {
        if f.is_optional {
          format!("Option<{ty}>")
        } else {
          ty.to_string()
        }
      };
      let ty = {
        if f.is_vec {
          format!("Vec<{ty}>")
        } else {
          ty.to_string()
        }
      };

      let field_name: syn::Ident = syn::parse_str(name).expect("Invalid name");
      let path: syn::Path = syn::parse_str(&ty).expect("Invalid type");

      Ok(quote::quote! {
        pub(crate) #field_name: #path
      })
    })
    .collect::<Result<Vec<TokenStream>, _>>()?;

  let struct_name =
    proc_macro2::Ident::new(&name, proc_macro2::Span::call_site());

  let semantics = {
    if args.has_semantic_tokens() {
      // The following can be overwritten later in the pipeline, so they are
      // optional If they are Some they will take precedence over the
      // normal value assigned to this Ident.
      quote! {
        pub semantic_token_type: Option<laburnum::protocol::lsp::SemanticTokenType>,
        pub semantic_token_modifier: Option<laburnum::protocol::lsp::SemanticTokenModifier>,
      }
    } else {
      quote! {}
    }
  };

  let ts_struct = quote! {
    #[derive(Clone, PartialEq)]
    pub struct #struct_name {
      /// The unique identifier for this node.
      pub(crate) id: crate::NodeId,
      /// The syntax node id that this node is based on.
      pub(crate) syntax: crate::SyntaxId,
      /// Any meta information that is associated with this node.
      pub(crate) meta: Option<crate::NodeId>,

      #semantics

      #(#field_data),*
    }

    impl #struct_name {
      pub fn get_id(&self) -> crate::NodeId {
        self.id
      }

      pub fn get_syntax(&self) -> crate::SyntaxId {
        self.syntax
      }

      pub fn get_meta(&self) -> Option<crate::NodeId> {
        self.meta
      }
    }
  };

  let current_mod_name = quote! {
    {
      let __mp = module_path!();
      __mp.find("::").map(|i| &__mp[i+2..]).unwrap_or(__mp)
    }
  };

  // -- Display, Debug impl ----------------------------------------------------

  let ts_impl_display = quote! {
      #[allow(unused)]
      impl std::fmt::Display for #struct_name {
        fn fmt(&self,f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
          f.debug_struct(&format!("{}::{}", #current_mod_name,#name))
            .finish()
        }
      }
  };

  let ts_impl_debug = quote! {
    #[allow(unused)]
    impl std::fmt::Debug for #struct_name {
      fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct(&format!("{}::{}", #current_mod_name, #name))
          .finish()
      }
    }
  };

  // -- Bluegum impl -----------------------------------------------------------

  let ts_impl_bluegum = quote! {
    #[allow(unused)]
    impl bluegum::Bluegum for #struct_name {
      fn node(&self,b: &mut bluegum::Builder){
        b.name(
          &crate::style_title(
            #current_mod_name,
            &#name,
          ),
        );
      }
    }
  };

  let ts_impl_bluegum_with_state = {
    let fields = fields
      .iter()
      .map(bluegum_map_field_quote)
      .collect::<Vec<TokenStream>>();

    quote! {
    #[allow(unused)]
    impl<'tokens, 'src>
      bluegum::BluegumWithState<crate::Printer<'tokens>> for #struct_name {
      fn node_with_state(
        &self,
        b: &mut bluegum::Builder,
        state: &crate::Printer<'tokens>,
      ){
        use owo_colors::OwoColorize;
        use unicode_width::UnicodeWidthStr;

        b.name(
          &crate::style_title(
            #current_mod_name,
            &#name,
          ),
        );

        // Add span debug info from CST syntax node
        if let Some(cst_node) = state.get_syntax(self.syntax) {
          let span = cst_node.span();
          if let Some(span_data) = span.data(state.get_span_cache()) {
            b.debug("span", format!("{}..{}", span_data.start, span_data.start + span_data.len));
          }
        }

        if let Some(meta) = self.meta {
          b.add_node_with_state(state, &"meta".blue().italic().to_string(), &meta);
        }

        { #(#fields)* };

        }
      }
    }
  };

  // -- Walk impl --------------------------------------------------------------

  let ts_impl_walk = {
    let fields_mod_name_str = format!("{}_field", heck::AsSnekCase(&name));

    let fields_mod_name = proc_macro2::Ident::new(
      &fields_mod_name_str,
      proc_macro2::Span::call_site(),
    );

    let (fields, valid_field_tys): (
      Vec<(TokenStream, TokenStream)>,
      Vec<TokenStream>,
    ) = fields
      .iter()
      .filter(|f| f.name != "syntax")
      .map(|f| -> Result<_, crate::error::Error> {
        let (field_ty, maybe_field_ty_name) = map_valid_field_type(f)?;

        let maybe_field_ty_name: Option<TokenStream> = match maybe_field_ty_name
        {
          | Some(ty_name) => {
            let ty_path = syn::parse_str::<syn::Path>(&format!(
              "{fields_mod_name_str}::{ty_name}"
            ))
            .ok();

            Some(quote! {#ty_path})
          },
          | None => None,
        };

        let (impl_walk, walk_fn) = map_impl_walk(f, maybe_field_ty_name);

        let validate = {
          match &f.ty {
            | FieldType::EnumNodeId(_ty) => {
              if f.is_vec {
                quote! {
                  self.#walk_fn(ast).iter().map(|f| f.validate(ast).unwrap());
                }
              } else {
                quote! {
                  self.#walk_fn(ast).unwrap().validate(ast).unwrap();
                }
              }
            },
            | FieldType::MultipleNode(_items) => {
              if f.is_vec {
                quote! {
                  self.#walk_fn(ast).iter().map(|f| f.validate(ast).unwrap());
                }
              } else {
                quote! {
                  self.#walk_fn(ast).unwrap().validate(ast).unwrap();
                }
              }
            },
            | _ => quote! {},
          }
        };

        Ok(((impl_walk, validate), field_ty))
      })
      .collect::<Result<Vec<_>, _>>()?
      .into_iter()
      .unzip();

    let (fields, field_methods): (Vec<TokenStream>, Vec<TokenStream>) =
      fields.into_iter().unzip();

    quote! {
      pub mod #fields_mod_name {
        #(#valid_field_tys)*
      }

      #[allow(unused)]
      impl #struct_name {
        #(#fields)*
      }

      #[allow(unused)]
      impl<'ast> #struct_name {
        pub fn validate(
          &self,
          ast: &'ast crate::AST,
        ) -> crate::Result<()>{

          #( #field_methods )*

          Ok(())
        }
      }
    }
  };

  // -- Assembly ---------------------------------------------------------------

  // Assemble the final token stream
  let ts = quote! {
    #ts_struct
    #ts_impl_walk

    #ts_impl_display
    #ts_impl_debug

    #ts_impl_bluegum
    #ts_impl_bluegum_with_state
  };

  // The following snippet is useful to see the generated code
  // eprintln!(
  //   "---/n{}/n---",
  //   prettyplease::unparse(&syn::parse_file(&ts.to_string()).unwrap())
  // );

  Ok(ts)
}
