use crate::*;

use indexmap::{IndexSet, IndexMap};


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

pub struct LabelAddress {
    pub previous: usize,
    pub current: usize,
    pub touched: bool,
}

#[derive(Eq, Hash, PartialEq, Clone)]
pub struct SymbolSignature {
    pub name: String,
    pub arg_count: usize,
}

impl std::fmt::Display for SymbolSignature {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(f, "{}::{}", self.name, self.arg_count)
    }
}

struct Environment {
    arguments: IndexMap<String, Tracked<IntermediateValue>>,
    id: usize,
}


pub fn parse_intermediate(semantic: Vec<Tracked<SemanticToken>>) -> Result<Vec<Tracked<IntermediateToken>>, Vec<Tracked<IntermediateError>>> {
    IntermediateParser::new(semantic).parse()
}

struct IntermediateParser {
    semantic: Vec<Tracked<SemanticToken>>,
    label_names: IndexSet<Tracked<SymbolSignature>>,
    label_addresses: IndexMap<Tracked<String>, LabelAddress>,
    macro_definitions: IndexMap<SymbolSignature, Tracked<MacroDefinition>>,
    environment_stack: Vec<Environment>,
    address: usize,
    intermediate: Vec<Tracked<IntermediateToken>>,
    errors: Vec<Tracked<IntermediateError>>,
}

impl IntermediateParser {
    /// Initialise a new parser.
    pub fn new(mut semantic: Vec<Tracked<SemanticToken>>) -> Self {
        // Collect all label and macro definitions ahead of time.
        let mut label_names = IndexSet::new();
        let mut macro_definitions = IndexMap::new();
        // Get the signatures of all defined labels.
        for symbol in SymbolParser::new().parse(&semantic) {
            if let SymbolRoleDetailed::LabelDefinition = symbol.role {
                let signature = SymbolSignature { name: symbol.name.clone(), arg_count: symbol.arg_count };
                if !label_names.insert(Tracked::from(signature.clone(), symbol.source)) {
                    unreachable!("Uncaught duplicate label definition '{signature}'");
                }
            }
        }
        // Strip all semantic macro definitions from the semantic tokens.
        let definitions = semantic.extract_if(..,
            |t| if let SemanticToken::MacroDefinition(_) = t.value { true } else { false });
        for definition in definitions {
            let source = definition.source;
            let SemanticToken::MacroDefinition(definition) = definition.value else { unreachable!() };
            let name = definition.name.value.clone();
            let arg_count = definition.arguments.len();
            let signature = SymbolSignature { name, arg_count };
            if macro_definitions.insert(signature.clone(), Tracked::from(definition, source)).is_some() {
                unreachable!("Uncaught duplicate macro definition '{signature}'")
            }
        }
        Self {
            semantic,
            label_names,
            label_addresses: IndexMap::new(),
            macro_definitions,
            environment_stack: Vec::new(),
            address: 0,
            intermediate: Vec::new(),
            errors: Vec::new(),
        }
    }

    pub fn parse(mut self) -> Result<Vec<Tracked<IntermediateToken>>, Vec<Tracked<IntermediateError>>> {
       let semantic = std::mem::take(&mut self.semantic);

       for i in 0..MAX_ITERATIONS_TO_STABILISE {
            info!("Attempting iteration {} of intermediate assembly stage", i+1);
            // Erase the previous parse attempt.
            self.address = 0;
            self.intermediate.clear();
            self.errors.clear();
            self.environment_stack.clear();
            unsafe { ENVIRONMENT_ID = 0; }
            // Update label addresses.
            for (_, address) in &mut self.label_addresses {
                address.previous = address.current;
                address.touched = false;
            }
            // Attempt to parse the program (which is now all block tokens).
            for token in &semantic {
                let source = &token.source;
                let SemanticToken::BlockToken(ref block_token) = token.value else { unreachable!() };
                let env = Environment { arguments: IndexMap::new(), id: next_id!() };
                self.environment_stack.push(env);
                let mut tokens = self.parse_block_token(&block_token, source);
                self.intermediate.append(&mut tokens);
            }
            // Return unrecoverable errors.
            if !self.errors.is_empty() {
                return Err(self.errors);
            }
            // Check label stability
            if self.check_for_instability(false) {
                continue;
            }
            // Program is stable, return.
            info!("Stabilised in iteration {} of intermediate assembly stage", i+1);
            return Ok(self.intermediate);
        }

        self.check_for_instability(true);
        return Err(self.errors);
    }

