// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: 2020 The Gleam contributors

use super::*;
use crate::{
    analyse::name::check_name_case,
    ast::{
        Layer, TypeAst, TypeAstConstructor, TypeAstConstructorName, TypeAstFn, TypeAstHole,
        TypeAstTuple, TypeAstVar,
    },
    reference::ReferenceKind,
};
use std::sync::Arc;

use im::hashmap;

/// The Hydrator takes an AST representing a type (i.e. a type annotation
/// for a function argument) and returns a Type for that annotation.
///
/// If a valid Type cannot be constructed it returns an error.
///
/// It keeps track of any type variables created. This is useful for:
///
/// - Determining if a generic type variable should be made into an
///   unbound type variable during type instantiation.
/// - Ensuring that the same type is constructed if the programmer
///   uses the same name for a type variable multiple times.
///
#[derive(Debug)]
pub struct Hydrator {
    created_type_variables: im::HashMap<EcoString, CreatedTypeVariable>,
    /// A rigid type is a generic type that was specified as being generic in
    /// an annotation. As such it should never be instantiated into an unbound
    /// variable. This type_id => name map is used for reporting the original
    /// annotated name on error.
    rigid_type_names: im::HashMap<u64, EcoString>,
    permit_new_type_variables: bool,
    permit_holes: bool,
}

#[derive(Debug)]
pub struct ScopeResetData {
    created_type_variables: im::HashMap<EcoString, CreatedTypeVariable>,
    rigid_type_names: im::HashMap<u64, EcoString>,
}

impl Default for Hydrator {
    fn default() -> Self {
        Self::new()
    }
}

impl Hydrator {
    pub fn new() -> Self {
        Self {
            created_type_variables: hashmap![],
            rigid_type_names: hashmap![],
            permit_new_type_variables: true,
            permit_holes: false,
        }
    }

    pub fn named_type_variables(&self) -> im::HashMap<EcoString, CreatedTypeVariable> {
        self.created_type_variables.clone()
    }

    pub fn open_new_scope(&mut self) -> ScopeResetData {
        let created_type_variables = self.created_type_variables.clone();
        let rigid_type_names = self.rigid_type_names.clone();
        ScopeResetData {
            created_type_variables,
            rigid_type_names,
        }
    }

    pub fn close_scope(&mut self, data: ScopeResetData) {
        self.created_type_variables = data.created_type_variables;
        self.rigid_type_names = data.rigid_type_names;
    }

    pub fn disallow_new_type_variables(&mut self) {
        self.permit_new_type_variables = false
    }

    pub fn permit_holes(&mut self, flag: bool) {
        self.permit_holes = flag
    }

    /// A rigid type is a generic type that was specified as being generic in
    /// an annotation. As such it should never be instantiated into an unbound
    /// variable.
    pub fn is_rigid(&self, id: &u64) -> bool {
        self.rigid_type_names.contains_key(id)
    }

    pub fn rigid_names(&self) -> im::HashMap<u64, EcoString> {
        self.rigid_type_names.clone()
    }

    pub fn type_from_option_ast(
        &mut self,
        ast: &Option<TypeAst>,
        environment: &mut Environment<'_>,
        problems: &mut Problems,
    ) -> Result<Arc<Type>, Error> {
        match ast {
            Some(ast) => self.type_from_ast(ast, environment, problems),
            None => Ok(environment.new_unbound_var()),
        }
    }

