use crate::vm::{
    reg::*,
    inst::*,
    flags::*,
};
use byteorder::{LE, ReadBytesExt, WriteBytesExt};
use std::{
    mem,
    io::Cursor,
};

pub type Word = u64;
pub type HalfWord = u32;
pub type Registers = [Word; 64];

pub struct Vm {
    code: Vec<Inst>,
    memory: Vec<u8>,
    registers: Registers,
}

#[allow(dead_code)]
impl Vm {
    pub fn new(code: Vec<Inst>, memory: Vec<u8>, registers: Registers) -> Self {
        Vm { code, memory, registers }
    }

    pub fn run(&mut self) {
        while !self.is_halted() {
            self.tick();
        }
    }

    pub fn tick(&mut self) {
        let inst = self.decode();
        let ip = self.ip();
        self.set_reg(IP, ip + 1);
        self.execute(inst);
    }

    pub fn resume(&mut self) {
        self.remove_flags(Flags::HALT);
    }

    pub fn is_halted(&self) -> bool {
        self.flags().contains(Flags::HALT)
    }

    pub fn get_word(&self, addr: Word) -> Word {
        let mut reader = Cursor::new(&self.memory[(addr as usize)..]);
        reader.read_u64::<LE>()
            .expect("word outside of address range")
    }

    pub fn get_halfword(&self, addr: Word) -> HalfWord {
        let mut reader = Cursor::new(&self.memory[(addr as usize)..]);
        reader.read_u32::<LE>()
            .expect("word outside of address range")
    }

    pub fn get_byte(&self, addr: Word) -> u8 {
        let mut reader = Cursor::new(&self.memory[(addr as usize)..]);
        reader.read_u8()
            .expect("word outside of address range")
    }

    pub fn set_word(&mut self, addr: Word, value: Word) {
        let mut writer = Cursor::new(&mut self.memory[(addr as usize)..]);
        writer.write_u64::<LE>(value)
            .expect("word outside of address range");
    }

    pub fn set_halfword(&mut self, addr: Word, value: HalfWord) {
        let mut writer = Cursor::new(&mut self.memory[(addr as usize)..]);
        writer.write_u32::<LE>(value)
            .expect("word outside of address range");
    }

    pub fn set_byte(&mut self, addr: Word, value: u8) {
        let mut writer = Cursor::new(&mut self.memory[(addr as usize)..]);
        writer.write_u8(value)
            .expect("word outside of address range");
    }

    pub fn load(&self, reg: Reg) -> Word {
        self.get_word(self.get_reg(reg))
    }

    pub fn store(&mut self, reg: Reg, value: Word) {
        let addr = self.get_reg(reg);
        self.set_word(addr, value);
    }

    pub fn get_reg(&self, reg: Reg) -> Word {
        self.registers[reg as usize]
    }

    pub fn set_reg(&mut self, reg: Reg, value: Word) -> Word {
        mem::replace(&mut self.registers[reg as usize], value)
    }

    pub fn ip(&self) -> Word {
        self.get_reg(IP)
    }

    pub fn flags(&self) -> Flags {
        // this is safe because it's OK if there are random bits flipped - this shouldn't happen
        // anyway, but if it does, they're ignored
        unsafe { Flags::from_bits_unchecked(self.get_reg(FLAGS)) }
    }
    
    pub fn insert_flags(&mut self, flags: Flags) {
        let mut new_flags = self.flags();
        new_flags.insert(flags);
        self.set_flags(new_flags);
    }

    pub fn remove_flags(&mut self, flags: Flags) {
        let mut new_flags = self.flags();
        new_flags.remove(flags);
        self.set_flags(new_flags);
    }

    pub fn set_flags(&mut self, flags: Flags) {
        self.set_reg(FLAGS, flags.bits());
    }

    pub fn map_flags<F>(&mut self, mapping: F)
        where F: FnOnce(Flags) -> Flags
    {
        let in_flags = self.flags();
        self.set_flags((mapping)(in_flags));
    }

    pub fn decode(&self) -> Inst {
        self.code[self.ip() as usize]
    }