    /// Check if any label is still stabilising.
    fn check_for_instability(&mut self, create_error: bool) -> bool {
        for (name, address) in &self.label_addresses {
            if address.touched && address.current != address.previous {
                info!("Label '{name}' was unstable, moving from address 0x{:04x} to 0x{:04x}",
                    address.previous, address.current);
                if create_error {
                    let error = IntermediateError::LabelNeverStabilised(name.to_string());
                    self.errors.push(Tracked::from(error, name.source.clone()));
                }
                return true;
            }
        }
        return false;
    }

    /// Get the current environment (the environment at the top of the stack).
    fn env(&self) -> &Environment {
        self.environment_stack.last().unwrap_or_else(||
            unreachable!("No environment on the stack"))
    }

    // Attach the environment ID to a local label name.
    fn tag_name(&self, name: &str) -> String {
        // If a local label belongs to a macro, the name of that macro
        // has been prefixed with the local label name in the
        // resolve_label_name method during the semantic parsing stage,
        // using a ':' character as a separator.
        match name.contains(':') {
            true => format!("{name}:{}", self.env().id),
            false => name.to_string(),
        }
    }

    fn parse_macro_definition_body(&mut self, body: &MacroDefinitionBody, source: &SourceSpan) -> Option<Tracked<IntermediateValue>> {
        match &body {
            MacroDefinitionBody::Integer(integer) => {
                let token = self.parse_integer_token(&integer, &source)?;
                let integer = IntermediateValue::Integer(token.value);
                Some(Tracked::from(integer, source.clone()))
            }
            MacroDefinitionBody::Invocation(invocation) => {
                self.parse_invocation(&invocation, &invocation.source)
            }
            MacroDefinitionBody::Block(blocks) => {
                let mut tokens = Vec::new();
                for block in blocks {
                    tokens.append(&mut self.parse_block_token(block, &block.source));
                }
                let value = IntermediateValue::Block(tokens);
                Some(Tracked::from(value, source.clone()))
            }
            MacroDefinitionBody::List(list) => {
                let list = self.parse_list_token(list, &source)?;
                let integer = IntermediateValue::List(list.value);
                Some(Tracked::from(integer, source.clone()))
            }
        }
    }

    fn parse_block_token(&mut self, block: &BlockToken, source: &SourceSpan) -> Vec<Tracked<IntermediateToken>> {
        let mut intermediate = Vec::new();
        match block {
            BlockToken::LabelDefinition(name) => {
                let signature = SymbolSignature { name: name.to_string(), arg_count: 0 };
                if !self.label_names.contains(&signature) {
                    unreachable!("Unrecognised name for label definition");
                }
                let tagged_name = self.tag_name(name);
                let tracked = Tracked::from(tagged_name.clone(), source.clone());
                self.label_addresses.entry(tracked)
                    .and_modify(|a| {
                        if a.touched { unreachable!("Label '{tagged_name}' was already touched during this cycle.") }
                        a.previous = a.current;
                        a.current = self.address;
                        a.touched = true;
                    })
                    .or_insert(LabelAddress {
                        previous: 0,
                        current: self.address,
                        touched: true,
                    });
            }
            BlockToken::PinnedAddress(integer) => {
                if let Some(pinned) = self.parse_integer_token(integer, &integer.source) {
                    let pinned = **pinned as usize;
                    if pinned < self.address {
                        let error = IntermediateError::PinnedAddressBacktrack(pinned, self.address);
                        self.errors.push(Tracked::from(error, source.clone()));
                    }
                    self.address = pinned;
                    let token = IntermediateToken::PinnedAddress(pinned);
                    intermediate.push(Tracked::from(token, source.clone()));
                }
            }
            BlockToken::ConditionalBlock(cond) => {
                if let Some(predicate) = self.parse_integer_token(&cond.predicate, &cond.predicate.source) {
                    if **predicate != 0 {
                        if let Some(label_source) = block_contains_label_definition(&cond.body, &cond.body.source) {
                            let error = IntermediateError::LabelDefinitionInConditionBody;
                            let new_source = cond.body.source.clone().wrap(label_source);
                            self.errors.push(Tracked::from(error, new_source));
                        } else {
                            let mut body = self.parse_block_token(&cond.body, &cond.body.source);
                            intermediate.append(&mut body);
                        }
                    }
                }
            }
            BlockToken::WordTemplate(word_template) => {
                let word_width = word_template.width;
                let mut word_value = word_template.value;
                for field in &word_template.fields {
                    let name = field.name.to_string();
                    let field_source = &field.source;
                    let invocation = Invocation { name, arguments: Vec::new() };
                    if let Some(result) = self.parse_integer_invocation(&invocation, field_source) {
                        let field_value = **result;
                        // Inject field value into real value.
                        let value_width = width(field_value);
                        if field.width < value_width {
                            let error = IntermediateError::ValueTooWide(field.width, value_width);
                            let new_source = field_source.wrap(result.source);
                            self.errors.push(Tracked::from(error, new_source));
                        } else {
                            let mask = 2_usize.pow(field.width as u32) - 1;
                            let clamped_value = (field_value as usize) & mask;
                            word_value |= (clamped_value << field.shift) as usize;
                        }
                    }
                }
                let word = Word { width: word_width, value: word_value };
                let token = IntermediateToken::Word(word);
                intermediate.push(Tracked::from(token, source.clone()));
                self.address += 1;
            }
            BlockToken::Block(blocks) => {
                for block in blocks {
                    let mut tokens = self.parse_block_token(block, &block.source);
                    intermediate.append(&mut tokens);
                }
            }
            BlockToken::Invocation(invocation) => {
                if let Some(mut tokens) = self.parse_block_invocation(invocation, source) {
                    intermediate.append(&mut tokens);
                }
            }
        }
        return intermediate;
    }