    /// Construct a Type from an AST Type annotation.
    ///
    pub fn type_from_ast(
        &mut self,
        ast: &TypeAst,
        environment: &mut Environment<'_>,
        problems: &mut Problems,
    ) -> Result<Arc<Type>, Error> {
        match ast {
            TypeAst::Constructor(TypeAstConstructor {
                location,
                name,
                arguments,
                start_parentheses,
            }) => {
                // Hydrate the type argument AST into types
                let mut argument_types = Vec::with_capacity(arguments.len());
                for argument in arguments {
                    let type_ = self.type_from_ast(argument, environment, problems)?;
                    argument_types.push((argument.location(), type_));
                }

                let module = match name {
                    TypeAstConstructorName::Unqualified { .. } => None,
                    TypeAstConstructorName::Qualified {
                        module,
                        module_location,
                        ..
                    } => Some((module.clone(), *module_location)),
                };

                // If the constructor has been typed incorrectly with no name,
                // we can't figure out what type it's supposed to be, so it's
                // going to be an error: `wibble.` <- the name is missing!
                let (name, name_location) = match name {
                    TypeAstConstructorName::Unqualified { name, location }
                    | TypeAstConstructorName::Qualified {
                        name: Some((name, location)),
                        ..
                    } => (name, location),
                    TypeAstConstructorName::Qualified {
                        name: None,
                        module_location,
                        dot_location,
                        ..
                    } => {
                        return Err(Error::QualifiedTypeMissingName {
                            location: SrcSpan::new(module_location.start, *dot_location),
                        });
                    }
                };

                // Look up the constructor
                let TypeConstructor {
                    parameters,
                    type_: return_type,
                    deprecation,
                    ..
                } = environment
                    .get_type_constructor(&module, name)
                    .map_err(|error| {
                        let parameter_names = generate_parameter_names(arguments);
                        convert_get_type_constructor_error(
                            error,
                            location,
                            module.as_ref().map(|(_, location)| *location),
                            parameter_names,
                        )
                    })?
                    .clone();

                if let Some((type_module, type_name)) = return_type.named_type_name() {
                    let reference_kind = if let Some((module_alias, module_location)) = &module {
                        ReferenceKind::Qualified {
                            module_alias: module_alias.clone(),
                            module_location: *module_location,
                        }
                    } else if name != &type_name {
                        ReferenceKind::Alias
                    } else {
                        ReferenceKind::Unqualified
                    };
                    environment.references.register_type_reference(
                        type_module,
                        type_name,
                        name,
                        *name_location,
                        reference_kind,
                    );
                } else {
                    environment
                        .references
                        .register_type_reference_in_call_graph(name.clone());
                }

                match deprecation {
                    Deprecation::NotDeprecated => {}
                    Deprecation::Deprecated { message } => {
                        problems.warning(Warning::DeprecatedItem {
                            location: *location,
                            message,
                            layer: Layer::Type,
                        })
                    }
                }

                // Register the type constructor as being used if it is unqualified.
                // We do not track use of qualified type constructors as they may be
                // used in another module.
                if module.is_none() {
                    environment.increment_usage(name);
                }

                // Ensure that the correct number of arguments have been given
                // to the constructor.
                //
                // This is a special case for when a type is being called as a
                // type constructor. For example: `Int()` or `Bool(a, b)`
                if let Some(start_parentheses) = start_parentheses
                    && parameters.is_empty()
                {
                    return Err(Error::TypeUsedAsAConstructor {
                        location: SrcSpan::new(*start_parentheses, location.end),
                        name: name.clone(),
                    });
                } else if arguments.len() != parameters.len() {
                    return Err(Error::IncorrectTypeArity {
                        location: *location,
                        name: name.clone(),
                        expected: parameters.len(),
                        given: arguments.len(),
                    });
                }

                // Instantiate the constructor type for this specific usage
                let mut type_vars = hashmap![];
                #[allow(clippy::needless_collect)] // Not needless, used for side effects
                let parameter_types: Vec<_> = parameters
                    .into_iter()
                    .map(|type_| environment.instantiate(type_, &mut type_vars, self))
                    .collect();

                let return_type = environment.instantiate(return_type, &mut type_vars, self);

                // Unify argument types with instantiated parameter types so that the correct types
                // are inserted into the return type
                for (parameter, (location, argument)) in
                    parameter_types.into_iter().zip(argument_types)
                {
                    unify(parameter, argument).map_err(|e| convert_unify_error(e, location))?;
                }

                Ok(return_type)
            }

            TypeAst::Tuple(TypeAstTuple { elements, .. }) => Ok(tuple(
                elements
                    .iter()
                    .map(|type_| self.type_from_ast(type_, environment, problems))
                    .try_collect()?,
            )),

            TypeAst::Fn(TypeAstFn {
                arguments, return_, ..
            }) => {
                let arguments = arguments
                    .iter()
                    .map(|type_| self.type_from_ast(type_, environment, problems))
                    .try_collect()?;
                let return_ = self.type_from_ast(return_, environment, problems)?;
                Ok(fn_(arguments, return_))
            }

            TypeAst::Var(TypeAstVar { name, location }) => {
                match self.created_type_variables.get_mut(name) {
                    Some(var) => {
                        var.usage_count += 1;
                        Ok(var.type_.clone())
                    }

                    None if self.permit_new_type_variables => {
                        if let Err(error) = check_name_case(*location, name, Named::TypeVariable) {
                            problems.error(error);
                        }
                        let t = environment.new_generic_var();
                        let _ = self
                            .rigid_type_names
                            .insert(environment.previous_uid(), name.clone());
                        environment
                            .names
                            .type_variable_in_scope(environment.previous_uid(), name.clone());
                        let _ = self.created_type_variables.insert(
                            name.clone(),
                            CreatedTypeVariable {
                                type_: t.clone(),
                                usage_count: 1,
                            },
                        );
                        Ok(t)
                    }

                    None => {
                        let hint = match environment.scope.contains_key(name) {
                            true => UnknownTypeHint::ValueInScopeWithSameName,
                            false => UnknownTypeHint::AlternativeTypes(
                                environment.module_types.keys().cloned().collect(),
                            ),
                        };

                        Err(Error::UnknownType {
                            name: name.clone(),
                            location: *location,
                            hint,
                            parameter_names: vec![],
                        })
                    }
                }
            }

            TypeAst::Hole(TypeAstHole { .. }) if self.permit_holes => {
                Ok(environment.new_unbound_var())
            }

            TypeAst::Hole(TypeAstHole { location, .. }) => Err(Error::UnexpectedTypeHole {
                location: *location,
            }),
        }
    }

