Framework for embedding localizations into Rust types
use std::cell::OnceCell;

use super::{Error, FunctionError, MessageContext};
use crate::macro_impl::derive::ReferenceKind;

use fixed_decimal::{Decimal, FloatPrecision};
use fluent_syntax::ast::{
    Expression, InlineExpression, Pattern, PatternElement, Variant, VariantKey,
};
use heck::{ToLowerCamelCase, ToPascalCase, ToSnakeCase};
use proc_macro2::{Literal, TokenStream};
use quote::{ToTokens, format_ident, quote};
use syn::parse_quote;

#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum VariantType {
    Integer,
    Plural,
}

struct Select<'variants> {
    variant_type: VariantType,
    default_pattern: TokenStream,
    default_expression: &'variants Pattern<String>,
    match_arms: Vec<(TokenStream, &'variants Pattern<String>)>,
}

pub fn message_body(message_context: MessageContext) -> Result<syn::Expr, Error> {
    let mut write_expressions: Vec<syn::Expr> =
        Vec::with_capacity(message_context.pattern.elements.len());

    for element in &message_context.pattern.elements {
        match element {
            PatternElement::TextElement { value } => {
                write_expressions.push(parse_quote!(buffer.push_str(#value)));
            }
            PatternElement::Placeable { expression } => {
                write_expressions.push(match expression {
                    Expression::Select { selector, variants } => {
                        let select = select(&message_context, variants)?;

                        let raw_match_target = inline_expression(selector, message_context)?;
                        let match_target = match select.variant_type {
                            VariantType::Integer { .. } => quote!(#raw_match_target),
                            VariantType::Plural { .. } => {
                                let category_reference = match selector {
                                    // Can't dereference value returned by function
                                    InlineExpression::FunctionReference { .. } => {
                                        quote!(#raw_match_target)
                                    }
                                    // Only dereference if enum variant
                                    _ => match message_context.derive_context.reference_kind {
                                        ReferenceKind::EnumField => quote!(*#raw_match_target),
                                        ReferenceKind::StructField => quote!(#raw_match_target),
                                    },
                                };
                                quote!(plural_rules.category_for(#category_reference))
                            }
                        };

                        let default_match_pattern = select.default_pattern;
                        let default_match_expression = message_body(MessageContext {
                            pattern: select.default_expression,
                            ..message_context
                        })?;

                        let match_patterns = select
                            .match_arms
                            .iter()
                            .map(|(pattern, _expression)| pattern)
                            .collect::<Vec<&TokenStream>>();
                        let match_expressions = select
                            .match_arms
                            .iter()
                            .map(|(_match_pattern, expression)| {
                                message_body(MessageContext {
                                    pattern: expression,
                                    ..message_context
                                })
                            })
                            .collect::<Result<Vec<syn::Expr>, Error>>()?;

                        parse_quote! {
                            match #match_target {
                                #(#match_patterns => #match_expressions),*
                                #default_match_pattern => #default_match_expression,
                            }
                        }
                    }
                    Expression::Inline(expression) => {
                        let expression = inline_expression(expression, message_context)?;
                        parse_quote!(buffer.push_str(&#expression.localize_for(locale)))
                    }
                });
            }
        }
    }

    Ok(parse_quote! {
        {
            #(#write_expressions;)*
        }
    })
}

fn select<'variants>(
    message_context: &MessageContext,
    variants: &'variants Vec<Variant<String>>,
) -> Result<Select<'variants>, Error> {
    let expected_variant_type = OnceCell::new();

    for variant in variants {
        let variant_type = match &variant.key {
            VariantKey::Identifier { name } => match name.as_str() {
                // "other" is ambiguous as it could be used for plurals or as the default integer case
                "other" => None,
                _ => Some(VariantType::Plural),
            },
            VariantKey::NumberLiteral { .. } => Some(VariantType::Integer),
        };

        // Make sure all variants are of the same type
        if let Some(variant_kind) = variant_type {
            let expected_type = expected_variant_type.get_or_init(|| variant_kind);
            if *expected_type != variant_kind {
                todo!()
            }
        }
    }

    match expected_variant_type.get().unwrap() {
        VariantType::Integer => select_integers(message_context, variants),
        VariantType::Plural => select_plurals(message_context, variants),
    }
}

fn select_integers<'variants>(
    message_context: &MessageContext,
    variants: &'variants Vec<Variant<String>>,
) -> Result<Select<'variants>, Error> {
    let mut default_expression = OnceCell::new();
    let mut match_arms = Vec::new();

    for variant in variants {
        match &variant.key {
            VariantKey::Identifier { name } => match name.to_lowercase().as_str() {
                "other" => default_expression
                    .set(&variant.value)
                    .expect("Multiple default variants"),
                _ => {
                    return Err(Error::InvalidIdentifierType {
                        invalid_identifier: name.clone(),
                        expected_type: VariantType::Integer,
                        found_type: VariantType::Plural,
                        source_code: message_context.source.named_source.clone(),
                    });
                }
            },
            VariantKey::NumberLiteral { value } => {
                let parsed_integer: u128 = match value.parse() {
                    Ok(integer) => integer,
                    Err(parse_error) => {
                        return Err(Error::InvalidNumberLiteral {
                            invalid_literal: value.clone(),
                            parse_error,
                            source_code: message_context.source.named_source.clone(),
                        });
                    }
                };
                let integer_literal = Literal::u128_unsuffixed(parsed_integer);

                match_arms.push((integer_literal.to_token_stream(), &variant.value));
            }
        }
    }

    let default_expression = default_expression.take().expect("No default variant");
    Ok(Select {
        variant_type: VariantType::Integer,
        default_pattern: quote!(_),
        default_expression,
        match_arms,
    })
}

fn select_plurals<'variants>(
    message_context: &MessageContext,
    variants: &'variants Vec<Variant<String>>,
) -> Result<Select<'variants>, Error> {
    let mut default_arm = OnceCell::new();
    let mut match_arms = Vec::new();

    for variant in variants {
        match &variant.key {
            VariantKey::Identifier { name } => {
                let name_pascal_case = name.to_pascal_case();
                let ident = syn::Ident::new(&name_pascal_case, proc_macro2::Span::call_site());
                let pattern =
                    quote!(::l10n_embed::macro_prelude::icu_plurals::PluralCategory::#ident);

                match variant.default {
                    // The default arm should include any enum variants not matched
                    true => default_arm
                        .set((quote!(#pattern | _), &variant.value))
                        .expect("Multiple default variants"),
                    false => match_arms.push((pattern, &variant.value)),
                }
            }
            VariantKey::NumberLiteral { value } => {
                return Err(Error::InvalidIdentifierType {
                    invalid_identifier: value.clone(),
                    expected_type: VariantType::Plural,
                    found_type: VariantType::Integer,
                    source_code: message_context.source.named_source.clone(),
                });
            }
        }
    }

    let (default_pattern, default_expression) = default_arm.take().expect("No default variant");
    Ok(Select {
        variant_type: VariantType::Plural,
        default_pattern,
        default_expression,
        match_arms,
    })
}

fn inline_expression(
    expression: &InlineExpression<String>,
    message_context: MessageContext,
) -> Result<syn::Expr, Error> {
    Ok(match expression {
        InlineExpression::StringLiteral { value } => {
            let byte_string_literal = proc_macro2::Literal::byte_string(value.as_bytes());
            parse_quote!(#byte_string_literal)
        }
        InlineExpression::NumberLiteral { value } => {
            let parsed_value: f64 = value.parse().unwrap();
            // Check validity at compile-time, so we avoid generating code that will break at runtime
            assert!(Decimal::try_from_f64(parsed_value, FloatPrecision::RoundTrip).is_ok());

            let float_literal = proc_macro2::Literal::f64_suffixed(parsed_value);
            parse_quote!(
                ::l10n_embed::macro_prelude::fixed_decimal::Decimal::try_from_f64(
                    #float_literal,
                    ::l10n_embed::macro_prelude::fixed_decimal::FloatPrecision::RoundTrip,
                )
                .unwrap();
            )
        }
        InlineExpression::VariableReference { id } => {
            // Make sure the referenced variable is in the set of valid variables
            let fluent_name = &id.name;
            let rust_name = id.name.to_snake_case();

            let ident = if let Some(variable) = message_context
                .derive_context
                .valid_references
                .get(&rust_name)
            {
                format_ident!("{variable}")
            } else {
                let source_code = message_context.source.named_source.clone();

                return Err(Error::InvalidReference {
                    fluent_name: fluent_name.clone(),
                    rust_name,
                    source_code,
                    valid_references: format!(
                        "the following references are valid:\n{}",
                        message_context
                            .derive_context
                            .valid_references
                            .iter()
                            .map(|field| format!("- ${}", field.to_lower_camel_case()))
                            .collect::<Vec<String>>()
                            .join("\n")
                    ),
                });
            };

            match message_context.derive_context.reference_kind {
                ReferenceKind::EnumField => parse_quote!(#ident),
                ReferenceKind::StructField => parse_quote!(self.#ident),
            }
        }
        InlineExpression::FunctionReference { id, arguments } => {
            match id.name.to_snake_case().as_str() {
                "len" => {
                    if !arguments.named.is_empty() {
                        return Err(Error::Function(FunctionError::UnexpectedNamedArguments {
                            source_code: message_context.source.named_source.clone(),
                        }));
                    }

                    let argument = match arguments.positional.as_slice() {
                        [single_argument] => single_argument,
                        _ => {
                            return Err(Error::Function(
                                FunctionError::IncorrectPositionalArgumentCount {
                                    expected_len: 1,
                                    actual_len: arguments.positional.len(),
                                    source_code: message_context.source.named_source.clone(),
                                },
                            ));
                        }
                    };

                    let variable_reference = match argument {
                        InlineExpression::VariableReference { .. } => {
                            inline_expression(argument, message_context)?
                        }
                        _ => {
                            return Err(Error::Function(FunctionError::NonVariableReference {
                                source_code: message_context.source.named_source.clone(),
                            }));
                        }
                    };

                    let vec_reference = match message_context.derive_context.reference_kind {
                        ReferenceKind::EnumField => quote!(#variable_reference),
                        ReferenceKind::StructField => quote!(&#variable_reference),
                    };

                    parse_quote!(Vec::len(#vec_reference))
                }
                _ => {
                    return Err(Error::Function(FunctionError::InvalidName {
                        name: id.name.clone(),
                        source_code: message_context.source.named_source.clone(),
                    }));
                }
            }
        }
        _ => {
            dbg!(expression);
            todo!()
        }
    })
}