use std::collections::HashMap;

use id_arena::Arena;
use indexmap::IndexMap;
use log::trace;

use crate::{Alias, FunctionType, Pair, Struct, Type, TypeId, Union, stringify_impl};

pub fn substitute(arena: &mut Arena<Type>, id: TypeId) -> TypeId {
    substitute_with_mappings(arena, id, &HashMap::new())
}

pub fn substitute_with_mappings(
    arena: &mut Arena<Type>,
    id: TypeId,
    mappings: &HashMap<TypeId, TypeId>,
) -> TypeId {
    let name = stringify_impl(arena, id, &mut IndexMap::new());
    trace!("Substituting {name}");
    let result = substitute_impl(arena, id, mappings, &mut IndexMap::new());
    let new_name = stringify_impl(arena, result, &mut IndexMap::new());
    trace!("Substitution from {name} yielded {new_name}");
    result
}

fn substitute_impl(
    arena: &mut Arena<Type>,
    mut old_id: TypeId,
    map: &HashMap<TypeId, TypeId>,
    stack: &mut IndexMap<TypeId, TypeId>,
) -> TypeId {
    if let Some(id) = map.get(&old_id) {
        old_id = *id;
    }

    if let Some(id) = stack.get(&old_id) {
        return *id;
    }

    let new_id = arena.alloc(Type::Unresolved);
    stack.insert(old_id, new_id);

    let result = match arena[old_id].clone() {
        Type::Apply(apply) => {
            let mut new_map = apply.generics.clone();

            for arg in new_map.values_mut() {
                *arg = substitute_impl(arena, *arg, map, stack);
            }

            substitute_impl(arena, apply.inner, &new_map, stack)
        }
        Type::Unresolved | Type::Generic(_) | Type::Atom(_) | Type::Never | Type::Any => old_id,
        Type::Ref(id) => substitute_impl(arena, id, map, stack),
        Type::Pair(pair) => {
            let first = substitute_impl(arena, pair.first, map, stack);
            let rest = substitute_impl(arena, pair.rest, map, stack);
            arena.alloc(Type::Pair(Pair::new(first, rest)))
        }
        Type::Struct(ty) => {
            let inner = substitute_impl(arena, ty.inner, map, stack);
            arena.alloc(Type::Struct(Struct { inner, ..ty }))
        }
        Type::Alias(alias) => {
            let inner = substitute_impl(arena, alias.inner, map, stack);
            arena.alloc(Type::Alias(Alias { inner, ..alias }))
        }
        Type::Function(function) => {
            let params = function
                .params
                .iter()
                .map(|(name, id)| (name.clone(), substitute_impl(arena, *id, map, stack)))
                .collect();
            let ret = substitute_impl(arena, function.ret, map, stack);
            arena.alloc(Type::Function(FunctionType {
                params,
                ret,
                ..function
            }))
        }
        Type::Union(union) => {
            let types = union
                .types
                .iter()
                .map(|id| substitute_impl(arena, *id, map, stack))
                .collect();
            arena.alloc(Type::Union(Union::new(types)))
        }
    };

    stack.pop().unwrap();

    *arena.get_mut(new_id).unwrap() = Type::Ref(result);

    result
}
