Experimenting with more structured ways to handle command-line input/output in Rust
use super::{ast, Error as FluentError, MessageContext, SourceFile};
use crate::macro_impl::derive;
use std::collections::HashMap;
use std::path::PathBuf;

use heck::ToKebabCase;
use icu_locale::Locale;
use miette::SourceSpan;
use proc_macro2::Ident;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum GroupError {
    #[error("error parsing Fluent code: {0}")]
    Fluent(#[from] FluentError),
    #[error("the field `{rust_name}` doesn't have a matching Fluent message `{fluent_name}`")]
    InvalidMessage {
        fluent_name: String,
        rust_name: String,
        span: proc_macro2::Span,
    },
    #[error("missing canonical locale `{canonical_locale}`")]
    MissingCanonicalLocale {
        canonical_locale: String,
        matched_locales: String,
    },
}

#[derive(Clone, Debug)]
pub struct Group {
    canonical_locale: Locale,
    canonical_resource: SourceFile,
    locales: HashMap<Locale, SourceFile>,
}

impl Group {
    pub fn new(
        canonical_locale: Locale,
        locale_paths: HashMap<Locale, PathBuf>,
    ) -> Result<Self, GroupError> {
        let mut locales = HashMap::with_capacity(locale_paths.len());

        for (locale, path) in locale_paths {
            let resource = SourceFile::new(path)?;

            // Insert resource (and make sure it's unique!)
            let previous_value = locales.insert(locale, resource);
            assert!(previous_value.is_none());
        }

        // Make sure the canonical locale exists
        let canonical_resource = if let Some(canonical_resource) = locales.remove(&canonical_locale)
        {
            canonical_resource
        } else {
            return Err(GroupError::MissingCanonicalLocale {
                canonical_locale: canonical_locale.to_string(),
                matched_locales: format!(
                    "the following locales were found:\n{}",
                    locales
                        .keys()
                        .map(|locale| format!("- {locale}"))
                        .collect::<Vec<_>>()
                        .join("\n")
                ),
            });
        };

        // Make sure canonical resource contains all message keys
        let canonical_keys: Vec<&str> = canonical_resource.message_ids().collect();
        for (locale, source_file) in &locales {
            // Find any keys in `locale_keys` missing from `canonical_keys`
            let mut unexpected_keys =
                source_file
                    .messages
                    .iter()
                    .filter(|(message_id, _message_value)| {
                        !canonical_keys.contains(&message_id.as_str())
                    });

            // TODO: group all unexpected keys into single error
            if let Some((unexpected_key, unexpected_message)) = unexpected_keys.next() {
                let source_code = source_file.named_source.clone();

                return Err(GroupError::Fluent(FluentError::UnexpectedKey {
                    unexpected_key: unexpected_key.to_string(),
                    locale: locale.to_string(),
                    canonical_locale: canonical_locale.to_string(),
                    source_code,
                    span: SourceSpan::from(unexpected_message.id.span.0.clone()),
                }));
            }
        }

        Ok(Self {
            canonical_locale,
            canonical_resource,
            locales,
        })
    }

    pub fn canonical_locale(&self) -> &Locale {
        &self.canonical_locale
    }

    pub fn remove_canonical_message(
        &mut self,
        ident: &Ident,
        derive_context: &derive::Context,
    ) -> Result<syn::Expr, GroupError> {
        let id = ident.to_string().to_kebab_case();
        let message = match self.canonical_resource.try_remove_expression(&id) {
            Some(canonical_message) => Ok(canonical_message),
            None => Err(GroupError::InvalidMessage {
                fluent_name: id.clone(),
                rust_name: ident.to_string(),
                span: ident.span(),
            }),
        }?;

        let message_context = MessageContext {
            source: &self.canonical_resource,
            // Any message where this value is `None` shouldn't be accessed directly, see:
            // https://docs.rs/fluent-syntax/latest/fluent_syntax/ast/struct.Attribute.html#example
            pattern: message.value.as_ref().unwrap(),
            derive_context,
        };

        let expression = ast::message_body(message_context)?;
        Ok(expression)
    }

    pub fn remove_additional_messages(
        &mut self,
        ident: &Ident,
        derive_context: &derive::Context,
    ) -> Result<HashMap<&Locale, syn::Expr>, GroupError> {
        // Create a message for every additional locale
        let mut messages = HashMap::with_capacity(self.locales.len());
        let id = ident.to_string().to_kebab_case();

        for (locale, compiled_messages) in self.locales.iter_mut() {
            // Include the message only if it exists in this locale
            if let Some(message) = compiled_messages.try_remove_expression(&id) {
                let expression = ast::message_body(MessageContext {
                    source: &compiled_messages,
                    pattern: message.value.as_ref().unwrap(),
                    derive_context,
                })?;

                // Insert this message (and make sure it's unique!)
                let previous_value = messages.insert(locale, expression);
                assert!(previous_value.is_none());
            }
        }

        Ok(messages)
    }

    pub fn locales_for_message<'a>(&'a self, id: &'a str) -> impl Iterator<Item = &'a Locale> {
        self.locales
            .iter()
            .filter(|(_locale, messages)| messages.contains_expression(id))
            .map(|(locale, _source_file)| locale)
    }
}