use crate::*;

use std::collections::HashMap;


/// Doesn't truncate trailing null bytes.
pub fn generate_bytecode(semantic: &Program) -> AssembledProgram {
    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 });
    }
    AssembledProgram {
        bytecode: generator.bytecode,
        symbols,
    }
}


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

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 DefinitionKind::LabelDefinition = definition.kind {
                let key = name.to_string();
                let value = LabelInformation { address: 0, slots: Vec::new() };
                labels.insert(key, value);
            }
        }
        Self {
            definitions,
            labels,
            stack: Vec::new(),
            bytecode: 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::Comment(_) => (),

                SemanticToken::LabelDefinition(name) => if in_macro {
                    unreachable!("Uncaught label definition in macro");
                } else {
                    let information = self.labels.get_mut(name).unwrap();
                    information.address = i;
                }
                SemanticToken::MacroDefinition{ .. } => if in_macro {
                    unreachable!("Uncaught macro definition in macro");
                }

                SemanticToken::RawValue(value) => match value {
                    Value::Byte(byte) => byte!(*byte),
                    Value::Double(double) => double!(*double),
                }
                SemanticToken::Instruction(instruction) => {
                    byte!(instruction.value)
                }
                SemanticToken::Invocation(name) => {
                    if let Some(definition) = self.definitions.get(name) {
                        match &definition.kind {
                            DefinitionKind::MacroDefinition(body) => {
                                self.parse(body, true);
                            }
                            DefinitionKind::LabelDefinition => {
                                let information = self.labels.get_mut(name).unwrap();
                                information.slots.push(i);
                                double!(0);
                            }
                        }
                    } else {
                        unreachable!("Uncaught undefined symbol '{name}'");
                    }
                }

                SemanticToken::Padding(value) => {
                    self.bytecode.resize(i + usize::from(value), 0);
                },
                SemanticToken::String(bytes) => {
                    self.bytecode.extend_from_slice(bytes)
                },

                SemanticToken::BlockOpen(_) => {
                    self.stack.push(i);
                    double!(0);
                }
                SemanticToken::BlockClose(_) => {
                    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;
                }
            }
        }
    }

    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;
            }
        }
    }
}