use crate::core::*;

use crate::gensym;

use crate::token::{Index, Span};

use std::collections::HashMap;

pub fn expand<'a>(fields: &mut Vec<ModuleField<'a>>) {

    let mut expander = Expander::default();

    expander.process(fields);

}

#[derive(Default)]

pub(crate) struct Expander<'a> {

    // Maps used to "intern" types. These maps are populated as type annotations

    // are seen and inline type annotations use previously defined ones if

    // there's a match.

    func_type_to_idx: HashMap<FuncKey<'a>, Index<'a>>,

    /// Fields, during processing, which should be prepended to the

    /// currently-being-processed field. This should always be empty after

    /// processing is complete.

    to_prepend: Vec<ModuleField<'a>>,

}

impl<'a> Expander<'a> {

    fn process(&mut self, fields: &mut Vec<ModuleField<'a>>) {

        // Next we expand "header" fields which are those like types and

        // imports. In this context "header" is defined by the previous

        // `process_imports_early` annotation.

        let mut cur = 0;

        while cur < fields.len() {

            self.expand_header(&mut fields[cur]);

            for item in self.to_prepend.drain(..) {

                fields.insert(cur, item);

                cur += 1;

            }

            cur += 1;

        }

        // Next after we've done that we expand remaining fields. Note that

        // after this we actually append instead of prepend. This is because

        // injected types are intended to come at the end of the type section

        // and types will be sorted before all other items processed here in the

        // final module anyway.

        for field in fields.iter_mut() {

            self.expand(field);

        }

        fields.append(&mut self.to_prepend);

    }

    fn expand_header(&mut self, item: &mut ModuleField<'a>) {

        match item {

            ModuleField::Type(ty) => {

                let id = gensym::fill(ty.span, &mut ty.id);

                match &mut ty.def {

                    TypeDef::Func(f) => {

                        f.key().insert(self, Index::Id(id));

                    }

                    TypeDef::Array(_) | TypeDef::Struct(_) => {}

                }

            }

            _ => {}

        }

    }

    fn expand(&mut self, item: &mut ModuleField<'a>) {

        match item {

            // This is pre-expanded above

            ModuleField::Type(_) => {}

            ModuleField::Rec(_) => {}

            ModuleField::Import(i) => {

                self.expand_item_sig(&mut i.item);

            }

            ModuleField::Func(f) => {

                self.expand_type_use(&mut f.ty);

                if let FuncKind::Inline { expression, .. } = &mut f.kind {

                    self.expand_expression(expression);

                }

            }

            ModuleField::Global(g) => {

                if let GlobalKind::Inline(expr) = &mut g.kind {

                    self.expand_expression(expr);

                }

            }

            ModuleField::Data(d) => {

                if let DataKind::Active { offset, .. } = &mut d.kind {

                    self.expand_expression(offset);

                }

            }

            ModuleField::Elem(e) => {

                if let ElemKind::Active { offset, .. } = &mut e.kind {

                    self.expand_expression(offset);

                }

                if let ElemPayload::Exprs { exprs, .. } = &mut e.payload {

                    for expr in exprs {

                        self.expand_expression(expr);

                    }

                }

            }

            ModuleField::Tag(t) => match &mut t.ty {

                TagType::Exception(ty) => {

                    self.expand_type_use(ty);

                }

            },

            ModuleField::Table(t) => match &mut t.kind {

                TableKind::Normal { init_expr, .. } => {

                    if let Some(expr) = init_expr {

                        self.expand_expression(expr);

                    }

                }

                TableKind::Import { .. } | TableKind::Inline { .. } => {}

            },

            ModuleField::Memory(_)

            | ModuleField::Start(_)

            | ModuleField::Export(_)

            | ModuleField::Custom(_) => {}

        }

    }

    fn expand_item_sig(&mut self, item: &mut ItemSig<'a>) {

        match &mut item.kind {

            ItemKind::Func(t) | ItemKind::Tag(TagType::Exception(t)) => {

                self.expand_type_use(t);

            }

            ItemKind::Global(_) | ItemKind::Table(_) | ItemKind::Memory(_) => {}

        }

    }

    fn expand_expression(&mut self, expr: &mut Expression<'a>) {

        for instr in expr.instrs.iter_mut() {

            self.expand_instr(instr);

        }

    }

