use crate::*;

use std::collections::HashMap;


pub fn parse_bytecode(tokens: Vec<Tracked<IntermediateToken>>, width: Option<u32>) -> Result<Vec<Segment>, Vec<Tracked<BytecodeError>>> {
    BytecodeParser::new(width).parse(tokens)
}


pub struct BytecodeParser {
    width: Option<u32>,
    addresses: HashMap<String, Tracked<usize>>,
    address: usize,
    segment_address: usize,
    segment_source: Option<SourceSpan>,
    segments: Vec<Segment>,
    words: Vec<Tracked<Word>>,
    errors: Vec<Tracked<BytecodeError>>,
}

impl BytecodeParser {
    pub fn new(width: Option<u32>) -> Self {
        Self {
            width,
            addresses: HashMap::new(),
            address: 0,
            segment_address: 0,
            segment_source: None,
            segments: Vec::new(),
            words: Vec::new(),
            errors: Vec::new(),
        }
    }

    pub fn parse(mut self, tokens: Vec<Tracked<IntermediateToken>>) -> Result<Vec<Segment>, Vec<Tracked<BytecodeError>>> {
        // Calculate all label addresses ahead of time.
        let mut address = 0;
        for token in &tokens {
            let source = &token.source;
            match &token.value {
                IntermediateToken::LabelDefinition(name) => {
                    let tracked = Tracked::from(address, source.clone());
                    if let Some(_) = self.addresses.insert(name.clone(), tracked) {
                        unreachable!("Uncaught duplicate label definition '{name}'");
                    }
                }
                IntermediateToken::Word(_) => {
                    address += 1;
                }
                IntermediateToken::PinnedAddress(pinned) => {
                    address = pinned.value;
                }
            }
        }
        for token in &tokens {
            let source = &token.source;
            match &token.value {
                IntermediateToken::Word(word) => {
                    let word = self.evaluate_word(word, source);
                    // Check that the word width fits the provided width.
                    if let Some(width) = self.width {
                        if word.width != width {
                            let error = BytecodeError::IncorrectWidth(width, word.width);
                            self.errors.push(Tracked::from(error, source.clone()));
                        }
                    }
                    self.words.push(word);
                    self.address += 1;
                }
                IntermediateToken::PinnedAddress(address) => {
                    let current = self.address;
                    let pinned = address.value;
                    if current > pinned {
                        let error = BytecodeError::PinnedAddressBacktrack(pinned, current);
                        self.errors.push(Tracked::from(error, address.source.clone()));
                    } else {
                        let words = std::mem::take(&mut self.words);
                        if !words.is_empty() {
                            let address = self.segment_address;
                            let source = std::mem::take(&mut self.segment_source);
                            let segment = Segment { address, source, words };
                            self.segments.push(segment);
                        }
                        self.segment_source = Some(address.source.clone());
                        self.address = pinned;
                        self.segment_address = pinned;
                    }
                }
                IntermediateToken::LabelDefinition(_) => (),
            }
        }
        // Finish final segment.
        let words = std::mem::take(&mut self.words);
        if !words.is_empty() {
            let address = self.segment_address;
            let source = std::mem::take(&mut self.segment_source);
            let segment = Segment { address, source, words };
            self.segments.push(segment);
        }

        match self.errors.is_empty() {
            true => Ok(self.segments),
            false => Err(self.errors),
        }
    }

    fn evaluate_expression(&mut self, expression: &IntermediateExpression, source: &SourceSpan) -> isize {
        let mut stack = ExpressionStack::new();
        for token in &expression.tokens {
            let source = &token.source;
            match &token.value {
                IntermediateExpressionToken::Integer(integer) => match integer {
                    IntermediateInteger::Integer(value) => {
                        stack.push(*value);
                    }
                    IntermediateInteger::Expression(expression) => {
                        stack.push(self.evaluate_expression(expression, source));
                    }
                    IntermediateInteger::LabelReference(name) => {
                        stack.push(self.evaluate_label_reference(name));
                    }
                }
                IntermediateExpressionToken::Operator(operator) => {
                    if let Err(err) = stack.apply(*operator, source) {
                        let error = BytecodeError::StackError(err);
                        self.errors.push(Tracked::from(error, source.clone()))
                    }
                }
            }
        }
        match stack.pull_result() {
            Ok(value) => value,
            Err(err) => {
                let error = BytecodeError::StackError(Tracked::from(err, source.clone()));
                self.errors.push(Tracked::from(error, source.clone()));
                0
            }
        }
    }

    fn evaluate_label_reference(&mut self, name: &Tracked<String>) -> isize {
        if let Some(address) = self.addresses.get(&name.to_string()) {
            address.value as isize
        } else {
            unreachable!("Uncaught unresolved label reference '{name}'")
        }
    }

    fn evaluate_word(&mut self, word: &IntermediateWord, source: &SourceSpan) -> Tracked<Word> {
        let mut word_value = word.value;
        for field in &word.fields {
            let field_source = &field.value.value.source;
            let field_value = match &field.value.value.value {
                IntermediateInteger::Expression(expression) => {
                    self.evaluate_expression(expression, source)
                }
                IntermediateInteger::LabelReference(name) => {
                    self.evaluate_label_reference(name)
                }
                IntermediateInteger::Integer(value) => {
                    *value
                }
            };
            let value_width = match field_value.cmp(&0) {
                std::cmp::Ordering::Less => (-field_value).ilog2() + 1,
                std::cmp::Ordering::Equal => 0,
                std::cmp::Ordering::Greater => field_value.ilog2() + 1,
            };
            if field.width < value_width {
                let error = BytecodeError::ValueTooWide(field.width, value_width);
                self.errors.push(Tracked::from(error, field_source.clone()));
            } 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 };
        return Tracked::from(word, source.clone());
    }
}