    fn parse_integer_token(&mut self, integer: &IntegerToken, source: &SourceSpan) -> Option<Tracked<IntermediateInteger>> {
        match integer {
            IntegerToken::IntegerLiteral(value) => {
                let tracked = Tracked::from(*value, source.clone());
                Some(Tracked::from(tracked, source.clone()))
            }
            IntegerToken::Expression(expression) => {
                self.parse_expression(expression, source)
            }
            IntegerToken::Invocation(invocation) => {
                self.parse_integer_invocation(invocation, source)
            }
        }
    }

    fn parse_list_token(&mut self, list: &ListToken, source: &SourceSpan) -> Option<Tracked<IntermediateList>> {
        match list {
            ListToken::StringLiteral(literal) => {
                Some(Tracked::from(literal.chars.clone(), source.clone()))
            }
            ListToken::ListLiteral(literal) => {
                let mut integers = Vec::new();
                for token in literal {
                    let integer = self.parse_integer_token(&token.value, &token.source)?;
                    integers.push(integer.value);
                }
                Some(Tracked::from(integers, source.clone()))
            }
            ListToken::Invocation(invocation) => {
                self.parse_list_invocation(&invocation, source)
            }
        }
    }

    fn parse_integer_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<IntermediateInteger>> {
        let result = self.parse_invocation(invocation, source)?;
        match result.value {
            IntermediateValue::Integer(integer) => {
                let source = integer.source.clone();
                Some(Tracked::from(integer, source))
            }
            IntermediateValue::Block(_) | IntermediateValue::List(_) => {
                let error = IntermediateError::ExpectedInteger;
                self.errors.push(Tracked::from(error, result.source));
                None
            }
        }
    }

    fn parse_block_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Vec<Tracked<IntermediateToken>>> {
        let result = self.parse_invocation(invocation, source)?;
        let source = result.source;
        match result.value {
            IntermediateValue::Block(tokens) => Some(tokens),
            IntermediateValue::Integer(_) | IntermediateValue::List(_) => {
                let error = IntermediateError::ExpectedBlock;
                self.errors.push(Tracked::from(error, source));
                None
            }
        }
    }

    fn parse_list_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<IntermediateList>> {
        let result = self.parse_invocation(invocation, source)?;
        let source = result.source;
        match result.value {
            IntermediateValue::List(literal) => Some(Tracked::from(literal, source)),
            IntermediateValue::Integer(_) | IntermediateValue::Block(_) => {
                let error = IntermediateError::ExpectedList;
                self.errors.push(Tracked::from(error, source));
                None
            }
        }
    }

