use std::collections::HashMap;

use wit_parser::*;

#[derive(Default)]
pub struct Types {
    type_info: HashMap<TypeId, TypeInfo>,
    equal_types: UnionFind,
}

#[derive(Default, Clone, Copy, Debug)]
pub struct TypeInfo {
    /// Whether this type is ever used (transitively) within the
    /// parameter of an imported function.
    ///
    /// This means that it's used in a context where ownership isn't
    /// relinquished.
    pub borrowed: bool,

    /// Whether this type is ever used (transitively) within the
    /// parameter or result of an export, or the result of an import.
    ///
    /// This means that it's used in a context where ownership is required and
    /// memory management is necessary.
    pub owned: bool,

    /// Whether this type is ever used (transitively) within the
    /// error case in the result of a function.
    pub error: bool,

    /// Whether this type (transitively) has a list (or string).
    pub has_list: bool,

    /// Whether this type (transitively) has a tuple.
    pub has_tuple: bool,

    /// Whether this type (transitively) has a resource (or handle).
    pub has_resource: bool,

    /// Whether this type (transitively) has a borrow handle.
    pub has_borrow_handle: bool,

    /// Whether this type (transitively) has an own handle.
    pub has_own_handle: bool,
}

impl std::ops::BitOrAssign for TypeInfo {
    fn bitor_assign(&mut self, rhs: Self) {
        self.borrowed |= rhs.borrowed;
        self.owned |= rhs.owned;
        self.error |= rhs.error;
        self.has_list |= rhs.has_list;
        self.has_tuple |= rhs.has_tuple;
        self.has_resource |= rhs.has_resource;
        self.has_borrow_handle |= rhs.has_borrow_handle;
        self.has_own_handle |= rhs.has_own_handle;
    }
}

impl TypeInfo {
    pub fn is_clone(&self) -> bool {
        !self.has_resource
    }
    pub fn is_copy(&self) -> bool {
        !self.has_list && !self.has_resource
    }
}

impl Types {
    pub fn analyze(&mut self, resolve: &Resolve) {
        for (t, _) in resolve.types.iter() {
            self.type_id_info(resolve, t);
        }
        for (_, world) in resolve.worlds.iter() {
            for (import, (_, item)) in world
                .imports
                .iter()
                .map(|i| (true, i))
                .chain(world.exports.iter().map(|i| (false, i)))
            {
                match item {
                    WorldItem::Function(f) => {
                        self.type_info_func(resolve, f, import);
                    }
                    WorldItem::Interface {
                        id, stability: _, ..
                    } => {
                        for (_, f) in resolve.interfaces[*id].functions.iter() {
                            self.type_info_func(resolve, f, import);
                        }
                    }
                    WorldItem::Type { .. } => {}
                }
            }
        }
    }

    /// Identify all the structurally equal types reachable from the `world_id`.
    ///
    /// The `may_alias_another_type` closure is used to determine whether the
    /// language's definition of the provided `TypeId` might possibly alias
    /// some other type in a language. This is a language-specific deduction
    /// which can also be affected by options to a binding generator. If a type
    /// can't be aliased by anything else then it can't be considered equal to
    /// anything else. Note that in this situations other types may still
    /// be equal to it, such as aliased types at the WIT level (e.g. `type foo
    /// = some-record`).
    pub fn collect_equal_types(
        &mut self,
        resolve: &Resolve,
        world_id: WorldId,
        may_alias_another_type: &dyn Fn(TypeId) -> bool,
    ) {
        let mut live_types = wit_parser::LiveTypes::default();
        live_types.add_world(resolve, world_id);
        for (i, ty) in live_types.iter().enumerate() {
            if !may_alias_another_type(ty) {
                continue;
            }
            // TODO: we could define a hash function for TypeDefKind to prevent the inner loop.
            for earlier in live_types.iter().take(i) {
                if self.equal_types.find(ty) == self.equal_types.find(earlier) {
                    continue;
                }
                // The correctness of is_structurally_equal relies on the fact
                // that live_types.iter() is in topological order.
                if self.is_structurally_equal(resolve, ty, earlier) {
                    self.equal_types.union(ty, earlier);
                    break;
                }
            }
        }

        // Merge TypeInfo for structurally equal types so that all members
        // of an equivalence class carry the union of all usage flags
        // (borrowed, owned, etc.).
        let mut merged: HashMap<TypeId, TypeInfo> = HashMap::new();
        for (&id, &info) in &self.type_info {
            let rep = self.equal_types.find(id);
            *merged.entry(rep).or_default() |= info;
        }
        for (&id, info) in &mut self.type_info {
            let rep = self.equal_types.find(id);
            if let Some(&merged_info) = merged.get(&rep) {
                *info = merged_info;
            }
        }
    }

