#![feature(option_zip)]

use crate::ast::statement::Statement;
use crate::constants_table::ConstantsTable;
use crate::diagnostic::Diagnostics;
use crate::ir::compile_dvm_function;
use crate::ir::ir_variable::{IrStackFrameVariables, IrVariable, IrVariableInfo};
use crate::lowering::util::to_ir_type_info;
use crate::lowering::{lower_to_ir_compilation_unit, lower_to_ir_synthetic_function};
use crate::parser::{ParseResult, parse_compilation_unit};
use crate::semantic_analysis::symbol::SymbolId;
use crate::semantic_analysis::{analyze_compilation_unit, analyze_statement};
use dvm_lib::vm::function::Function;
use semantic_analysis::analysis_context::AnalysisContext;
use std::collections::HashMap;
use std::rc::Rc;

pub mod ast;
pub mod constants_table;
pub mod diagnostic;
mod diagnostic_factories;
mod error_codes;
pub mod ir;
pub mod lexer;
pub mod lowering;
pub mod parser;
pub mod semantic_analysis;
pub mod source_range;
pub mod token;

pub type Filename = Rc<str>;
pub type FileId = usize;

pub struct CompileCompilationUnitResult {
    pub functions: HashMap<Rc<str>, Function>,
}

pub fn compile_compilation_unit(
    input: &str,
    analysis_context: &mut AnalysisContext,
    register_count: usize,
    constants_table: &mut ConstantsTable,
) -> Result<CompileCompilationUnitResult, Diagnostics> {
    let ParseResult(compilation_unit, mut parse_diagnostics) = parse_compilation_unit(input, None);
    let mut analysis_diagnostics = analyze_compilation_unit(&compilation_unit, analysis_context);

    if !parse_diagnostics.is_empty() || !analysis_diagnostics.is_empty() {
        let mut all_diagnostics = Vec::new();
        all_diagnostics.append(&mut parse_diagnostics);
        all_diagnostics.append(&mut analysis_diagnostics);
        return Err(all_diagnostics);
    }

    let lower_to_ir_result = lower_to_ir_compilation_unit(&compilation_unit, analysis_context);

    let mut dvm_functions = HashMap::new();

    for ir_function in &lower_to_ir_result.functions {
        let dvm_function = compile_dvm_function(ir_function, register_count, constants_table);
        dvm_functions.insert(dvm_function.name_owned(), dvm_function);
    }

    Ok(CompileCompilationUnitResult {
        functions: dvm_functions,
    })
}

pub struct SyntheticFunctionSession {
    fqn: Rc<str>,
    ctx: Box<AnalysisContext>,
    env: Box<SyntheticEnvironment>,
}

impl SyntheticFunctionSession {
    pub fn new(fqn: &str) -> Self {
        Self {
            fqn: fqn.into(),
            ctx: AnalysisContext::new().into(),
            env: SyntheticEnvironment::new().into(),
        }
    }

    pub fn ctx(&self) -> &AnalysisContext {
        &self.ctx
    }

    pub fn ctx_mut(&mut self) -> &mut AnalysisContext {
        &mut self.ctx
    }

    pub fn compile_statement(
        &mut self,
        statement: &Statement,
        register_count: usize,
        constants_table: &mut ConstantsTable,
    ) -> Result<Function, Diagnostics> {
        let diagnostics = analyze_statement(statement, &mut self.ctx);
        if !diagnostics.is_empty() {
            return Err(diagnostics);
        }

        // Allocate stack frame variable for let statements only
        match statement {
            Statement::Let(let_statement) => {
                let ir_variable_info = IrVariableInfo::new(
                    let_statement.declared_name_owned(),
                    to_ir_type_info(self.ctx.get_type_info_for_node(let_statement.node_id())),
                );
                let ir_stack_frame_variable_id = self
                    .env
                    .ir_stack_frame_variables_mut()
                    .push(ir_variable_info);
                self.env.symbols_to_variables_mut().insert(
                    self.ctx.nodes_to_symbols()[&let_statement.node_id()],
                    IrVariable::StackFrame(ir_stack_frame_variable_id),
                );
            }
            _ => {}
        }

        let ir_function = lower_to_ir_synthetic_function(statement, &self);
        Ok(compile_dvm_function(
            &ir_function,
            register_count,
            constants_table,
        ))
    }
}

pub struct SyntheticEnvironment {
    ir_stack_frame_variables: IrStackFrameVariables,
    symbols_to_variables: HashMap<SymbolId, IrVariable>,
}

impl SyntheticEnvironment {
    pub fn new() -> Self {
        Self {
            ir_stack_frame_variables: IrStackFrameVariables::new(),
            symbols_to_variables: HashMap::new(),
        }
    }

    pub fn ir_stack_frame_variables(&self) -> &IrStackFrameVariables {
        &self.ir_stack_frame_variables
    }

    pub fn ir_stack_frame_variables_mut(&mut self) -> &mut IrStackFrameVariables {
        &mut self.ir_stack_frame_variables
    }

    pub fn symbols_to_variables(&self) -> &HashMap<SymbolId, IrVariable> {
        &self.symbols_to_variables
    }

    pub fn symbols_to_variables_mut(&mut self) -> &mut HashMap<SymbolId, IrVariable> {
        &mut self.symbols_to_variables
    }
}