use crate::*;
use AssemblerErrorVariant as ErrVar;

use indexmap::IndexMap;


static mut ID: usize = 0;
macro_rules! new_id {
    () => { unsafe {
        let id = ID;
        ID += 1;
        id
    }};
}


impl SemanticProgram {
    pub fn assemble(&self) -> Vec<AssembledToken> {
        let environment = Environment {
            macro_definitions: &self.macro_definitions,
            label_definitions: &self.label_definitions,
            arguments: &IndexMap::new(),
            id: new_id!(),
        };
        let mut assembled_tokens = Vec::new();
        for token in &self.body {
            let tokens = environment.reify_semantic_token(token);
            assembled_tokens.extend(tokens);
        }
        return assembled_tokens;
    }
}


pub struct Environment<'a> {
    pub macro_definitions: &'a IndexMap<String, MacroDefinition>,
    pub label_definitions: &'a IndexMap<String, LabelDefinition>,
    pub arguments: &'a IndexMap<String, Argument>,
    pub id: usize,
}

impl<'a> Environment<'a> {
    // This is only ever called for the highest level body tokens, never for invocations.
    fn reify_semantic_token(&self, token: &SemanticToken) -> Vec<AssembledToken> {
        let mut assembled_tokens = Vec::new();
        match token {
            SemanticToken::Word(pbl) => {
                let word = self.reify_packed_binary_literal(pbl);
                assembled_tokens.push(AssembledToken::Word(word));
            }
            SemanticToken::Invocation(invocation) => {
                match self.reify_invocation(invocation) {
                    Ok(argument) => match argument {
                        Argument::Block(block) => assembled_tokens.extend(block),
                        Argument::Integer(_) => {
                            let variant = AssemblerErrorVariant::NotABlock;
                            let source = invocation.source.clone();
                            let error = AssemblerError { source, variant };
                            assembled_tokens.push(AssembledToken::Error(error))
                        }
                    }
                    Err(error) => assembled_tokens.push(AssembledToken::Error(error)),
                }
            }
            SemanticToken::LabelDefinition(definition) => {
                assembled_tokens.push(AssembledToken::LabelDefinition(definition.clone()));
            }
            SemanticToken::PinnedAddress(address) => {
                assembled_tokens.push(AssembledToken::PinnedAddress(address.clone()));
            }
            SemanticToken::Error(_) => (),
        }
        return assembled_tokens;
    }

    fn reify_packed_binary_literal(&self, pbl: &PackedBinaryLiteral) -> AssembledWord {
        let mut assembled_fields = Vec::new();
        let mut errors = Vec::new();
        for field in &pbl.fields {
            let name = field.name.to_string();
            match self.reify_integer_reference(&name, &field.source) {
                Ok(value) => assembled_fields.push(
                    AssembledField {
                        source: field.source.clone(),
                        value,
                        bits: field.bits,
                        shift: field.shift,
                    }
                ),
                Err(error) => errors.push(error),
            };
        }
        let source = pbl.source.clone();
        let value = pbl.value;
        let bits = pbl.bits;
        AssembledWord { source, bits, fields: assembled_fields, value, errors }
    }

    fn reify_integer_reference(&self, name: &str, source: &SourceSpan) -> Result<IntegerArgument, AssemblerError> {
        match self.reify_reference(name, source)? {
            Argument::Integer(integer) => Ok(integer),
            Argument::Block(_) => Err(
                AssemblerError {
                    source: source.clone(),
                    variant: ErrVar::NotAnInteger,
                }
            ),
        }
    }

    fn reify_reference(&self, name: &str, source: &SourceSpan) -> Result<Argument, AssemblerError> {
        let source = source.clone();
        if let Some(argument) = self.arguments.get(name) {
            Ok(argument.clone())
        } else if let Some(definition) = self.macro_definitions.get(name) {
            self.reify_value(&definition.value)
        } else if let Some(label) = self.label_definitions.get(name) {
            let name = Tracked::from(self.tag_label_name(&label.name), &source);
            Ok(Argument::Integer(IntegerArgument::LabelReference(name)))
        } else {
            let variant = ErrVar::DefinitionNotFound(name.to_string());
            Err(AssemblerError { source, variant })
        }
    }

    fn tag_label_name(&self, name: &str) -> String {
        match name.contains(':') {
            true => format!("{name}:{}", self.id),
            false => name.to_string(),
        }
    }