    fn type_info_func(&mut self, resolve: &Resolve, func: &Function, import: bool) {
        let mut live = LiveTypes::default();
        for Param { ty, .. } in func.params.iter() {
            self.type_info(resolve, ty);
            live.add_type(resolve, ty);
        }
        for id in live.iter() {
            if resolve.types[id].name.is_some() {
                let info = self.type_info.get_mut(&id).unwrap();
                if import {
                    info.borrowed = true;
                } else {
                    info.owned = true;
                }
            }
        }
        let mut live = LiveTypes::default();
        if let Some(ty) = &func.result {
            self.type_info(resolve, ty);
            live.add_type(resolve, ty);
        }
        for id in live.iter() {
            if resolve.types[id].name.is_some() {
                self.type_info.get_mut(&id).unwrap().owned = true;
            }
        }

        if let Some(Type::Id(id)) = func.result {
            let err = match &resolve.types[id].kind {
                TypeDefKind::Result(Result_ { err, .. }) => err,
                _ => return,
            };
            if let Some(Type::Id(id)) = err {
                // When an interface `use`s a type from another interface, it creates a new typeid
                // referring to the definition typeid. Chase any chain of references down to the
                // typeid of the definition.
                fn resolve_type_definition_id(resolve: &Resolve, mut id: TypeId) -> TypeId {
                    loop {
                        match resolve.types[id].kind {
                            TypeDefKind::Type(Type::Id(def_id)) => id = def_id,
                            _ => return id,
                        }
                    }
                }
                let id = resolve_type_definition_id(resolve, *id);
                self.type_info.get_mut(&id).unwrap().error = true;
            }
        }
    }

    pub fn get(&self, id: TypeId) -> TypeInfo {
        self.type_info[&id]
    }

    pub fn type_id_info(&mut self, resolve: &Resolve, ty: TypeId) -> TypeInfo {
        if let Some(info) = self.type_info.get(&ty) {
            return *info;
        }
        let mut info = TypeInfo::default();
        match &resolve.types[ty].kind {
            TypeDefKind::Record(r) => {
                for field in r.fields.iter() {
                    info |= self.type_info(resolve, &field.ty);
                }
            }
            TypeDefKind::Resource => {
                info.has_resource = true;
            }
            TypeDefKind::Handle(handle) => {
                match handle {
                    Handle::Borrow(_) => info.has_borrow_handle = true,
                    Handle::Own(_) => info.has_own_handle = true,
                }
                info.has_resource = true;
            }
            TypeDefKind::Tuple(t) => {
                for ty in t.types.iter() {
                    info |= self.type_info(resolve, ty);
                }
                info.has_tuple = true;
            }
            TypeDefKind::Flags(_) => {}
            TypeDefKind::Enum(_) => {}
            TypeDefKind::Variant(v) => {
                for case in v.cases.iter() {
                    info |= self.optional_type_info(resolve, case.ty.as_ref());
                }
            }
            TypeDefKind::List(ty) => {
                info = self.type_info(resolve, ty);
                info.has_list = true;
            }
            TypeDefKind::Type(ty) => {
                info = self.type_info(resolve, ty);
            }
            TypeDefKind::Option(ty) => {
                info = self.type_info(resolve, ty);
            }
            TypeDefKind::Result(r) => {
                info = self.optional_type_info(resolve, r.ok.as_ref());
                info |= self.optional_type_info(resolve, r.err.as_ref());
            }
            TypeDefKind::Future(_) | TypeDefKind::Stream(_) => {
                // These are all u32 handles regardless of payload type, so no
                // need to recurse.
                info.has_resource = true;

                // Flag these as being onwed handles as lowering these values
                // should use the same ownership semantics as `own<T>`
                info.has_own_handle = true;
            }
            TypeDefKind::FixedLengthList(ty, _) => {
                info = self.type_info(resolve, ty);
            }
            TypeDefKind::Map(key, value) => {
                info = self.type_info(resolve, key);
                info |= self.type_info(resolve, value);
                info.has_list = true;
            }
            TypeDefKind::Unknown => unreachable!(),
        }
        let prev = self.type_info.insert(ty, info);
        assert!(prev.is_none());
        info
    }

