use sexp::Sexp;

use crate::asm::data::Register;
use crate::asm::error::CrsnError;
use crate::asm::instr::{Flatten, Routine};
use crate::asm::parse::{parse_instructions, ParserContext};
use crate::asm::parse::arg_parser::TokenParser;
use crate::asm::parse::parse_data::parse_reg_alias;
use crate::asm::parse::sexp_expect::expect_string_atom;
use crate::asm::patches::SexpIsA;
use crate::builtin::parse::parse_routine_name;

pub fn parse_routine(mut toki: impl Iterator<Item=Sexp> + Clone, pcx: &ParserContext) -> Result<Box<dyn Flatten>, CrsnError> {
    let name = expect_string_atom(toki.next())?;
    let mut name = parse_routine_name(name)?;

    let arg_tokens = TokenParser::new(toki.clone().take_while(|e| e.is_atom()).collect(), pcx);

    // If arity is explicitly given, then either no named argument must be provided,
    // or their count must match the arity. If no arity is given, then arity is determined
    // by the number of named arguments.
    if name.arity == 0 && arg_tokens.len() != 0 {
        name.arity = arg_tokens.len() as u8;
    } else if arg_tokens.len() != 0 && name.arity as usize != arg_tokens.len() {
        return Err(CrsnError::Parse(format!("arity mismatch in routine {}", name.name).into()));
    }

    let toki = toki.skip_while(|e| e.is_atom());

    {
        let mut pstate = pcx.state.borrow_mut();

        let old = std::mem::replace(&mut pstate.reg_aliases, Default::default());
        pstate.reg_alias_stack.push(old);

        for (n, tok) in arg_tokens.into_iter().enumerate() {
            let alias = parse_reg_alias(Some(tok))?;

            if pstate.constants.contains_key(&alias) {
                return Err(CrsnError::Parse(format!("Symbol \"{}\" already used for a constant.", alias).into()));
            }

            pstate.reg_aliases.insert(alias, Register::Arg(n as u8));
        }
    }

    let body = parse_instructions(toki, pcx)?;

    {
        let mut pstate = pcx.state.borrow_mut();
        let old = pstate.reg_alias_stack.pop().unwrap();
        pstate.reg_aliases = old;
    }

    return Ok(Box::new(Routine {
        name,
        body,
    }));
}