    pub fn clear_ridgid_type_names(&mut self) {
        self.rigid_type_names.clear();
    }

    /// All the type variables that were created but never used.
    pub fn unused_type_variables(&self) -> impl Iterator<Item = &EcoString> {
        self.created_type_variables
            .iter()
            .filter(|(_, var)| var.usage_count == 0)
            .map(|(name, _)| name)
    }

    /// Create a new type variable with the given name.
    pub fn add_type_variable(
        &mut self,
        name: &EcoString,
        environment: &mut Environment<'_>,
    ) -> Result<Arc<Type>, Arc<Type>> {
        let t = environment.new_generic_var();
        let v = CreatedTypeVariable {
            type_: t.clone(),
            usage_count: 0,
        };

        environment
            .names
            .type_variable_in_scope(environment.previous_uid(), name.clone());
        match self.created_type_variables.insert(name.clone(), v) {
            Some(_) => Err(t),
            None => Ok(t),
        }
    }
}

#[derive(Debug, Clone)]
pub struct CreatedTypeVariable {
    pub type_: Arc<Type>,
    pub usage_count: usize,
}

/// Names for the parameters of an unknown type, used to suggest a definition.
/// Type variables keep their name; other parameters get a generated one.
fn generate_parameter_names(arguments: &[TypeAst]) -> Vec<EcoString> {
    // Reserve the type variable names so generated ones don't collide with them.
    let mut used: HashSet<EcoString> = arguments
        .iter()
        .filter_map(|argument| match argument {
            TypeAst::Var(TypeAstVar { name, .. }) => Some(name.clone()),
            TypeAst::Constructor(_) | TypeAst::Fn(_) | TypeAst::Tuple(_) | TypeAst::Hole(_) => None,
        })
        .collect();

    let mut next = 0;
    arguments
        .iter()
        .map(|argument| match argument {
            TypeAst::Var(TypeAstVar { name, .. }) => name.clone(),
            TypeAst::Constructor(_) | TypeAst::Fn(_) | TypeAst::Tuple(_) | TypeAst::Hole(_) => {
                loop {
                    let candidate = type_parameter_name(next);
                    next += 1;
                    if used.insert(candidate.clone()) {
                        break candidate;
                    }
                }
            }
        })
        .collect()
}

/// `0 -> "a"`, `25 -> "z"`, `26 -> "aa"`, ...
fn type_parameter_name(index: usize) -> EcoString {
    let alphabet_length = 26;
    let char_offset = 97;
    let mut chars = vec![];
    let mut rest = index;

    loop {
        let n = rest % alphabet_length;
        rest /= alphabet_length;
        chars.push((n as u8 + char_offset) as char);

        if rest == 0 {
            break;
        }
        rest -= 1;
    }

    chars.into_iter().rev().collect()
}