    pub fn type_info(&mut self, resolve: &Resolve, ty: &Type) -> TypeInfo {
        let mut info = TypeInfo::default();
        match ty {
            Type::String => info.has_list = true,
            Type::ErrorContext => info.has_resource = true,
            Type::Id(id) => return self.type_id_info(resolve, *id),
            _ => {}
        }
        info
    }

    fn optional_type_info(&mut self, resolve: &Resolve, ty: Option<&Type>) -> TypeInfo {
        match ty {
            Some(ty) => self.type_info(resolve, ty),
            None => TypeInfo::default(),
        }
    }

    fn is_structurally_equal(&mut self, resolve: &Resolve, a: TypeId, b: TypeId) -> bool {
        let a_def = &resolve.types[a].kind;
        let b_def = &resolve.types[b].kind;
        if self.equal_types.find(a) == self.equal_types.find(b) {
            return true;
        }
        match (a_def, b_def) {
            // Peel off typedef layers and continue recursing.
            (TypeDefKind::Type(a), _) => self.type_id_equal_to_type(resolve, b, a),
            (_, TypeDefKind::Type(b)) => self.type_id_equal_to_type(resolve, a, b),

            (TypeDefKind::Record(ra), TypeDefKind::Record(rb)) => {
                ra.fields.len() == rb.fields.len()
                  // Fields are ordered in WIT, so record {a: T, b: U} is different from {b: U, a: T}
                  && ra.fields.iter().zip(rb.fields.iter()).all(|(fa, fb)| {
                      fa.name == fb.name && self.types_equal(resolve, &fa.ty, &fb.ty)
                  })
            }
            (TypeDefKind::Record(_), _) => false,
            (TypeDefKind::Variant(va), TypeDefKind::Variant(vb)) => {
                va.cases.len() == vb.cases.len()
                    && va.cases.iter().zip(vb.cases.iter()).all(|(ca, cb)| {
                        ca.name == cb.name && self.optional_types_equal(resolve, &ca.ty, &cb.ty)
                    })
            }
            (TypeDefKind::Variant(_), _) => false,
            (TypeDefKind::Enum(ea), TypeDefKind::Enum(eb)) => {
                ea.cases.len() == eb.cases.len()
                    && ea
                        .cases
                        .iter()
                        .zip(eb.cases.iter())
                        .all(|(ca, cb)| ca.name == cb.name)
            }
            (TypeDefKind::Enum(_), _) => false,
            (TypeDefKind::Flags(fa), TypeDefKind::Flags(fb)) => {
                fa.flags.len() == fb.flags.len()
                    && fa
                        .flags
                        .iter()
                        .zip(fb.flags.iter())
                        .all(|(fa, fb)| fa.name == fb.name)
            }
            (TypeDefKind::Flags(_), _) => false,
            (TypeDefKind::Tuple(ta), TypeDefKind::Tuple(tb)) => {
                ta.types.len() == tb.types.len()
                    && ta
                        .types
                        .iter()
                        .zip(tb.types.iter())
                        .all(|(a, b)| self.types_equal(resolve, a, b))
            }
            (TypeDefKind::Tuple(_), _) => false,
            (TypeDefKind::List(la), TypeDefKind::List(lb)) => self.types_equal(resolve, la, lb),
            (TypeDefKind::List(_), _) => false,
            (TypeDefKind::FixedLengthList(ta, sa), TypeDefKind::FixedLengthList(tb, sb)) => {
                sa == sb && self.types_equal(resolve, ta, tb)
            }
            (TypeDefKind::FixedLengthList(..), _) => false,
            (TypeDefKind::Option(oa), TypeDefKind::Option(ob)) => self.types_equal(resolve, oa, ob),
            (TypeDefKind::Option(_), _) => false,
            (TypeDefKind::Result(ra), TypeDefKind::Result(rb)) => {
                self.optional_types_equal(resolve, &ra.ok, &rb.ok)
                    && self.optional_types_equal(resolve, &ra.err, &rb.err)
            }
            (TypeDefKind::Result(_), _) => false,
            (TypeDefKind::Map(ak, av), TypeDefKind::Map(bk, bv)) => {
                self.types_equal(resolve, ak, bk) && self.types_equal(resolve, av, bv)
            }
            (TypeDefKind::Map(..), _) => false,
            (TypeDefKind::Future(a), TypeDefKind::Future(b)) => {
                self.optional_types_equal(resolve, a, b)
            }
            (TypeDefKind::Future(..), _) => false,
            (TypeDefKind::Stream(a), TypeDefKind::Stream(b)) => {
                self.optional_types_equal(resolve, a, b)
            }
            (TypeDefKind::Stream(..), _) => false,
            (TypeDefKind::Handle(a), TypeDefKind::Handle(b)) => match (a, b) {
                (Handle::Own(a), Handle::Own(b)) | (Handle::Borrow(a), Handle::Borrow(b)) => {
                    self.is_structurally_equal(resolve, *a, *b)
                }
                (Handle::Own(_) | Handle::Borrow(_), _) => false,
            },
            (TypeDefKind::Handle(_), _) => false,
            (TypeDefKind::Unknown, _) => unreachable!(),

            // Resources are only equal if their original ids are equal,
            // otherwise all resources are un-equal to each other.
            (TypeDefKind::Resource, TypeDefKind::Resource) => a == b,
            (TypeDefKind::Resource, _) => false,
        }
    }

