use crate::*;

use std::collections::HashMap;


/// Doesn't truncate trailing null bytes.
pub fn generate_bytecode(semantic: &Program) -> Result<AssembledProgram, Vec<Tracked<BytecodeError>>> {
    let mut generator = BytecodeGenerator::new(&semantic.definitions);
    generator.parse(&semantic.tokens, false);
    generator.fill_slots();
    let mut symbols = Vec::new();
    for (name, information) in generator.labels {
        let source = semantic.definitions.get(&name).unwrap().source.clone();
        let address = information.address;
        symbols.push(AssembledSymbol { name, address, source });
    }

    match generator.errors.is_empty() {
        true => Ok(
            AssembledProgram {
                bytecode: generator.bytecode,
                symbols,
            }
        ),
        false => Err(generator.errors),
    }
}


pub struct BytecodeGenerator<'a> {
    definitions: &'a HashMap<String, Tracked<Definition>>,
    labels: HashMap<String, LabelInformation>,
    stack: Vec<usize>,
    bytecode: Vec<u8>,
    errors: Vec<Tracked<BytecodeError>>,
}

struct LabelInformation {
    address: usize,
    slots: Vec<usize>,
}

impl<'a> BytecodeGenerator<'a> {
    pub fn new(definitions: &'a HashMap<String, Tracked<Definition>>) -> Self {
        let mut labels = HashMap::new();
        for (name, definition) in definitions {
            if let DefinitionVariant::LabelDefinition = definition.variant {
                // Use fake address for now.
                let information = LabelInformation { address: 0, slots: Vec::new() };
                labels.insert(name.to_string(), information);
            }
        }
        Self {
            definitions,
            labels,
            stack: Vec::new(),
            bytecode: Vec::new(),
            errors: Vec::new(),
        }
    }

    pub fn parse(&mut self, tokens: &[Tracked<SemanticToken>], in_macro: bool) {
        macro_rules! byte {
            ($byte:expr) => {
                self.bytecode.push($byte)
            };
        }
        macro_rules! double {
            ($double:expr) => {{
                let [high, low] = u16::to_be_bytes($double);
                self.bytecode.push(high); self.bytecode.push(low);
            }};
        }

        for token in tokens {
            let i = self.bytecode.len();
            match &token.value {
                SemanticToken::Literal(value) => match value {
                    Value::Byte(byte) => byte!(*byte),
                    Value::Double(double) => double!(*double),
                }
                SemanticToken::Pad(value) => {
                    self.bytecode.resize(i + usize::from(value), 0);
                },
                SemanticToken::String(bytes) => {
                    self.bytecode.extend_from_slice(bytes)
                },
                SemanticToken::Comment(_) => (),
                SemanticToken::BlockOpen(_) => {
                    self.stack.push(i);
                    // Use a fake index for now.
                    double!(0);
                }
                SemanticToken::BlockClose(_) => {
                    if i > 0xFFFF {
                        let error = BytecodeError::InvalidBlockAddress(i);
                        self.errors.push(Tracked::from(error, token.source.clone()));
                    }
                    let Some(addr) = self.stack.pop() else {
                        unreachable!("Uncaught unmatched block terminator");
                    };
                    let [high, low] = (i as u16).to_be_bytes();
                    self.bytecode[addr] = high;
                    self.bytecode[addr+1] = low;
                }
                SemanticToken::Symbol(name) => {
                    if let Some(definition) = self.definitions.get(name) {
                        match &definition.variant {
                            DefinitionVariant::MacroDefinition(body) => {
                                self.parse(body, true);
                            }
                            DefinitionVariant::LabelDefinition => {
                                let information = self.labels.get_mut(name).unwrap();
                                information.slots.push(i);
                                // Use a fake index for now.
                                double!(0);
                            }
                        }
                    } else {
                        unreachable!("Uncaught undefined symbol '{name}'");
                    }
                }
                SemanticToken::Instruction(instruction) => {
                    byte!(instruction.value)
                }
                SemanticToken::LabelDefinition(name) => if in_macro {
                    unreachable!("Uncaught label definition in macro");
                } else {
                    if i > 0xFFFF {
                        let error = BytecodeError::InvalidLabelAddress(i);
                        self.errors.push(Tracked::from(error, token.source.clone()));
                    }
                    let information = self.labels.get_mut(name).unwrap();
                    // Replace fake index with real index.
                    information.address = i;
                }
                SemanticToken::MacroDefinition{ .. } => if in_macro {
                    unreachable!("Uncaught macro definition in macro");
                }
            }
        }

        if !self.stack.is_empty() {
            unreachable!("Uncaught unterminated block");
        }
    }

    /// Fill each label slot with a real label address.
    pub fn fill_slots(&mut self) {
        for information in self.labels.values() {
            let [high, low] = (information.address as u16).to_be_bytes();
            for addr in &information.slots {
                self.bytecode[*addr] = high;
                self.bytecode[*addr + 1] = low;
            }
        }
    }
}