use std::{
    collections::{hash_map::Entry, HashMap, HashSet},
    rc::Rc,
    str::FromStr,
};

use super::{Fqsn, NameKind, NameSpec, ScopeSegment, SymbolError, SymbolSpec, SymbolTable, TypeKind};
use crate::{
    ast::{
        Declaration, Expression, ExpressionKind, ItemId, Statement, StatementKind, TypeBody,
        TypeSingletonName, Variant, VariantKind, AST,
    },
    builtin::Builtin,
    parsing::Location,
    type_inference::{self, PendingType, TypeBuilder, TypeContext, VariantBuilder},
};

pub(super) struct SymbolTablePopulator<'a> {
    pub(super) type_context: &'a mut TypeContext,
    pub(super) table: &'a mut SymbolTable,
}

impl<'a> SymbolTablePopulator<'a> {
    /* note: this adds names for *forward reference* but doesn't actually create any types. solve that problem
     * later */

    fn add_symbol(&mut self, id: &ItemId, fqsn: Fqsn, spec: SymbolSpec) {
        self.table.add_symbol(id, fqsn, spec)
    }

    /// This function traverses the AST and adds symbol table entries for
    /// constants, functions, types, and modules defined within. This simultaneously
    /// checks for dupicate definitions (and returns errors if discovered), and sets
    /// up name tables that will be used by further parts of the compiler
    pub fn populate_definition_tables(&mut self, ast: &AST) -> Vec<SymbolError> {
        let mut scope_stack = vec![];
        self.add_from_scope(ast.statements.as_ref(), &mut scope_stack, false)
    }

    fn add_from_scope(
        &mut self,
        statements: &[Statement<StatementKind>],
        scope_stack: &mut Vec<ScopeSegment>,
        function_scope: bool,
    ) -> Vec<SymbolError> {
        let mut errors = vec![];

        for statement in statements {
            let Statement { id, kind, location } = statement;
            let location = *location;
            if let Err(err) = self.add_single_statement(id, kind, location, scope_stack, function_scope) {
                errors.push(err);
            } else {
                let decl = match kind {
                    StatementKind::Declaration(decl) => decl,
                    _ => continue,
                };
                // If there's an error with a name, don't recurse into subscopes of that name
                let recursive_errs = match decl {
                    Declaration::FuncDecl(signature, body) => {
                        let new_scope = ScopeSegment::Name(signature.name.clone());
                        scope_stack.push(new_scope);
                        let output = self.add_from_scope(body.as_ref(), scope_stack, true);
                        scope_stack.pop();
                        output
                    }
                    Declaration::Module { name, items } => {
                        let new_scope = ScopeSegment::Name(name.clone());
                        scope_stack.push(new_scope);
                        let output = self.add_from_scope(items.as_ref(), scope_stack, false);
                        scope_stack.pop();
                        output
                    }
                    Declaration::TypeDecl { name, body, mutable } => {
                        let type_fqsn = Fqsn::from_scope_stack(scope_stack, name.name.clone());
                        self.add_type_members(name, body, mutable, location, type_fqsn)
                    }

                    Declaration::Impl { type_name, interface_name, block } => {
                        let mut errors = vec![];
                        let new_scope = ScopeSegment::Name(Rc::new(format!("<impl-block>{}", type_name)));
                        scope_stack.push(new_scope);

                        for decl_stmt in block.iter() {
                            let Statement { id, kind, location } = decl_stmt;
                            let location = *location;
                            match kind {
                                decl @ Declaration::FuncDecl(signature, body) => {
                                    let output =
                                        self.add_single_declaration(id, decl, location, scope_stack, true);
                                    if let Err(e) = output {
                                        errors.push(e);
                                    };
                                    let new_scope = ScopeSegment::Name(signature.name.clone());
                                    scope_stack.push(new_scope);
                                    let output = self.add_from_scope(body.as_ref(), scope_stack, true);
                                    scope_stack.pop();
                                    errors.extend(output.into_iter());
                                }
                                _other => errors.push(SymbolError::BadImplBlockEntry),
                            };
                        }
                        scope_stack.pop();
                        errors
                    }
                    _ => vec![],
                };
                errors.extend(recursive_errs.into_iter());
            }
        }

        errors
    }