    fn parse_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<IntermediateValue>> {
        let received_count = invocation.arguments.len();
        let signature = SymbolSignature {
            name: invocation.name.clone(),
            arg_count: invocation.arguments.len(),
        };
        if let Some(argument) = self.env().arguments.get(&invocation.name) {
            // This invocation is a macro argument.
            // Check that no arguments were provided.
            if received_count != 0 {
                let error = IntermediateError::IncorrectArgumentCount(0, received_count);
                self.errors.push(Tracked::from(error, source.clone()));
                None
            } else {
                Some(argument.clone())
            }
        } else if let Some(label) = self.label_names.get(&signature) {
            // This invocation is a label reference.
            // Check that no arguments were provided.
            if received_count != 0 {
                let error = IntermediateError::IncorrectArgumentCount(0, received_count);
                self.errors.push(Tracked::from(error, source.clone()));
                None
            } else {
                let tagged_name = self.tag_name(&signature.name);
                let address = self.label_addresses.get(&tagged_name)
                    .and_then(|a| Some(a.current)).or(Some(0)).unwrap();
                let tracked = Tracked::from(address as isize, label.source.clone());
                let integer = IntermediateValue::Integer(tracked);
                Some(Tracked::from(integer, source.clone()))
            }
        } else if let Some(definition) = self.macro_definitions.get(&signature) {
            // This invocation is a macro reference.
            let definition = definition.clone();
            // Check that the correct number of arguments were provided.
            let expected_count = definition.arguments.len();
            if received_count != expected_count {
                let error = IntermediateError::IncorrectArgumentCount(expected_count, received_count);
                self.errors.push(Tracked::from(error, source.clone()));
                None
            } else {
                // Gather and type-check the provided arguments.
                let mut arguments = Vec::new();
                for (i, argument) in invocation.arguments.iter().enumerate() {
                    let expected_type = definition.arguments[i].variant;
                    let received_value = self.parse_invocation_argument(argument, expected_type)?;
                    arguments.push(received_value);
                }
                // Invoke the invocation multiple times.
                let invocations = arguments.iter().map(|a| a.len()).max().unwrap_or(1);
                let mut values = Vec::new();
                for i in 0..invocations {
                    // Construct an argument map for this invocation.
                    let mut argument_map = IndexMap::new();
                    for (a, argument) in arguments.iter().enumerate() {
                        let name = definition.arguments[a].name.clone();
                        let source = invocation.arguments[a].source.clone();
                        let value = match argument {
                            RepeatedArgument::Loop(value) => {
                                Tracked::from(value.clone(), source)
                            }
                            RepeatedArgument::List(list) => match list.get(i) {
                                Some(value) => {
                                    Tracked::from(value.clone(), source)
                                }
                                None => {
                                    let error = IntermediateError::ListExhausted;
                                    let source = invocation.arguments[a].source.clone();
                                    self.errors.push(Tracked::from(error, source));
                                    return None;
                                }
                            }
                        };
                        if argument_map.insert(name.clone(), value).is_some() {
                            unreachable!("Uncaught duplicate macro argument name '{name}'");
                        };
                    }
                    // Invoke the macro once.
                    let env = Environment { arguments: argument_map, id: next_id!() };
                    self.environment_stack.push(env);
                    let result = self.parse_macro_definition_body(&definition.body, source);
                    self.environment_stack.pop().unwrap();
                    values.push(result?);
                }
                if values.len() == 1 {
                    // If the macro was invoked once, return the value.
                    values.pop()
                } else {
                    // If the macro was invoked multiple times, create a list of
                    // block tokens from the returned values.
                    let mut block = Vec::new();
                    for value in values {
                        match value.value {
                            IntermediateValue::Integer(_) | IntermediateValue::List(_) => {
                                let error = IntermediateError::ExpectedBlock;
                                self.errors.push(Tracked::from(error, value.source));
                                return None;
                            }
                            IntermediateValue::Block(mut tokens) => {
                                block.append(&mut tokens);
                            }
                        }
                    }
                    Some(Tracked::from(IntermediateValue::Block(block), source.clone()))
                }
            }
        } else {
            unreachable!("Uncaught unresolved reference '{signature}'");
        }
    }

