use crate::variable_manager::{self, LocationRequest, VariableLocation, VariableScope};
use parser::{
    Parser as ASTParser,
    tree_node::{
        BlockExpression, DeviceDeclarationExpression, Expression, FunctionExpression,
        InvocationExpression, Literal,
    },
};
use quick_error::quick_error;
use std::{
    collections::HashMap,
    io::{BufWriter, Write},
};

macro_rules! debug {
    ($self: expr, $debug_value: expr) => {
        if $self.config.debug {
            format!($debug_value)
        } else {
            "".into()
        }
    };
}

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        ParseError(error: parser::Error) {
            from()
        }
        IoError(error: std::io::Error) {
            from()
        }
        ScopeError(error: variable_manager::Error) {
            from()
        }
        DuplicateIdentifier(func_name: String) {
            display("`{func_name}` has already been defined")
        }
        UnknownIdentifier(ident: String) {
            display("`{ident}` is not found in the current scope.")
        }
        InvalidDevice(device: String) {
            display("`{device}` is not valid")
        }
        AgrumentMismatch(func_name: String) {
            display("Incorrect number of arguments passed into `{func_name}`")
        }
        Unknown(reason: String) {
            display("{reason}")
        }
    }
}

#[derive(Default)]
#[repr(C)]
pub struct CompilerConfig {
    pub debug: bool,
}

pub struct Compiler<'a, W: std::io::Write> {
    parser: ASTParser,
    function_locations: HashMap<String, usize>,
    function_metadata: HashMap<String, Vec<String>>,
    devices: HashMap<String, String>,
    output: &'a mut BufWriter<W>,
    current_line: usize,
    declared_main: bool,
    config: CompilerConfig,
}

impl<'a, W: std::io::Write> Compiler<'a, W> {
    pub fn new(
        parser: ASTParser,
        writer: &'a mut BufWriter<W>,
        config: Option<CompilerConfig>,
    ) -> Self {
        Self {
            parser,
            function_locations: HashMap::new(),
            function_metadata: HashMap::new(),
            devices: HashMap::new(),
            output: writer,
            current_line: 1,
            declared_main: false,
            config: config.unwrap_or_default(),
        }
    }

    pub fn compile(mut self) -> Result<(), Error> {
        let expr = self.parser.parse_all()?;

        let Some(expr) = expr else { return Ok(()) };

        self.write_output("j main")?;
        self.expression(expr, &mut VariableScope::default())
    }

    fn write_output(&mut self, output: impl Into<String>) -> Result<(), Error> {
        self.output.write_all(output.into().as_bytes())?;
        self.output.write_all(b"\n")?;
        self.current_line += 1;
        Ok(())
    }

