use std::collections::HashMap;

use crate::{
    operation::{Operation, RegisterOperation},
    parser::{parse, Instruction, RegisterContents, ResultValue, Value},
    Error,
};

fn push_u16(bytecode: &mut Vec<u8>, num: u16) {
    bytecode.push((num >> 8) as u8);
    bytecode.push(num as u8);
}

/// Generate bytecode from a single line of code.
fn compile(left: u8, op: Operation, right: u16, is_register: bool) -> Result<Vec<u8>, String> {
    let mut bytecode = Vec::new();
    bytecode.push(left);
    bytecode.push(op.to_u8() << 5 | is_register as u8);
    push_u16(&mut bytecode, right);
    Ok(bytecode)
}

/// Assemble the hase programm and return the bytecode.
pub fn assemble(source: &str) -> Result<Vec<u8>, Error> {
    let (lines, register_definitions) = parse(source)?;
    let mut bytecode = Vec::new();
    bytecode.push(register_definitions.len().try_into().unwrap());
    for register in register_definitions.iter() {
        push_u16(
            &mut bytecode,
            match &register.1.contents {
                RegisterContents::Size(size) => *size,
                RegisterContents::Values(values) => values.len() as u16,
            },
        );
    }
    for (num, register) in register_definitions.iter().enumerate() {
        if let RegisterContents::Values(values) = &register.1.contents {
            let reg_num = (num + 1).try_into().unwrap();
            for value in values {
                bytecode.append(
                    &mut compile(
                        reg_num,
                        Operation::Register(RegisterOperation::Add),
                        *value,
                        false,
                    )
                    .map_err(|e| Error {
                        message: e,
                        line: register.0,
                    })?,
                );
                if values.len() > 1 {
                    bytecode.append(
                        &mut compile(
                            reg_num,
                            Operation::Register(RegisterOperation::Right),
                            1,
                            false,
                        )
                        .map_err(|e| Error {
                            message: e,
                            line: register.0,
                        })?,
                    );
                }
            }
        }
    }
    let registers = register_definitions.iter().enumerate().try_fold(
        HashMap::<char, u8>::new(),
        |mut m, (num, (line_num, reg))| {
            if m.get(&reg.name).is_some() {
                Err(Error {
                    message: format!("Register {} already defined", reg.name),
                    line: *line_num,
                })
            } else {
                m.insert(reg.name, (num + 1).try_into().unwrap());
                Ok(m)
            }
        },
    )?;
    let register_definition_instructions = (bytecode.len() - 1 - (registers.len() * 2)) / 4;
    let found_labels = lines
        .iter()
        .enumerate()
        .filter_map(|(num, (line_num, line))| {
            line.label
                .map(|name| (name, register_definition_instructions + num, line_num))
        })
        .try_fold(
            HashMap::<u16, usize>::new(),
            |mut m, (name, offset, line_num)| {
                if m.get(&(name as u16)).is_some() {
                    Err(Error {
                        message: format!("Label {} already defined", name),
                        line: *line_num,
                    })
                } else {
                    m.insert(name.try_into().unwrap(), offset);
                    Ok(m)
                }
            },
        )?;
    for (num, (line_num, line)) in lines.into_iter().enumerate() {
        let register = match line.instruction {
            Instruction::Register(name, _, _) => Some(name),
            Instruction::Jump(reg, _) => reg,
            Instruction::Take(reg, _) => Some(reg),
            Instruction::Result(ResultValue::Register(r)) => Some(r),
            Instruction::Result(ResultValue::Value(_)) => None,
            Instruction::Result(ResultValue::All) => None,
        }
        .map(|name| {
            registers.get(&name).cloned().ok_or_else(|| Error {
                message: format!("Referencing undeclared register: {}", name),
                line: line_num,
            })
        })
        .unwrap_or(Ok(0))?;

        let (right, is_register) = match line.instruction {
            Instruction::Register(_, _, value) => Some(value),
            Instruction::Result(ResultValue::Register(reg)) => Some(Value::Register(reg)),
            Instruction::Result(ResultValue::Value(value)) => Some(value),
            Instruction::Result(ResultValue::All) => None,
            Instruction::Jump(_, to) => Some(Value::Number(to.try_into().unwrap())),
            Instruction::Take(_, reg) => Some(Value::Register(reg)),
        }
        .map(|value| match value {
            Value::Register(name) => match registers.get(&name) {
                Some(reg) => Ok((*reg as u16, true)),
                None => panic!(
                    "Referencing undeclared register: {} at line {line_num}",
                    name
                ),
            },
            Value::Number(value) => {
                if matches!(line.instruction, Instruction::Jump(_, _)) {
                    found_labels
                        .get(&value)
                        .ok_or(Error {
                            message: format!("Label {value} not defined"),
                            line: line_num,
                        })
                        .map(|v| (*v as u16, false))
                } else {
                    Ok((value, false))
                }
            }
        })
        .unwrap_or(Ok((0, false)))?;
        let op = match line.instruction {
            Instruction::Register(_, op, _) => Operation::Register(op),
            Instruction::Result(_) => Operation::Result,
            Instruction::Jump(_, _) => Operation::Jump,
            Instruction::Take(_, _) => Operation::Take,
        };
        bytecode.append(
            &mut compile(register, op, right, is_register).map_err(|e| Error {
                message: e,
                line: num,
            })?,
        );
    }
    Ok(bytecode)
}