use crate::ast::assign_statement::AssignStatement;
use crate::ast::constructor::Constructor;
use crate::ast::expression::Expression;
use crate::ast::field::Field;
use crate::ast::fqn_context::FqnContext;
use crate::ast::fqn_util::fqn_parts_to_string;
use crate::ast::function::Function;
use crate::ast::generic_parameter::GenericParameter;
use crate::ast::helpers::{collect_diagnostics_mut, collect_diagnostics_single, resolve_ctor_name};
use crate::ast::statement::Statement;
use crate::diagnostic::Diagnostic;
use crate::error_codes::{FIELD_MULTIPLE_INIT, FIELD_UNINIT};
use crate::ir::ir_class::{IrClass, IrField};
use crate::ir::ir_function::IrFunction;
use crate::source_range::SourceRange;
use crate::symbol::Symbol;
use crate::symbol::class_symbol::ClassSymbol;
use crate::symbol::constructor_symbol::ConstructorSymbol;
use crate::symbol::variable_symbol::VariableSymbol;
use crate::symbol_table::SymbolTable;
use crate::type_info::TypeInfo;
use crate::types_table::TypesTable;
use crate::{diagnostics_result, handle_diagnostics, ok_or_err_diagnostics};
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

pub struct Class {
    declared_name: Rc<str>,
    declared_name_source_range: SourceRange,
    generic_parameters: Vec<GenericParameter>,
    constructor: Option<Constructor>,
    fields: Vec<Field>,
    functions: Vec<Function>,
    scope_id: Option<usize>,
    self_class_scope_id: Option<usize>,
    self_class_body_scope_id: Option<usize>,
}