    fn expression<'v>(
        &mut self,
        expr: Expression,
        scope: &mut VariableScope<'v>,
    ) -> Result<(), Error> {
        match expr {
            Expression::Function(expr_func) => self.expression_function(expr_func, scope)?,
            Expression::Block(expr_block) => self.expression_block(expr_block, scope)?,
            Expression::DeviceDeclaration(expr_dev) => self.expression_device(expr_dev)?,
            Expression::Declaration(var_name, expr) => {
                self.expression_declaration(var_name, *expr, scope)?
            }
            Expression::Invocation(expr_invoke) => {
                self.expression_function_invocation(expr_invoke, scope)?
            }
            _ => todo!(),
        };

        Ok(())
    }

    fn expression_declaration<'v>(
        &mut self,
        var_name: String,
        expr: Expression,
        scope: &mut VariableScope<'v>,
    ) -> Result<(), Error> {
        match expr {
            Expression::Literal(Literal::Number(num)) => {
                let var_location =
                    scope.add_variable(var_name.clone(), LocationRequest::Persist)?;
                let num_str = num.to_string();

                if let VariableLocation::Temporary(reg) | VariableLocation::Persistant(reg) =
                    var_location
                {
                    self.write_output(format!(
                        "move r{reg} {num_str} {}",
                        debug!(self, "#{var_name}")
                    ))?;
                } else {
                    self.write_output(format!("push {num_str} {}", debug!(self, "#{var_name}")))?;
                }
            }
            Expression::Invocation(invoke_expr) => {
                self.expression_function_invocation(invoke_expr, scope)?;

                // Return value _should_ be in VariableScope::RETURN_REGISTER
                match scope.add_variable(var_name.clone(), LocationRequest::Persist)? {
                    VariableLocation::Temporary(reg) | VariableLocation::Persistant(reg) => self
                        .write_output(format!(
                            "move r{reg} r{} {}",
                            VariableScope::RETURN_REGISTER,
                            debug!(self, "#{var_name}")
                        ))?,
                    VariableLocation::Stack(_) => self.write_output(format!(
                        "push r{} {}",
                        VariableScope::RETURN_REGISTER,
                        debug!(self, "#{var_name}")
                    ))?,
                }
            }
            _ => {
                return Err(Error::Unknown(
                    "`{var_name}` declaration of this type is not supported.".into(),
                ));
            }
        }

        Ok(())
    }

    fn expression_function_invocation(
        &mut self,
        invoke_expr: InvocationExpression,
        stack: &mut VariableScope,
    ) -> Result<(), Error> {
        if !self.function_locations.contains_key(&invoke_expr.name) {
            return Err(Error::UnknownIdentifier(invoke_expr.name));
        }

        let Some(args) = self.function_metadata.get(&invoke_expr.name) else {
            return Err(Error::UnknownIdentifier(invoke_expr.name));
        };

        if args.len() != invoke_expr.arguments.len() {
            return Err(Error::AgrumentMismatch(invoke_expr.name));
        }

        // backup all used registers to the stack
        let active_registers = stack.registers().cloned().collect::<Vec<_>>();
        for register in &active_registers {
            stack.add_variable(format!("temp_{register}"), LocationRequest::Stack)?;
            self.write_output(format!("push r{register}"))?;
        }
        for arg in invoke_expr.arguments {
            match arg {
                Expression::Literal(Literal::Number(num)) => {
                    let num_str = num.to_string();
                    self.write_output(format!("push {num_str}"))?;
                }
                Expression::Variable(var_name) => match stack.get_location_of(var_name)? {
                    VariableLocation::Persistant(reg) | VariableLocation::Temporary(reg) => {
                        self.write_output(format!("push r{reg}"))?;
                    }
                    VariableLocation::Stack(stack_offset) => {
                        self.write_output(format!(
                            "sub r{0} sp {stack_offset}",
                            VariableScope::TEMP_STACK_REGISTER
                        ))?;
                        self.write_output(format!(
                            "get r{0} db r{0}",
                            VariableScope::TEMP_STACK_REGISTER
                        ))?;
                        self.write_output(format!(
                            "push r{0}",
                            VariableScope::TEMP_STACK_REGISTER
                        ))?;
                    }
                },
                _ => {
                    return Err(Error::Unknown(format!(
                        "Attempted to call `{}` with an unsupported argument",
                        invoke_expr.name
                    )));
                }
            }
        }

        // jump to the function and store current line in ra
        self.write_output(format!("jal {}", invoke_expr.name))?;

        for register in active_registers {
            let VariableLocation::Stack(stack_offset) =
                stack.get_location_of(format!("temp_{register}"))?
            else {
                return Err(Error::UnknownIdentifier(format!("temp_{register}")));
            };
            self.write_output(format!(
                "sub r{0} sp {stack_offset}",
                VariableScope::TEMP_STACK_REGISTER
            ))?;
            self.write_output(format!(
                "get r{register} db r{0}",
                VariableScope::TEMP_STACK_REGISTER
            ))?;
        }

        if stack.stack_offset() > 0 {
            self.write_output(format!("sub sp sp {}", stack.stack_offset()))?;
        }

        Ok(())
    }

    fn expression_device(&mut self, expr: DeviceDeclarationExpression) -> Result<(), Error> {
        if self.devices.contains_key(&expr.name) {
            return Err(Error::DuplicateIdentifier(expr.name));
        }
        self.devices.insert(expr.name, expr.device);

        Ok(())
    }

    fn expression_block<'v>(
        &mut self,
        mut expr: BlockExpression,
        scope: &mut VariableScope<'v>,
    ) -> Result<(), Error> {
        // First, sort the expressions to ensure functions are hoisted
        expr.0.sort_by(|a, b| {
            if matches!(b, Expression::Function(_)) && matches!(a, Expression::Function(_)) {
                std::cmp::Ordering::Equal
            } else if matches!(a, Expression::Function(_)) {
                std::cmp::Ordering::Less
            } else {
                std::cmp::Ordering::Greater
            }
        });

        for expr in expr.0 {
            if !self.declared_main
                && !matches!(expr, Expression::Function(_))
                && !scope.has_parent()
            {
                self.write_output("main:")?;
                self.declared_main = true;
            }

            self.expression(expr, scope)?;
        }

        Ok(())
    }

    /// Takes the result of the expression and stores it in VariableScope::RETURN_REGISTER
    fn expression_return<'v>(
        &mut self,
        expr: Expression,
        scope: &mut VariableScope<'v>,
    ) -> Result<(), Error> {
        match expr {
            Expression::Variable(var_name) => match scope.get_location_of(var_name)? {
                VariableLocation::Temporary(reg) | VariableLocation::Persistant(reg) => {
                    self.write_output(format!(
                        "move r{} r{reg} {}",
                        VariableScope::RETURN_REGISTER,
                        debug!(self, "#returnValue")
                    ))?;
                }
                VariableLocation::Stack(offset) => {
                    self.write_output(format!(
                        "sub r{} sp {offset}",
                        VariableScope::TEMP_STACK_REGISTER
                    ))?;
                    self.write_output(format!(
                        "get r{} db r{}",
                        VariableScope::RETURN_REGISTER,
                        VariableScope::TEMP_STACK_REGISTER
                    ))?;
                }
            },
            Expression::Literal(Literal::Number(num)) => {
                self.write_output(format!("move r{} {}", VariableScope::RETURN_REGISTER, num))?;
            }
            _ => return Err(Error::Unknown("Unsupported `return` statement.".into())),
        }

        Ok(())
    }

    /// Compile a function declaration.
    /// Calees are responsible for backing up any registers they wish to use.
    fn expression_function<'v>(
        &mut self,
        expr: FunctionExpression,
        scope: &mut VariableScope<'v>,
    ) -> Result<(), Error> {
        let FunctionExpression {
            name,
            arguments,
            body,
        } = expr;

        if self.function_locations.contains_key(&name) {
            return Err(Error::DuplicateIdentifier(name));
        }

        self.function_metadata
            .insert(name.clone(), arguments.clone());

        // Declare the function as a line identifier
        self.write_output(format!("{}:", name))?;

        self.function_locations
            .insert(name.clone(), self.current_line);

        // Create a new block scope for the function body
        let mut block_scope = VariableScope::scoped(scope);

        let mut saved_variables = 0;

        // do a reverse pass to pop variables from the stack and put them into registers
        for var_name in arguments
            .iter()
            .rev()
            .take(VariableScope::PERSIST_REGISTER_COUNT as usize)
        {
            let loc = block_scope.add_variable(var_name, LocationRequest::Persist)?;
            // we don't need to imcrement the stack offset as it's already on the stack from the
            // previous scope

            match loc {
                VariableLocation::Persistant(loc) => {
                    self.write_output(format!("pop r{loc} {}", debug!(self, "#{var_name}")))?;
                }
                VariableLocation::Stack(_) => {
                    return Err(Error::Unknown(
                        "Attempted to save to stack without tracking in scope".into(),
                    ));
                }

                _ => {
                    return Err(Error::Unknown(
                        "Attempted to return a Temporary scoped variable from a Persistant request"
                            .into(),
                    ));
                }
            }
            saved_variables += 1;
        }

        // now do a forward pass in case we have spilled into the stack. We don't need to push
        // anything as they already exist on the stack, but we DO need to let our block_scope be
        // aware that the variables exist on the stack (left to right)
        for var_name in arguments.iter().take(arguments.len() - saved_variables) {
            block_scope.add_variable(var_name, LocationRequest::Stack)?;
        }

        self.write_output("push ra")?;
        block_scope.add_variable(format!("{name}_ra"), LocationRequest::Stack)?;

        for expr in body.0 {
            match expr {
                Expression::Return(ret_expr) => {
                    self.expression_return(*ret_expr, &mut block_scope)?
                }
                _ => self.expression(expr, &mut block_scope)?,
            }
        }

        // Get the saved return address and save it back into `ra`
        let VariableLocation::Stack(ra_stack_offset) =
            block_scope.get_location_of(format!("{name}_ra"))?
        else {
            return Err(Error::Unknown(
                "Stored return address not in stack as expected".into(),
            ));
        };

        self.write_output(format!(
            "sub r{0} sp {ra_stack_offset}",
            VariableScope::TEMP_STACK_REGISTER
        ))?;
        self.write_output(format!(
            "get ra db r{0}",
            VariableScope::TEMP_STACK_REGISTER
        ))?;

        if block_scope.stack_offset() > 0 {
            self.write_output(format!("sub sp sp {}", block_scope.stack_offset()))?;
        }

        self.write_output("j ra")?;
        Ok(())
    }
}