    pub fn execute(&mut self, inst: Inst) {
        macro_rules! map {
            ($r0:ident = $($tail:tt)*) => {{
                let w0 = map!($($tail)+);
                self.set_reg($r0, w0);
            }};

            ($r1:ident $sym:tt $r2:ident) => {{
                let w1 = self.get_reg($r1);
                let w2 = self.get_reg($r2);
                (w1 $sym w2)
            }}
        }

        match inst {
            Inst::Add(r1, r2) => {
                let w1 = self.get_reg(r1);
                let w2 = self.get_reg(r2);
                let value = w1.wrapping_add(w2);
                self.set_reg(r1, value);
            }
            Inst::Mul(r1, r2) => {
                let w1 = self.get_reg(r1);
                let w2 = self.get_reg(r2);
                let value = w1.wrapping_mul(w2);
                self.set_reg(r1, value);
            }
            Inst::Div(r1, r2) => map!(r1 = r1 / r2),
            Inst::Mod(r1, r2) => map!(r1 = r1 % r2),
            Inst::INeg(r1) => {
                let value = self.get_reg(r1);
                let ivalue = -i64::from_ne_bytes(value.to_ne_bytes());
                let uvalue = u64::from_ne_bytes(ivalue.to_ne_bytes());
                self.set_reg(r1, uvalue);
            }
            Inst::And(r1, r2) => map!(r1 = r1 & r2),
            Inst::Or(r1, r2) => map!(r1 = r1 | r2),
            Inst::Xor(r1, r2) => map!(r1 = r1 ^ r2),
            Inst::Shl(r1, r2) => map!(r1 = r1 << r2),
            Inst::Shr(r1, r2) => map!(r1 = r1 >> r2),
            Inst::CmpEq(r1, r2) => {
                if map!(r1 == r2) {
                    self.insert_flags(Flags::COMPARE);
                } else {
                    self.remove_flags(Flags::COMPARE);
                }
            }
            Inst::CmpLt(r1, r2) => {
                if map!(r1 < r2) {
                    self.insert_flags(Flags::COMPARE);
                } else {
                    self.remove_flags(Flags::COMPARE);
                }
            }
            Inst::Jz(r1) => {
                if !self.flags().contains(Flags::COMPARE) {
                    let w1 = self.get_reg(r1);
                    self.set_reg(IP, w1);
                }
            }
            Inst::Jnz(r1) => {
                if self.flags().contains(Flags::COMPARE) {
                    let w1 = self.get_reg(r1);
                    self.set_reg(IP, w1);
                }
            }
            Inst::Load(r1, r2) => {
                let value = self.load(r2);
                self.set_reg(r1, value);
            }
            Inst::Store(r1, r2) => {
                let value = self.get_reg(r1);
                self.store(r2, value);
            }
            Inst::StoreImm(r1, value) => {
                self.set_reg(r1, value as u64);
            }
            Inst::RegCopy(r1, r2) => {
                let w1 = self.get_reg(r2);
                self.set_reg(r1, w1);
            }
            Inst::MemCopy(r1, r2) => {
                let w1 = self.load(r2);
                self.store(r1, w1);
            }
            Inst::Nop => {}
            Inst::Halt => {
                self.insert_flags(Flags::HALT);
            }
        }
    }

    fn with_registers<F, B>(&self, r1: Reg, r2: Reg, map_fn: F) -> B
        where F: FnOnce(Word, Word) -> B
    {
        let w1 = self.get_reg(r1);
        let w2 = self.get_reg(r2);
        (map_fn)(w1, w2)
    }
}

#[cfg(test)]
mod test {
    use super::*;

    macro_rules! itou {
        { $expr:expr } => { u64::from_ne_bytes(($expr).to_ne_bytes()) };
    }

    #[test]
    fn test_arithmetic() {
        use Inst::*;
        use Reg::*;

        let code = vec![
            StoreImm(R01, 25),

            Add(R00, R01),
            Halt,

            Mul(R00, R01),
            Halt,

            Div(R00, R01),
            Halt,

            StoreImm(R01, 15),
            And(R00, R01),
            Halt, // should be 9

            // Storing large numbers
            StoreImm(R00, 0xABCD_EF98),
            StoreImm(R01, 32),
            Shl(R00, R01),
            Or(R00, R01),
            Halt,

            Shr(R00, R01),
            Halt,

            Xor(R00, R00),
            Halt,
        ];

        let mut vm = Vm::new(code, Default::default(), [0; 64]);

        vm.run();
        assert_eq!(vm.get_reg(R00), 25);
        assert_eq!(vm.get_reg(R01), 25);
        vm.resume();

        vm.run();
        assert_eq!(vm.get_reg(R00), 625);
        assert_eq!(vm.get_reg(R01), 25);
        vm.resume();

        vm.run();
        assert_eq!(vm.get_reg(R00), 25);
        assert_eq!(vm.get_reg(R01), 25);
        vm.resume();

        vm.run();
        assert_eq!(vm.get_reg(R00), 9);
        assert_eq!(vm.get_reg(R01), 15);
        vm.resume();

        vm.run();
        assert_eq!(vm.get_reg(R00), 0xABCD_EF98_0000_0020);
        assert_eq!(vm.get_reg(R01), 32);
        vm.resume();

        // TODO : signed instructions

        vm.run();
        assert_eq!(vm.get_reg(R00), 0xABCD_EF98);
        vm.resume();

        vm.run();
        assert_eq!(vm.get_reg(R00), 0);
        vm.resume();
    }
}