    fn add_single_statement(
        &mut self,
        id: &ItemId,
        kind: &StatementKind,
        location: Location,
        scope_stack: &[ScopeSegment],
        function_scope: bool,
    ) -> Result<(), SymbolError> {
        match kind {
            StatementKind::Declaration(decl) =>
                self.add_single_declaration(id, decl, location, scope_stack, function_scope),
            _ => return Ok(()),
        }
    }

    fn add_single_declaration(
        &mut self,
        id: &ItemId,
        decl: &Declaration,
        location: Location,
        scope_stack: &[ScopeSegment],
        function_scope: bool,
    ) -> Result<(), SymbolError> {
        match decl {
            Declaration::FuncSig(signature) => {
                let fq_function = Fqsn::from_scope_stack(scope_stack, signature.name.clone());
                self.table
                    .fq_names
                    .register(fq_function.clone(), NameSpec { location, kind: NameKind::Function })?;
                self.table
                    .types
                    .register(fq_function.clone(), NameSpec { location, kind: TypeKind::Function })?;

                self.add_symbol(id, fq_function, SymbolSpec::Func { method: None });
            }
            Declaration::FuncDecl(signature, ..) => {
                let fn_name = &signature.name;
                let fq_function = Fqsn::from_scope_stack(scope_stack, fn_name.clone());
                self.table
                    .fq_names
                    .register(fq_function.clone(), NameSpec { location, kind: NameKind::Function })?;
                self.table
                    .types
                    .register(fq_function.clone(), NameSpec { location, kind: TypeKind::Function })?;

                self.add_symbol(id, fq_function, SymbolSpec::Func { method: None });
            }
            Declaration::TypeDecl { name, .. } => {
                let fq_type = Fqsn::from_scope_stack(scope_stack, name.name.clone());
                self.table.types.register(fq_type, NameSpec { location, kind: TypeKind::Constructor })?;
            }
            //TODO handle type aliases
            Declaration::TypeAlias { .. } => (),
            Declaration::Binding { name, .. } => {
                let fq_binding = Fqsn::from_scope_stack(scope_stack, name.clone());
                self.table
                    .fq_names
                    .register(fq_binding.clone(), NameSpec { location, kind: NameKind::Binding })?;
                if !function_scope {
                    self.add_symbol(id, fq_binding, SymbolSpec::GlobalBinding);
                }
            }
            //TODO implement interfaces
            Declaration::Interface { .. } => (),
            Declaration::Impl { .. } => (),
            Declaration::Module { name, .. } => {
                let fq_module = Fqsn::from_scope_stack(scope_stack, name.clone());
                self.table.fq_names.register(fq_module, NameSpec { location, kind: NameKind::Module })?;
            }
            Declaration::Annotation { name, arguments, inner } => {
                let inner = inner.as_ref();
                self.add_single_statement(
                    &inner.id,
                    &inner.kind,
                    inner.location,
                    scope_stack,
                    function_scope,
                )?;
                self.process_annotation(name.as_ref(), arguments.as_slice(), scope_stack, inner)?;
            }
        }
        Ok(())
    }

    fn process_annotation(
        &mut self,
        name: &str,
        arguments: &[Expression],
        scope_stack: &[ScopeSegment],
        inner: &Statement<StatementKind>,
    ) -> Result<(), SymbolError> {
        if name == "register_builtin" {
            if let Statement {
                id: _,
                location: _,
                kind: StatementKind::Declaration(Declaration::FuncDecl(sig, _)),
            } = inner
            {
                let fqsn = Fqsn::from_scope_stack(scope_stack, sig.name.clone());
                let builtin_name = match arguments {
                    [Expression { kind: ExpressionKind::Value(qname), .. }]
                        if qname.components.len() == 1 =>
                        qname.components[0].clone(),
                    _ =>
                        return Err(SymbolError::BadAnnotation {
                            name: name.to_string(),
                            msg: "Bad argument for register_builtin".to_string(),
                        }),
                };

                let builtin =
                    Builtin::from_str(builtin_name.as_str()).map_err(|_| SymbolError::BadAnnotation {
                        name: name.to_string(),
                        msg: format!("Invalid builtin: {}", builtin_name),
                    })?;

                self.table.populate_single_builtin(fqsn, builtin);
                Ok(())
            } else {
                Err(SymbolError::BadAnnotation {
                    name: name.to_string(),
                    msg: "register_builtin not annotating a function".to_string(),
                })
            }
        } else {
            Err(SymbolError::UnknownAnnotation { name: name.to_string() })
        }
    }