    fn reify_value(&self, value: &Value) -> Result<Argument, AssemblerError> {
        match value {
            Value::Integer(integer) => {
                let value = match &integer {
                    Integer::Literal(integer) => {
                        IntegerArgument::Integer(integer.clone())
                    }
                    Integer::Expression(expr) => {
                        let expr = self.reify_constant_expression(expr)?;
                        IntegerArgument::Expression(expr)
                    }
                    Integer::LabelReference(name) => {
                        let name = Tracked::from(self.tag_label_name(name), &name.source);
                        IntegerArgument::LabelReference(name)
                    }
                };
                Ok(Argument::Integer(value))
            }
            Value::Block(block) => {
                let mut assembled_tokens = Vec::new();
                for token in block {
                    match &token {
                        SemanticToken::Word(pbl) => {
                            let word = self.reify_packed_binary_literal(pbl);
                            assembled_tokens.push(AssembledToken::Word(word));
                        }
                        SemanticToken::Invocation(invocation) => {
                            match self.reify_invocation(invocation)? {
                                Argument::Block(block) => assembled_tokens.extend(block),
                                Argument::Integer(_) => {
                                    let source = invocation.source.clone();
                                    let variant = AssemblerErrorVariant::IntegerInBlock;
                                    return Err(AssemblerError { source, variant});
                                }
                            }
                        }
                        SemanticToken::LabelDefinition(definition) => {
                            let mut definition = definition.clone();
                            definition.name.push_str(&format!(":{}", self.id));
                            let token = AssembledToken::LabelDefinition(definition);
                            assembled_tokens.push(token);
                        }
                        SemanticToken::PinnedAddress(address) => {
                            let token = AssembledToken::PinnedAddress(address.to_owned());
                            assembled_tokens.push(token);
                        }
                        SemanticToken::Error(_) => (),
                    }
                }
                Ok(Argument::Block(assembled_tokens))
            }
            Value::Invocation(invocation) => {
                self.reify_invocation(invocation)
            }
        }
    }

    fn reify_invocation(&self, invocation: &Invocation) -> Result<Argument, AssemblerError> {
        macro_rules! err {
            ($variant:expr) => { Err(AssemblerError {
                source: invocation.source.clone(), variant: $variant
            }) };
        }
        if let Some(argument) = self.arguments.get(&invocation.name) {
            let expected = 0;
            let received = invocation.arguments.len();
            if received != expected {
                return err!(ErrVar::IncorrectArgumentCount(expected, received));
            }
            Ok(argument.clone())
        } else if let Some(definition) = self.macro_definitions.get(&invocation.name) {
            // Check that the correct number of arguments were provided.
            let received = invocation.arguments.len();
            let expected = definition.arguments.len();
            if received != expected {
                return err!(ErrVar::IncorrectArgumentCount(expected, received));
            }
            let mut arguments = IndexMap::new();
            for (i, argument) in invocation.arguments.iter().enumerate() {
                // Check that the correct types of arguments were provided.
                let arg_invocation = self.reify_value(&argument.value)?;
                let arg_invocation_type = match &arg_invocation {
                    Argument::Integer(_) => ArgumentVariant::Integer,
                    Argument::Block(_) => ArgumentVariant::Block,
                };
                let arg_definition_type = definition.arguments[i].variant;
                if arg_invocation_type != arg_definition_type {
                    let variant = ErrVar::IncorrectArgumentType(
                        arg_definition_type, arg_invocation_type
                    );
                    return Err(AssemblerError { source: argument.source.clone(), variant });
                }
                let name = definition.arguments[i].name.clone();
                arguments.insert(name, arg_invocation);
            }
            let environment = Environment {
                macro_definitions: &self.macro_definitions,
                label_definitions: &self.label_definitions,
                arguments: &arguments,
                id: new_id!(),
            };
            environment.reify_value(&definition.value)
        } else if let Some(label) = self.label_definitions.get(&invocation.name) {
            let expected = 0;
            let received = invocation.arguments.len();
            if received != expected {
                return err!(ErrVar::IncorrectArgumentCount(expected, received));
            }
            let name = Tracked::from(self.tag_label_name(&label.name), &label.source);
            Ok(Argument::Integer(IntegerArgument::LabelReference(name)))
        } else {
            err!(ErrVar::DefinitionNotFound(invocation.name.to_string()))
        }
    }

    fn reify_constant_expression(&self, expr: &Expression) -> Result<AssembledExpression, AssemblerError> {
        use ExpressionTokenVariant as ExprVar;

        let mut assembled_tokens = Vec::new();
        for token in &expr.tokens {
            let assembled_token = match &token.variant {
                ExprVar::Literal(value) => {
                    let source = token.source.clone();
                    let integer = TrackedInteger { source, value: *value };
                    AssembledExpressionToken::Integer(integer)
                }
                ExprVar::Operator(operator) => {
                    AssembledExpressionToken::Operator(*operator)
                }
                ExprVar::Invocation(name) => {
                    match self.reify_integer_reference(&name, &token.source)? {
                        IntegerArgument::LabelReference(name) => {
                            AssembledExpressionToken::LabelReference(name)
                        }
                        IntegerArgument::Integer(integer) => {
                            AssembledExpressionToken::Integer(integer)
                        }
                        IntegerArgument::Expression(expr) => {
                            AssembledExpressionToken::Expression(Box::new(expr))
                        },
                    }
                }
                ExprVar::Error(_) => continue,
            };
            assembled_tokens.push(assembled_token);
        }
        Ok(AssembledExpression { source: expr.source.clone(), tokens: assembled_tokens })
    }
}