use crate::ast::ir_builder::IrBuilder;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::ast::type_use::TypeUse;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_parameter::IrParameter;
use crate::source_range::SourceRange;
use crate::symbol::Symbol;
use crate::symbol::function_symbol::FunctionSymbol;
use crate::symbol_table::{SymbolInsertError, SymbolTable};
use crate::type_info::TypeInfo;
use std::cell::RefCell;
use std::ops::Neg;
use std::rc::Rc;

pub struct Function {
    declared_name: String,
    declared_name_source_range: SourceRange,
    is_public: bool,
    parameters: Vec<Parameter>,
    return_type: Option<TypeUse>,
    statements: Vec<Statement>,
    function_symbol: Option<Rc<RefCell<FunctionSymbol>>>,
}

impl Function {
    pub fn new(
        declared_name: &str,
        declared_name_source_range: SourceRange,
        is_public: bool,
        parameters: Vec<Parameter>,
        return_type: Option<TypeUse>,
        statements: Vec<Statement>,
    ) -> Self {
        Self {
            declared_name: declared_name.to_string(),
            declared_name_source_range,
            is_public,
            parameters,
            return_type,
            statements,
            function_symbol: None,
        }
    }

    pub fn declared_name(&self) -> &str {
        &self.declared_name
    }

    pub fn statements(&self) -> Vec<&Statement> {
        self.statements.iter().collect()
    }

    pub fn gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];

        if !diagnostics.is_empty() {
            return Err(diagnostics);
        }

        // insert function symbol
        let insert_result = symbol_table.insert_function_symbol(FunctionSymbol::new(
            self.declared_name(),
            self.declared_name_source_range.clone(),
            false,
        ));

        // get function symbol if successful
        let function_symbol = match insert_result {
            Ok(function_symbol) => function_symbol,
            Err(symbol_insert_error) => {
                return match symbol_insert_error {
                    SymbolInsertError::AlreadyDeclared(already_declared) => {
                        diagnostics.push(Diagnostic::new(
                            &format!(
                                "Function {} already declared in current scope",
                                already_declared.symbol().borrow().declared_name()
                            ),
                            self.declared_name_source_range.start(),
                            self.declared_name_source_range.end(),
                        ));
                        Err(diagnostics)
                    }
                };
            }
        };

        // save function symbol for later
        self.function_symbol = Some(function_symbol.clone());

        // handle parameters
        symbol_table.push_function_scope(&format!("function_scope({})", self.declared_name));
        let mut parameter_symbols = vec![];
        for parameter in &mut self.parameters {
            match parameter.gather_declared_names(symbol_table) {
                Ok(parameter_symbol) => {
                    parameter_symbols.push(parameter_symbol);
                }
                Err(mut parameter_diagnostics) => {
                    diagnostics.append(&mut parameter_diagnostics);
                }
            }
        }
        function_symbol
            .borrow_mut()
            .set_parameters(parameter_symbols);

        // return type
        if let Some(type_use) = &mut self.return_type {
            match type_use.gather_declared_names(symbol_table) {
                Ok(_) => {}
                Err(mut type_use_diagnostics) => {
                    diagnostics.append(&mut type_use_diagnostics);
                }
            }
        }

        symbol_table.push_block_scope(&format!("main_block_scope({})", self.declared_name));
        for statement in &mut self.statements {
            match statement.gather_declared_names(symbol_table) {
                Ok(_) => {}
                Err(mut statement_diagnostics) => {
                    diagnostics.append(&mut statement_diagnostics);
                }
            }
        }
        symbol_table.pop_scope(); // main block scope
        symbol_table.pop_scope(); // function scope

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];

        // parameters
        diagnostics.append(
            &mut self
                .parameters
                .iter_mut()
                .map(|parameter| parameter.check_name_usages(symbol_table))
                .filter_map(|result| result.err())
                .flatten()
                .collect(),
        );

        // return type
        if let Some(type_use) = &mut self.return_type {
            match type_use.check_name_usages(symbol_table) {
                Ok(_) => {
                    // set return type info on function symbol
                    self.function_symbol
                        .as_mut()
                        .unwrap()
                        .borrow_mut()
                        .set_return_type_info(type_use.to_type_info());
                }
                Err(mut type_use_diagnostics) => {
                    diagnostics.append(&mut type_use_diagnostics);
                }
            }
        } else {
            // we don't have a given return type, so it's void
            self.function_symbol
                .as_mut()
                .unwrap()
                .borrow_mut()
                .set_return_type_info(TypeInfo::Void);
        }

        // statements
        for statement in &mut self.statements {
            match statement.check_name_usages(symbol_table) {
                Ok(_) => {}
                Err(mut statement_diagnostics) => {
                    diagnostics.append(&mut statement_diagnostics);
                }
            }
        }

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn type_check(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];

        // parameters
        diagnostics.append(
            &mut self
                .parameters
                .iter_mut()
                .map(|parameter| parameter.type_check(symbol_table))
                .filter_map(Result::err)
                .flatten()
                .collect(),
        );

        let function_symbol = self.function_symbol.as_ref().unwrap();
        let return_type_info = function_symbol.borrow().return_type_info().clone();
        let statements_len = self.statements.len();

        // statements
        diagnostics.append(
            &mut self
                .statements
                .iter_mut()
                .enumerate()
                .map(|(i, statement)| {
                    let is_last = i == statements_len - 1;
                    if is_last {
                        statement.type_check(symbol_table, Some(&return_type_info))
                    } else {
                        statement.type_check(symbol_table, None)
                    }
                })
                .filter_map(Result::err)
                .flatten()
                .collect(),
        );

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn to_ir(&self, symbol_table: &SymbolTable) -> IrFunction {
        let mut builder = IrBuilder::new();

        // parameters
        for (i, parameter) in self.parameters.iter().enumerate() {
            let parameter_symbol = parameter.parameter_symbol();
            let stack_offset = (self.parameters.len() as isize).neg() + (i as isize);
            let ir_parameter = IrParameter::new(
                parameter_symbol.borrow().declared_name(),
                parameter_symbol.borrow().type_info().clone(),
                stack_offset,
            );
            let as_rc = Rc::new(ir_parameter);
            builder.parameters_mut().push(as_rc.clone());
            parameter_symbol.borrow_mut().set_ir_parameter(as_rc);
        }

        let entry_block_id = builder.new_block();

        let function_symbol = self.function_symbol.as_ref().unwrap().borrow();
        let return_type_info = function_symbol.return_type_info();

        let should_return_value = !matches!(return_type_info, TypeInfo::Void);

        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.to_ir(&mut builder, symbol_table, should_return_value && is_last);
        }
        builder.finish_block();

        let entry_block = builder.get_block(entry_block_id).clone();
        IrFunction::new(
            self.function_symbol
                .as_ref()
                .unwrap()
                .borrow()
                .declared_name_owned(), // ok for now... but we need to start using the fqn
            builder.parameters(),
            return_type_info,
            entry_block,
        )
    }
}