    fn add_type_members(
        &mut self,
        type_name: &TypeSingletonName,
        type_body: &TypeBody,
        _mutable: &bool,
        location: Location,
        type_fqsn: Fqsn,
    ) -> Vec<SymbolError> {
        let (variants, immediate_variant) = match type_body {
            TypeBody::Variants(variants) => (variants.clone(), false),
            TypeBody::ImmediateRecord { id, fields } => (
                vec![Variant {
                    id: *id,
                    name: type_name.name.clone(),
                    kind: VariantKind::Record(fields.clone()),
                }],
                true,
            ),
        };

        // Check for duplicates before registering any types with the TypeContext
        let mut seen_variants = HashSet::new();
        let mut errors = vec![];

        for variant in variants.iter() {
            if seen_variants.contains(&variant.name) {
                errors.push(SymbolError::DuplicateVariant {
                    type_fqsn: type_fqsn.clone(),
                    name: variant.name.as_ref().to_string(),
                })
            }
            seen_variants.insert(variant.name.clone());

            if let VariantKind::Record(ref members) = variant.kind {
                let mut seen_members = HashMap::new();
                for (member_name, _) in members.iter() {
                    match seen_members.entry(member_name.as_ref()) {
                        Entry::Occupied(o) => {
                            let location = *o.get();
                            errors.push(SymbolError::DuplicateRecord {
                                type_fqsn: type_fqsn.clone(),
                                location,
                                record: variant.name.as_ref().to_string(),
                                member: member_name.as_ref().to_string(),
                            });
                        }
                        //TODO eventually this should track meaningful locations
                        Entry::Vacant(v) => {
                            v.insert(location);
                        }
                    }
                }
            }
        }

        if !errors.is_empty() {
            return errors;
        }

        let mut type_builder = TypeBuilder::new(type_name.name.as_ref());

        let mut variant_name_map = HashMap::new();
        for variant in variants.iter() {
            let Variant { name, kind, id } = variant;

            variant_name_map.insert(name.clone(), id);

            let mut variant_builder = VariantBuilder::new(name.as_ref());
            match kind {
                VariantKind::UnitStruct => (),
                VariantKind::TupleStruct(items) =>
                    for type_identifier in items {
                        let pending: PendingType = type_identifier.into();
                        variant_builder.add_member(pending);
                    },
                VariantKind::Record(members) =>
                    for (field_name, type_identifier) in members.iter() {
                        let pending: PendingType = type_identifier.into();
                        variant_builder.add_record_member(field_name.as_ref(), pending);
                    },
            }
            type_builder.add_variant(variant_builder);
        }

        let type_id = self.type_context.register_type(type_builder);
        let type_definition = self.type_context.lookup_type(&type_id).unwrap();

        // This index is guaranteed to be the correct tag
        for (index, variant) in type_definition.variants.iter().enumerate() {
            let id = variant_name_map.get(&variant.name).unwrap();
            let tag = index as u32;
            let spec = match &variant.members {
                type_inference::VariantMembers::Unit => SymbolSpec::DataConstructor { tag, type_id },
                type_inference::VariantMembers::Tuple(..) => SymbolSpec::DataConstructor { tag, type_id },
                type_inference::VariantMembers::Record(..) => SymbolSpec::RecordConstructor { tag, type_id },
            };
            self.table.add_symbol(id, type_fqsn.extend(&variant.name), spec);
        }

        if immediate_variant {
            let variant = &type_definition.variants[0];
            let id = variant_name_map.get(&variant.name).unwrap();
            let spec = SymbolSpec::RecordConstructor { tag: 0, type_id };
            self.table.add_symbol(id, type_fqsn, spec);
        }

        vec![]
    }
}