    fn types_equal(&mut self, resolve: &Resolve, a: &Type, b: &Type) -> bool {
        match (a, b) {
            // Peel off typedef layers and continue recursing.
            (Type::Id(a), b) => self.type_id_equal_to_type(resolve, *a, b),
            (a, Type::Id(b)) => self.type_id_equal_to_type(resolve, *b, a),

            // When both a and b are primitives, they're only equal of
            // the primitives are the same.
            (
                Type::Bool
                | Type::U8
                | Type::S8
                | Type::U16
                | Type::S16
                | Type::U32
                | Type::S32
                | Type::U64
                | Type::S64
                | Type::F32
                | Type::F64
                | Type::Char
                | Type::String
                | Type::ErrorContext,
                _,
            ) => a == b,
        }
    }

    fn type_id_equal_to_type(&mut self, resolve: &Resolve, a: TypeId, b: &Type) -> bool {
        let ak = &resolve.types[a].kind;
        match (ak, b) {
            (TypeDefKind::Type(a), b) => self.types_equal(resolve, a, b),
            (_, Type::Id(b)) => self.is_structurally_equal(resolve, a, *b),

            // Type `a` isn't a typedef, and type `b` is a primitive, so it's no
            // longer possible for them to be equal.
            _ => false,
        }
    }

    fn optional_types_equal(
        &mut self,
        resolve: &Resolve,
        a: &Option<Type>,
        b: &Option<Type>,
    ) -> bool {
        match (a, b) {
            (Some(a), Some(b)) => self.types_equal(resolve, a, b),
            (Some(_), None) | (None, Some(_)) => false,
            (None, None) => true,
        }
    }

    pub fn get_representative_type(&mut self, id: TypeId) -> TypeId {
        self.equal_types.find(id)
    }
}

#[derive(Default)]
pub struct UnionFind {
    parent: HashMap<TypeId, TypeId>,
}
impl UnionFind {
    fn find(&mut self, id: TypeId) -> TypeId {
        // Path compression
        let parent = self.parent.get(&id).copied().unwrap_or(id);
        if parent != id {
            let root = self.find(parent);
            self.parent.insert(id, root);
            root
        } else {
            id
        }
    }
    fn union(&mut self, a: TypeId, b: TypeId) {
        let ra = self.find(a);
        let rb = self.find(b);
        if ra != rb {
            // Use smaller id as root for determinism
            if ra < rb {
                self.parent.insert(rb, ra);
            } else {
                self.parent.insert(ra, rb);
            }
        }
    }
}