impl Class {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        generic_parameters: Vec<GenericParameter>,
        constructor: Option<Constructor>,
        fields: Vec<Field>,
        functions: Vec<Function>,
    ) -> Self {
        Self {
            declared_name: declared_name.into(),
            declared_name_source_range,
            generic_parameters,
            constructor,
            fields,
            functions,
            scope_id: None,
            self_class_scope_id: None,
            self_class_body_scope_id: None,
        }
    }

    pub fn init_scopes(&mut self, symbol_table: &mut SymbolTable, container_scope: usize) {
        self.scope_id = Some(container_scope);

        let class_scope_id =
            symbol_table.push_class_scope(&format!("class_scope({})", self.declared_name));
        self.self_class_scope_id = Some(class_scope_id);

        for generic_parameter in &mut self.generic_parameters {
            generic_parameter.init_scopes(symbol_table, class_scope_id);
        }

        let class_body_scope_id = symbol_table
            .push_class_body_scope(&format!("class_body_scope({})", self.declared_name));
        self.self_class_body_scope_id = Some(class_body_scope_id);

        for field in &mut self.fields {
            field.init_scopes(symbol_table, class_body_scope_id);
        }

        if let Some(constructor) = &mut self.constructor {
            constructor.init_scopes(symbol_table, class_body_scope_id);
        }
        for function in &mut self.functions {
            function.init_scopes(symbol_table, class_body_scope_id);
        }

        symbol_table.pop_scope();
        symbol_table.pop_scope();
    }

    pub fn make_symbols(&self, fqn_context: &mut FqnContext) -> Vec<Symbol> {
        let mut all_symbols: Vec<Symbol> = Vec::new();

        let mut generic_parameter_symbols = Vec::new();
        for generic_parameter in &self.generic_parameters {
            let symbol = Rc::new(generic_parameter.make_symbol());
            all_symbols.push(Symbol::GenericParameter(symbol.clone()));
            generic_parameter_symbols.push(symbol);
        }

        let mut field_symbols = Vec::new();
        for (field_index, field) in self.fields.iter().enumerate() {
            let symbol = Rc::new(field.make_symbol(field_index));
            all_symbols.push(Symbol::Field(symbol.clone()));
            field_symbols.push(symbol);
        }

        fqn_context.push(self.declared_name.clone()); // namespace this class' members

        let constructor_symbol = if let Some(constructor) = &self.constructor {
            let (constructor_symbol, mut symbols) = constructor.make_symbols(fqn_context);
            all_symbols.append(&mut symbols);
            constructor_symbol
        } else {
            Rc::new(ConstructorSymbol::new(
                &self.declared_name_source_range,
                resolve_ctor_name(fqn_context),
                false,
                true,
                self.self_class_body_scope_id.unwrap(),
                vec![],
            ))
        };
        all_symbols.push(Symbol::Constructor(constructor_symbol.clone()));

        let mut function_symbols = Vec::new();
        for function in &self.functions {
            let (function_symbol, mut symbols) = function.make_symbols(fqn_context, true);
            all_symbols.append(&mut symbols);
            function_symbols.push(function_symbol);
        }

        fqn_context.pop(); // un-namespace

        let class_symbol = Rc::new(ClassSymbol::new(
            &self.declared_name,
            Some(self.declared_name_source_range.clone()),
            fqn_context.resolve(&self.declared_name),
            false,
            self.scope_id.unwrap(),
            generic_parameter_symbols,
            Some(constructor_symbol),
            field_symbols,
            function_symbols,
        ));
        all_symbols.push(Symbol::Class(class_symbol.clone()));

        all_symbols
    }

    pub fn check_names(&self, symbol_table: &mut SymbolTable) -> Vec<Diagnostic> {
        let mut diagnostics: Vec<Diagnostic> = Vec::new();

        for generic_parameter in &self.generic_parameters {
            diagnostics.append(&mut generic_parameter.check_names(symbol_table));
        }

        for field in &self.fields {
            diagnostics.append(&mut field.check_names(symbol_table));
        }

        if let Some(constructor) = &self.constructor {
            diagnostics.append(&mut constructor.check_names(symbol_table));
        }

        for function in &self.functions {
            diagnostics.append(&mut function.check_names(symbol_table));
        }

        diagnostics
    }

    fn get_class_symbol_owned(&self, symbol_table: &SymbolTable) -> Rc<ClassSymbol> {
        symbol_table
            .get_class_symbol(self.scope_id.unwrap(), &self.declared_name)
            .cloned()
            .unwrap()
    }

    pub fn check_field_initializer_names(&self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        let class_symbol = self.get_class_symbol_owned(symbol_table);
        self.fields
            .iter()
            .flat_map(|field| field.check_field_initializer_names(symbol_table, &class_symbol))
            .collect()
    }

    pub fn analyze_local_names(&self, symbol_table: &mut SymbolTable) -> Vec<Diagnostic> {
        let class_symbol = self.get_class_symbol_owned(symbol_table);
        let mut diagnostics: Vec<Diagnostic> = Vec::new();
        if let Some(constructor) = &self.constructor {
            diagnostics.append(&mut constructor.analyze_local_names(symbol_table, &class_symbol));
        }
        for function in &self.functions {
            diagnostics
                .append(&mut function.analyze_method_local_names(symbol_table, &class_symbol));
        }
        diagnostics
    }

    pub fn gather_types(
        &self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        // instance type
        let class_symbol = self.get_class_symbol_owned(symbol_table);
        types_table.class_instance_types_mut().insert(
            class_symbol.clone(),
            TypeInfo::ClassInstance(class_symbol.clone()),
        );

        // constructor return type
        // this works for both declared and default constructors
        let constructor_symbol = symbol_table
            .get_constructor_symbol_owned(self.self_class_body_scope_id.unwrap())
            .unwrap();
        types_table
            .constructor_return_types_mut()
            .insert(constructor_symbol, TypeInfo::ClassInstance(class_symbol));

        let mut diagnostics = Vec::new();

        // generic params
        for generic_parameter in &self.generic_parameters {
            handle_diagnostics!(
                generic_parameter.gather_types(symbol_table, types_table),
                diagnostics
            );
        }

        // field types
        for field in &self.fields {
            handle_diagnostics!(field.gather_types(symbol_table, types_table), diagnostics);
        }

        // function return types
        for function in &self.functions {
            function.gather_types(symbol_table, types_table);
        }

        diagnostics_result!(diagnostics)
    }

    fn type_check_generics(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        collect_diagnostics_mut(&mut self.generic_parameters, |gp| {
            gp.type_check(symbol_table, types_table)
        })
    }

    fn type_check_fields(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        collect_diagnostics_mut(&mut self.fields, |f| {
            f.type_check(symbol_table, types_table)
        })
    }

    fn type_check_constructor(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        if let Some(constructor) = &mut self.constructor {
            constructor.type_check(symbol_table, types_table)?;
        }
        Ok(())
    }

    fn type_check_functions(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        collect_diagnostics_mut(&mut self.functions, |f| {
            f.type_check(symbol_table, types_table)
        })
    }

    /// Returns all field names with declared initializers.
    fn field_names_with_initializers(&self) -> HashSet<&str> {
        let mut set: HashSet<&str> = HashSet::new();
        for field in &self.fields {
            if field.initializer().is_some() {
                set.insert(field.declared_name());
            }
        }
        set
    }

    /// If the destination of the given [AssignStatement] matches a field, returns an
    /// `Ok(Some(field_name))` only if the field is not already in the `fields_already_init` set,
    /// AND, if the field is immutable, the field is not initialized more than once in the
    /// constructor. Otherwise, returns an `Err(Diagnostic)`.
    fn check_ctor_assign_statement<'a>(
        &self,
        assign_statement: &'a AssignStatement,
        fields_already_init: &HashSet<&&str>,
        class_symbol: &ClassSymbol,
    ) -> Result<Option<&'a str>, Diagnostic> {
        match assign_statement.destination() {
            Expression::Identifier(identifier) => {
                // find matching field symbol, if there is one
                if let Some(field_symbol) = class_symbol.fields().get(identifier.name()) {
                    // check that we don't init more than once IF field is immutable
                    if fields_already_init.contains(&identifier.name()) && !field_symbol.is_mut() {
                        let diagnostic = Diagnostic::new(
                            &format!("Immutable field {} cannot be initialized more than once in constructor.", identifier.name()),
                            identifier.source_range().start(),
                            identifier.source_range().end(),
                        ).with_reporter(file!(), line!())
                            .with_error_code(FIELD_MULTIPLE_INIT);
                        Err(diagnostic)
                    } else {
                        Ok(Some(identifier.name()))
                    }
                } else {
                    Ok(None)
                }
            }
            _ => panic!("Found a non-L Value destination"),
        }
    }

    /// Returns an `Ok(HashSet<&str>)` containing the names of all fields initialized in the
    /// constructor, provided that the following are true:
    ///
    /// - The field is not initialized more than once in the constructor
    /// - The field is not also initialized with a declared initializer.
    ///
    /// If the above are not met, returns `Err(diagnostics)`.
    fn get_fields_init_in_ctor<'a>(
        &self,
        fields_with_declared_initializers: &HashSet<&str>,
        symbol_table: &SymbolTable,
    ) -> Result<HashSet<&str>, Vec<Diagnostic>> {
        let mut constructor_inits: HashSet<&str> = HashSet::new();
        let mut diagnostics: Vec<Diagnostic> = vec![];
        if let Some(constructor) = &self.constructor {
            let class_symbol = symbol_table
                .get_class_symbol(self.scope_id.unwrap(), &self.declared_name)
                .unwrap();
            for statement in constructor.statements() {
                match statement {
                    Statement::Assign(assign_statement) => {
                        let fields_init_so_far = constructor_inits
                            .union(fields_with_declared_initializers)
                            .collect::<HashSet<_>>();
                        match self.check_ctor_assign_statement(
                            assign_statement,
                            &fields_init_so_far,
                            &class_symbol,
                        ) {
                            Ok(maybe_init_field) => match maybe_init_field {
                                None => {}
                                Some(init_field) => {
                                    constructor_inits.insert(init_field);
                                }
                            },
                            Err(diagnostic) => {
                                diagnostics.push(diagnostic);
                            }
                        }
                    }
                    _ => {}
                }
            }
        }
        ok_or_err_diagnostics!(constructor_inits, diagnostics)
    }

    /// Checks that all declared fields in this `Class` are present in the `all_inits` set. If so,
    /// returns `Ok`, else `Err`.
    fn check_all_fields_in_init_set(
        &self,
        all_inits: &HashSet<&&str>,
    ) -> Result<(), Vec<Diagnostic>> {
        collect_diagnostics_single(&self.fields, |field| {
            if all_inits.contains(&field.declared_name()) {
                Ok(())
            } else {
                Err(Diagnostic::new(
                    &format!("Field {} is not initialized.", field.declared_name()),
                    field.declared_name_source_range().start(),
                    field.declared_name_source_range().end(),
                )
                .with_primary_label_message("Must be initialized in declaration or constructor.")
                .with_reporter(file!(), line!())
                .with_error_code(FIELD_UNINIT))
            }
        })
    }

    /// Checks that all fields are initialized, either at their declaration or in the constructor.
    /// Immutable fields may be only initialized once, either at their declaration or once in the
    /// constructor. Mutable fields may be initialized either at their declaration, or at least once
    /// in the constructor.
    fn check_field_initialization(
        &self,
        symbol_table: &SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        // We need to determine if fields are initialized or not (the latter is an error).
        // First phase: check all fields, then check constructor, leaving pending those things that
        // are fields <- initialized by constructor. Then circle back to fields and check all are
        // initialized
        let field_names_with_initializers = self.field_names_with_initializers();
        let field_names_init_in_constructor =
            self.get_fields_init_in_ctor(&field_names_with_initializers, symbol_table)?;
        let combined = field_names_with_initializers
            .union(&field_names_init_in_constructor)
            .collect::<HashSet<_>>();

        // check that all fields are present in the hash set
        self.check_all_fields_in_init_set(&combined)?;

        Ok(())
    }

    pub fn type_check(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        self.type_check_generics(symbol_table, types_table)?;
        self.type_check_fields(symbol_table, types_table)?;
        self.type_check_constructor(symbol_table, types_table)?;
        self.type_check_functions(symbol_table, types_table)?;
        self.check_field_initialization(symbol_table)?;
        Ok(())
    }

    pub fn to_ir(
        &self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> (IrClass, Vec<IrFunction>) {
        let self_class_symbol = symbol_table
            .get_class_symbol(self.scope_id.unwrap(), &self.declared_name)
            .unwrap();

        let mut ir_functions: Vec<IrFunction> = vec![];
        if let Some(constructor) = &self.constructor {
            ir_functions.push(constructor.to_ir(
                self_class_symbol,
                &self.fields,
                symbol_table,
                types_table,
            ))
        }

        for function in &self.functions {
            ir_functions.push(function.to_ir(symbol_table, types_table, Some(self_class_symbol)));
        }

        let ir_class = IrClass::new(
            self_class_symbol.declared_name_owned(),
            fqn_parts_to_string(self_class_symbol.fqn_parts()).into(),
            self.fields
                .iter()
                .map(|field| {
                    let field_symbol = symbol_table
                        .get_field_symbol_owned(field.scope_id(), field.declared_name())
                        .unwrap();
                    let field_type = types_table.field_types().get(&field_symbol).unwrap();
                    IrField::new(
                        field.declared_name().into(),
                        field_symbol.field_index(),
                        field_type.clone(),
                    )
                })
                .collect(),
        );

        (ir_class, ir_functions)
    }
}