    fn parse_invocation_argument(&mut self, argument: &Tracked<InvocationArgument>, expected_type: ArgumentType) -> Option<RepeatedArgument> {
        let source = &argument.source;
        let (received_value, received_type) = match &argument.value {
            InvocationArgument::ListToken(list) => {
                let list = self.parse_list_token(list, source)?;
                self.parse_invocation_list_argument(list, expected_type)?
            }
            InvocationArgument::IntegerToken(integer) => {
                let tracked = self.parse_integer_token(&integer, &argument.source)?;
                let value = IntermediateValue::Integer(tracked.value);
                (RepeatedArgument::Loop(value), ArgumentType::Integer)
            }
            InvocationArgument::BlockToken(block) => {
                let tokens = self.parse_block_token(&block, &argument.source);
                let value = IntermediateValue::Block(tokens);
                (RepeatedArgument::Loop(value), ArgumentType::Block)
            }
            InvocationArgument::Invocation(invocation) => {
                let value = self.parse_invocation(&invocation, &argument.source)?;
                match value.value {
                    IntermediateValue::Integer(_) =>
                        (RepeatedArgument::Loop(value.value), ArgumentType::Integer),
                    IntermediateValue::Block(_) =>
                        (RepeatedArgument::Loop(value.value), ArgumentType::Block),
                    IntermediateValue::List(list) =>
                        self.parse_invocation_list_argument(Tracked::from(list, value.source), expected_type)?
                }
            }
        };
        if expected_type != received_type {
            let error = IntermediateError::IncorrectArgumentType(expected_type, received_type);
            self.errors.push(Tracked::from(error, argument.source.clone()));
            return None;
        }
        return Some(received_value);
    }

    fn parse_invocation_list_argument(&mut self, list: Tracked<IntermediateList>, expected_type: ArgumentType) -> Option<(RepeatedArgument, ArgumentType)> {
        if let ArgumentType::Integer = expected_type {
            let mut values = Vec::new();
            for value in &list.value {
                values.push(IntermediateValue::Integer(value.clone()));
            }
            Some((RepeatedArgument::List(values), ArgumentType::Integer))
        } else {
            let value = IntermediateValue::List(list.value);
            Some((RepeatedArgument::Loop(value), ArgumentType::List))
        }
    }

    fn parse_expression(&mut self, expression: &Expression, source: &SourceSpan) -> Option<Tracked<IntermediateInteger>> {
        let mut stack = ExpressionStack::new();
        for token in &expression.tokens {
            let source = &token.source;
            match &token.value {
                ExpressionToken::IntegerToken(token) => {
                    let integer = self.parse_integer_token(token, source)?;
                    stack.push(IntermediateValue::Integer(integer.value));
                }
                ExpressionToken::ListToken(token) => {
                    let list = self.parse_list_token(token, source)?;
                    stack.push(IntermediateValue::List(list.value));
                }
                ExpressionToken::Invocation(invocation) => {
                    if let Some(value) = self.parse_invocation(invocation, source) {
                        if let IntermediateValue::Block(_) = &value.value {
                            let error = IntermediateError::InvalidBlockInExpression;
                            let new_source = source.clone().wrap(value.source.clone());
                            self.errors.push(Tracked::from(error, new_source));
                            return None;
                        }
                        stack.push(value.value);
                    }
                }
                ExpressionToken::Operator(operator) => {
                    if let Err(expr_error) = stack.apply(*operator, source) {
                        let error = IntermediateError::ExpressionError(expr_error);
                        self.errors.push(Tracked::from(error, source.clone()));
                        return None;
                    }
                }
            }
        }
        match stack.pull_result() {
            Ok(value) => {
                let tracked = Tracked::from(value, source.clone());
                Some(Tracked::from(tracked, source.clone()))
            }
            Err(expr_error) => {
                let tracked = Tracked::from(expr_error, source.clone());
                let error = IntermediateError::ExpressionError(tracked);
                self.errors.push(Tracked::from(error, source.clone()));
                None
            }
        }
    }
}



macro_rules! return_some {
    ($option:expr) => {
        if $option.is_some() { return $option; }
    };
}


// Check if a block token contains a label definition.
fn block_contains_label_definition(block: &BlockToken, source: &SourceSpan) -> Option<SourceSpan> {
    match &block {
        BlockToken::LabelDefinition(_) =>
            return Some(source.clone()),
        BlockToken::Invocation(invocation) =>
            return_some!(invocation_contains_label_definition(invocation)),
        BlockToken::Block(blocks) =>
            for block in blocks {
                return_some!(block_contains_label_definition(block, &block.source))
            },
        _ => (),
    }
    return None;
}

// Check if the arguments passed to an invocation contain a label definition.
fn invocation_contains_label_definition(invocation: &Invocation) -> Option<SourceSpan> {
    for argument in &invocation.arguments {
        match &argument.value {
            InvocationArgument::BlockToken(block) => {
                return_some!(block_contains_label_definition(&block, &argument.source))
            }
            InvocationArgument::Invocation(invocation) => {
                return_some!(invocation_contains_label_definition(&invocation))
            }
            _ => (),
        }
    }
    return None;
}