Framework for embedding localizations into Rust types
use crate::fluent;
use std::collections::HashSet;

use heck::ToKebabCase;
use proc_macro2::{Ident, TokenStream};
use quote::quote;
use syn::punctuated::Punctuated;

use super::{MacroError, UnsupportedError};

#[derive(Clone, Copy, Debug)]
pub enum ReferenceKind {
    EnumField,
    StructField,
}

#[derive(Clone, Debug)]
pub struct Context {
    pub reference_kind: ReferenceKind,
    pub valid_references: HashSet<String>,
}

fn expr_for_message(
    group: &mut fluent::Group,
    ident: &Ident,
    context: &Context,
) -> Result<TokenStream, fluent::GroupError> {
    let canonical_message = group.remove_canonical_message(ident, context)?;

    let (additional_locales, additional_messages): (Vec<_>, Vec<_>) = group
        .remove_additional_messages(ident, context)?
        .into_iter()
        .unzip();

    let additional_locales = additional_locales
        .iter()
        .map(|locale| locale.id.to_string())
        .collect::<Vec<_>>();

    Ok(quote! {
        let mut buffer = String::new();

        // TODO: handle different rule types according to fluent code (not just cardinal)
        // TODO: only generate this when needed in message
        const PLURAL_RULE_TYPE: ::l10n_embed::macro_prelude::icu_plurals::PluralRuleType =
            ::l10n_embed::macro_prelude::icu_plurals::PluralRuleType::Cardinal;
        let plural_options = ::l10n_embed::macro_prelude::icu_plurals::PluralRulesOptions::default()
            .with_type(PLURAL_RULE_TYPE);
        let plural_rules = ::l10n_embed::macro_prelude::icu_plurals::PluralRules::try_new(
            locale.into(),
            plural_options,
        )
        .unwrap();


        // Handle any additional locales
        #(if locale.normalizing_eq(#additional_locales) {
            #additional_messages
            return buffer;
        }) else*
        // Fall back to the canonical locale, if no other valid locale was matched
        #canonical_message

        buffer
    })
}