    fn expand_instr(&mut self, instr: &mut Instruction<'a>) {

        match instr {

            Instruction::Block(bt)

            | Instruction::If(bt)

            | Instruction::Loop(bt)

            | Instruction::Let(LetType { block: bt, .. })

            | Instruction::Try(bt) => {

                // No expansion necessary, a type reference is already here.

                // We'll verify that it's the same as the inline type, if any,

                // later.

                if bt.ty.index.is_some() {

                    return;

                }

                match &bt.ty.inline {

                    // Only actually expand `TypeUse` with an index which appends a

                    // type if it looks like we need one. This way if the

                    // multi-value proposal isn't enabled and/or used we won't

                    // encode it.

                    Some(inline) => {

                        if inline.params.len() == 0 && inline.results.len() <= 1 {

                            return;

                        }

                    }

                    // If we didn't have either an index or an inline type

                    // listed then assume our block has no inputs/outputs, so

                    // fill in the inline type here.

                    //

                    // Do not fall through to expanding the `TypeUse` because

                    // this doesn't force an empty function type to go into the

                    // type section.

                    None => {

                        bt.ty.inline = Some(FunctionType::default());

                        return;

                    }

                }

                self.expand_type_use(&mut bt.ty);

            }

            Instruction::FuncBind(b) => {

                self.expand_type_use(&mut b.ty);

            }

            Instruction::CallIndirect(c) | Instruction::ReturnCallIndirect(c) => {

                self.expand_type_use(&mut c.ty);

            }

            _ => {}

        }

    }

    fn expand_type_use<T>(&mut self, item: &mut TypeUse<'a, T>) -> Index<'a>

    where

        T: TypeReference<'a>,

    {

        if let Some(idx) = &item.index {

            return *idx;

        }

        let key = match item.inline.as_mut() {

            Some(ty) => {

                ty.expand(self);

                ty.key()

            }

            None => T::default().key(),

        };

        let span = Span::from_offset(0); // FIXME(#613): don't manufacture

        let idx = self.key_to_idx(span, key);

        item.index = Some(idx);

        idx

    }

    fn key_to_idx(&mut self, span: Span, key: impl TypeKey<'a>) -> Index<'a> {

        // First see if this `key` already exists in the type definitions we've

        // seen so far...

        if let Some(idx) = key.lookup(self) {

            return idx;

        }

        // ... and failing that we insert a new type definition.

        let id = gensym::gen(span);

        self.to_prepend.push(ModuleField::Type(Type {

            span,

            id: Some(id),

            name: None,

            def: key.to_def(span),

            parent: None,

            final_type: None,

        }));

        let idx = Index::Id(id);

        key.insert(self, idx);

        idx

    }

}

pub(crate) trait TypeReference<'a>: Default {

    type Key: TypeKey<'a>;

    fn key(&self) -> Self::Key;

    fn expand(&mut self, cx: &mut Expander<'a>);

}

pub(crate) trait TypeKey<'a> {

    fn lookup(&self, cx: &Expander<'a>) -> Option<Index<'a>>;

    fn to_def(&self, span: Span) -> TypeDef<'a>;

    fn insert(&self, cx: &mut Expander<'a>, id: Index<'a>);

}

pub(crate) type FuncKey<'a> = (Box<[ValType<'a>]>, Box<[ValType<'a>]>);

impl<'a> TypeReference<'a> for FunctionType<'a> {

    type Key = FuncKey<'a>;

    fn key(&self) -> Self::Key {

        let params = self.params.iter().map(|p| p.2).collect();

        let results = self.results.clone();

        (params, results)

    }

    fn expand(&mut self, _cx: &mut Expander<'a>) {}

}

impl<'a> TypeKey<'a> for FuncKey<'a> {

    fn lookup(&self, cx: &Expander<'a>) -> Option<Index<'a>> {

        cx.func_type_to_idx.get(self).cloned()

    }

    fn to_def(&self, _span: Span) -> TypeDef<'a> {

        TypeDef::Func(FunctionType {

            params: self.0.iter().map(|t| (None, None, *t)).collect(),

            results: self.1.clone(),

        })

    }

    fn insert(&self, cx: &mut Expander<'a>, idx: Index<'a>) {

        cx.func_type_to_idx.entry(self.clone()).or_insert(idx);

    }

}