fn expr_for_unnamed_fields(
    unnamed_fields: &syn::FieldsUnnamed,
    ident: &syn::Ident,
    context: &Context,
) -> Result<TokenStream, UnsupportedError> {
    let field_count = unnamed_fields.unnamed.iter().count();
    if field_count != 1 {
        return Err(UnsupportedError::UnnamedFields {
            span: ident.clone(),
            field_count,
        });
    }

    let field_reference = match context.reference_kind {
        ReferenceKind::EnumField => quote!(unnamed_field),
        ReferenceKind::StructField => quote!(self.0),
    };

    Ok(quote!(#field_reference.localize_for(locale)))
}

/// Create a list of unique field names that can be referenced
fn unique_named_fields(named_fields: &syn::FieldsNamed) -> HashSet<String> {
    named_fields
        .named
        .iter()
        // Get the `syn::Ident` for each field
        .map(|field| {
            field
                .ident
                .as_ref()
                .expect("Named fields should have an associated ident")
        })
        .map(std::string::ToString::to_string)
        .collect::<HashSet<String>>()
}

pub fn locales_for_ident(
    group: &fluent::Group,
    fields: &syn::Fields,
    reference_kind: ReferenceKind,
    ident: &syn::Ident,
) -> TokenStream {
    match fields {
        // Provide available locales based on what Fluent source files are available
        syn::Fields::Named(_) | syn::Fields::Unit => {
            let id = ident.to_string().to_kebab_case();

            let locale_literals = group
                .locales_for_message(&id)
                .map(|locale| locale.id.to_string())
                .map(|locale_string| {
                    syn::LitStr::new(&locale_string, proc_macro2::Span::call_site())
                });

            quote!(
                vec![self.canonical_locale(), #(::l10n_embed::macro_prelude::icu_locale::locale!(#locale_literals)),*]
            )
        }
        // Forward to the unnamed field's implementation
        syn::Fields::Unnamed(_unnamed_fields) => {
            let unnamed_field_ident = match reference_kind {
                ReferenceKind::EnumField => quote!(unnamed_field),
                ReferenceKind::StructField => quote!(self.0),
            };

            quote!(#unnamed_field_ident.available_locales())
        }
    }
}

pub fn message_for_struct(
    mut group: fluent::Group,
    ident: &syn::Ident,
    fields: &syn::Fields,
) -> Result<TokenStream, MacroError> {
    // Turn the struct fields into a list of valid references
    let context = match fields {
        syn::Fields::Named(named_fields) => {
            Context {
                // Reference using `self.{reference_name}`
                reference_kind: ReferenceKind::StructField,
                // Create a list of unique field names that can be referenced
                valid_references: unique_named_fields(named_fields),
            }
        }
        syn::Fields::Unit | syn::Fields::Unnamed(_) => Context {
            reference_kind: ReferenceKind::StructField,
            valid_references: HashSet::new(),
        },
    };

    Ok(match fields {
        syn::Fields::Named(_) | syn::Fields::Unit => expr_for_message(&mut group, ident, &context)?,
        syn::Fields::Unnamed(unnamed_fields) => {
            expr_for_unnamed_fields(unnamed_fields, ident, &context)?
        }
    })
}

pub fn locales_for_enum(
    group: &fluent::Group,
    enum_variants: &Punctuated<syn::Variant, syn::token::Comma>,
) -> TokenStream {
    let mut match_arms: Vec<TokenStream> = Vec::with_capacity(enum_variants.len());

    for enum_variant in enum_variants {
        let variant_ident = &enum_variant.ident;

        let destructuring_pattern = match &enum_variant.fields {
            // Simplify match code by always ignoring enum fields (even if they don't exist)
            // We are matching the variant name, not any data, so each arm will have something like:
            // Self::VariantName { .. }
            // Even if `Self::VariantName` doesn't contain any data
            syn::Fields::Named(_) | syn::Fields::Unit => quote!(Self::#variant_ident { .. }),
            // Bind the single unnamed field so we can use its implementation of `Localize::available_locales()`
            syn::Fields::Unnamed(_unnamed_fields) => quote!(Self::#variant_ident(unnamed_field)),
        };
        let locales_for_variant = locales_for_ident(
            group,
            &enum_variant.fields,
            ReferenceKind::EnumField,
            variant_ident,
        );
        match_arms.push(quote!(#destructuring_pattern => #locales_for_variant));
    }

    quote! {
        match self {
            #(#match_arms),*
        }
    }
}

pub fn messages_for_enum(
    mut group: fluent::Group,
    enum_variants: &Punctuated<syn::Variant, syn::token::Comma>,
) -> Result<TokenStream, MacroError> {
    // let mut match_arms: HashMap<String, syn::ExprBlock> = HashMap::with_capacity(enum_variants.len());
    let mut match_arms: Vec<TokenStream> = Vec::with_capacity(enum_variants.len());

    for enum_variant in enum_variants {
        let variant_pascal_case = &enum_variant.ident;

        // Destructure fields of the enum variant
        // E.g. for the variant:
        // Emails { unread_emails: u64 }
        // Create the expression:
        // Self::Emails { unread_emails }
        let destructuring_pattern = match &enum_variant.fields {
            syn::Fields::Named(named_fields) => {
                let named_field_idents = named_fields.named.iter().map(|field| &field.ident);
                quote!(#variant_pascal_case { #(#named_field_idents),* })
            }
            syn::Fields::Unit => quote!(#variant_pascal_case),
            syn::Fields::Unnamed(unnamed_fields) => {
                let field_count = unnamed_fields.unnamed.iter().count();
                if field_count != 1 {
                    return Err(MacroError::Unsupported(UnsupportedError::UnnamedFields {
                        span: enum_variant.ident.clone(),
                        field_count,
                    }));
                }

                quote!(#variant_pascal_case(unnamed_field))
            }
        };

        let context = match &enum_variant.fields {
            syn::Fields::Named(named_fields) => Context {
                reference_kind: ReferenceKind::EnumField,
                valid_references: unique_named_fields(named_fields),
            },
            syn::Fields::Unit | syn::Fields::Unnamed(_) => Context {
                reference_kind: ReferenceKind::EnumField,
                valid_references: HashSet::new(),
            },
        };

        let arm_body = match &enum_variant.fields {
            syn::Fields::Named(_) | syn::Fields::Unit => {
                expr_for_message(&mut group, &enum_variant.ident, &context)?
            }
            syn::Fields::Unnamed(unnamed_fields) => {
                expr_for_unnamed_fields(unnamed_fields, &enum_variant.ident, &context)?
            }
        };
        match_arms.push(quote!(Self::#destructuring_pattern => { #arm_body }));
    }

    Ok(quote! {
        match self {
            #(#match_arms),*
        }
    })
}