#![feature(macro_metavar_expr)]

use mnemonics::mnemonics;

mod devices;
mod pc;
mod ram;
mod stack;
mod value;

use self::devices::*;
use self::pc::*;
use self::ram::*;
use self::stack::*;
use self::value::*;

pub struct Uxn {
    pub wst: Stack,
    pub rst: Stack,
    pub pc: u16,
    pub ram: [u8; 65536],
}

pub trait DeviceHandler {
    fn device_in(&mut self, uxn: &mut Uxn, addr: u8) -> u8;
    fn device_out(&mut self, uxn: &mut Uxn, addr: u8, val: u8);
}

impl Uxn {
    pub fn new() -> Self {
        Self {
            wst: Stack::default(),
            rst: Stack::default(),
            pc: Default::default(),
            ram: [Default::default(); 65536],
        }
    }

    #[must_use]
    pub fn with_ram(mut self, ram: [u8; 65536]) -> Self {
        self.ram = ram;
        self
    }
}

#[derive(Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum Trap {
    StackUnderflow = 1,
    StackOverflow = 2,
    DivisionByZero = 3,
}

impl Uxn {
    pub fn run<D: DeviceHandler>(&mut self, from: u16, device_handler: &mut D) -> Result<(), Trap> {
        self.pc = from;

        let mut executor = Executor {
            dev: Devices::new(device_handler),
            uxn: self,
        };

        executor.run()
    }
}

struct Executor<'a, D> {
    uxn: &'a mut Uxn,
    dev: Devices<'a, D>,
}

enum ControlFlow<T> {
    Next(T),
    Break(T),
}

use ControlFlow::*;

impl<T> ControlFlow<T> {
    fn decouple(self) -> (ControlFlow<()>, T) {
        match self {
            Self::Next(t) => (ControlFlow::Next(()), t),
            Self::Break(t) => (ControlFlow::Break(()), t),
        }
    }
}

impl<'a, D: DeviceHandler> Executor<'a, D> {
    pub fn run(&mut self) -> Result<(), Trap> {
        loop {
            let opcode: u8 = self.uxn.ram.peek::<Byte>(Short::from(self.uxn.pc)).into();
            self.uxn.pc.advance_by(1u16);

            if let Break(()) = Self::INSTR_FNS[opcode as usize](self)? {
                return Ok(());
            }
        }
    }

    fn map<const KEEP: bool, const RETURN: bool, const CROSS_STACK: bool, I, O, F>(
        &mut self,
        f: F,
    ) -> Result<ControlFlow<()>, Trap>
    where
        I: Pop + StackPeek,
        O: Push,
        F: FnOnce(&mut Self, I) -> Result<ControlFlow<O>, Trap>,
    {
        let input = match (KEEP, RETURN) {
            (true, true) => self.uxn.rst.peek()?,
            (true, false) => self.uxn.wst.peek()?,
            (false, true) => self.uxn.rst.pop()?,
            (false, false) => self.uxn.wst.pop()?,
        };

        let (control_flow, output) = f(self, input)?.decouple();

        match (RETURN, CROSS_STACK) {
            (true, true) => self.uxn.wst.push(output)?,
            (true, false) => self.uxn.rst.push(output)?,
            (false, true) => self.uxn.rst.push(output)?,
            (false, false) => self.uxn.wst.push(output)?,
        }

        Ok(control_flow)
    }

    const INSTR_FNS: [fn(&mut Self) -> Result<ControlFlow<()>, Trap>; 256] = mnemonics![
        /*      00      01      02      03      04      05      06      07      08      09      0a      0b      0c      0d      0e      0f    */
        /* 00 */BRK     INC     POP     NIP     SWP     ROT     DUP     OVR     EQU     NEQ     GTH     LTH     JMP     JCN     JSRx    STHx
        /* 10 */LDZ     STZ     LDR     STR     LDA     STA     DEI     DEO     ADD     SUB     MUL     DIV     AND     ORA     EOR     SFT
        /* 20 */JCI     INC2    POP2    NIP2    SWP2    ROT2    DUP2    OVR2    EQU2    NEQ2    GTH2    LTH2    JMP2    JCN2    JSR2x   STH2x
        /* 30 */LDZ2    STZ2    LDR2    STR2    LDA2    STA2    DEI2    DEO2    ADD2    SUB2    MUL2    DIV2    AND2    ORA2    EOR2    SFT2
        /* 40 */JMI     INCr    POPr    NIPr    SWPr    ROTr    DUPr    OVRr    EQUr    NEQr    GTHr    LTHr    JMPr    JCNr    JSRrx   STHrx
        /* 50 */LDZr    STZr    LDRr    STRr    LDAr    STAr    DEIr    DEOr    ADDr    SUBr    MULr    DIVr    ANDr    ORAr    EORr    SFTr
        /* 60 */JSI     INC2r   POP2r   NIP2r   SWP2r   ROT2r   DUP2r   OVR2r   EQU2r   NEQ2r   GTH2r   LTH2r   JMP2r   JCN2r   JSR2rx  STH2rx
        /* 70 */LDZ2r   STZ2r   LDR2r   STR2r   LDA2r   STA2r   DEI2r   DEO2r   ADD2r   SUB2r   MUL2r   DIV2r   AND2r   ORA2r   EOR2r   SFT2r
        /* 80 */LIT     INCk    POPk    NIPk    SWPk    ROTk    DUPk    OVRk    EQUk    NEQk    GTHk    LTHk    JMPk    JCNk    JSRkx   STHkx
        /* 90 */LDZk    STZk    LDRk    STRk    LDAk    STAk    DEIk    DEOk    ADDk    SUBk    MULk    DIVk    ANDk    ORAk    EORk    SFTk
        /* a0 */LIT2    INC2k   POP2k   NIP2k   SWP2k   ROT2k   DUP2k   OVR2k   EQU2k   NEQ2k   GTH2k   LTH2k   JMP2k   JCN2k   JSR2kx  STH2kx
        /* b0 */LDZ2k   STZ2k   LDR2k   STR2k   LDA2k   STA2k   DEI2k   DEO2k   ADD2k   SUB2k   MUL2k   DIV2k   AND2k   ORA2k   EOR2k   SFT2k
        /* c0 */LITr    INCkr   POPkr   NIPkr   SWPkr   ROTkr   DUPkr   OVRkr   EQUkr   NEQkr   GTHkr   LTHkr   JMPkr   JCNkr   JSRkrx  STHkrx
        /* d0 */LDZkr   STZkr   LDRkr   STRkr   LDAkr   STAkr   DEIkr   DEOkr   ADDkr   SUBkr   MULkr   DIVkr   ANDkr   ORAkr   EORkr   SFTkr
        /* e0 */LIT2r   INC2kr  POP2kr  NIP2kr  SWP2kr  ROT2kr  DUP2kr  OVR2kr  EQU2kr  NEQ2kr  GTH2kr  LTH2kr  JMP2kr  JCN2kr  JSR2krx STH2krx
        /* f0 */LDZ2kr  STZ2kr  LDR2kr  STR2kr  LDA2kr  STA2kr  DEI2kr  DEO2kr  ADD2kr  SUB2kr  MUL2kr  DIV2kr  AND2kr  ORA2kr  EOR2kr  SFT2kr
    ];
}

impl<D: DeviceHandler> Executor<'_, D> {
    fn brk<V: Value>(&mut self, (): ()) -> Result<ControlFlow<()>, Trap> {
        Ok(Break(()))
    }

    fn jci<V: Value>(&mut self, cond: bool) -> Result<ControlFlow<()>, Trap> {
        let offs = self.uxn.ram.peek::<Short>(Short::from(self.uxn.pc));
        self.uxn.pc.advance_by(2u16);
        if cond {
            self.uxn.pc.advance_by(offs)
        }
        Ok(Next(()))
    }

    fn jmi<V: Value>(&mut self, (): ()) -> Result<ControlFlow<()>, Trap> {
        let offs = self.uxn.ram.peek::<Short>(Short::from(self.uxn.pc));
        self.uxn.pc.advance_by(2u16);
        self.uxn.pc.advance_by(offs);
        Ok(Next(()))
    }

    fn jsi<V: Value>(&mut self, (): ()) -> Result<ControlFlow<Short>, Trap> {
        let offs = self.uxn.ram.peek::<Short>(Short::from(self.uxn.pc));
        self.uxn.pc.advance_by(2u16);
        let ret = Short::from(self.uxn.pc);
        self.uxn.pc.advance_by(offs);
        Ok(Next(ret))
    }

    fn lit<V: Value>(&mut self, (): ()) -> Result<ControlFlow<V>, Trap> {
        let val = self.uxn.ram.peek(Short::from(self.uxn.pc));
        self.uxn.pc.advance_by(V::size());
        Ok(Next(val))
    }

    fn inc<V: Value>(&mut self, a: V) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a + 1))
    }

    fn pop<V: Value>(&mut self, _: V) -> Result<ControlFlow<()>, Trap> {
        Ok(Next(()))
    }

    fn nip<V: Value>(&mut self, (_, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(b))
    }

    fn swp<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<(V, V)>, Trap> {
        Ok(Next((b, a)))
    }

    fn rot<V: Value>(&mut self, (a, b, c): (V, V, V)) -> Result<ControlFlow<(V, V, V)>, Trap> {
        Ok(Next((b, c, a)))
    }

    fn dup<V: Value>(&mut self, a: V) -> Result<ControlFlow<(V, V)>, Trap> {
        Ok(Next((a, a)))
    }

    fn ovr<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<(V, V, V)>, Trap> {
        Ok(Next((a, b, a)))
    }

    fn equ<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<bool>, Trap> {
        Ok(Next(a == b))
    }

    fn neq<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<bool>, Trap> {
        Ok(Next(a != b))
    }

    fn gth<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<bool>, Trap> {
        Ok(Next(a > b))
    }

    fn lth<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<bool>, Trap> {
        Ok(Next(a < b))
    }

    fn jmp<V: Value>(&mut self, val: V) -> Result<ControlFlow<()>, Trap> {
        self.uxn.pc.jmp(val);
        Ok(Next(()))
    }

    fn jcn<V: Value>(&mut self, (cond, val): (bool, V)) -> Result<ControlFlow<()>, Trap> {
        if cond {
            self.uxn.pc.jmp(val);
        }
        Ok(Next(()))
    }

    fn jsr<V: Value>(&mut self, val: V) -> Result<ControlFlow<Short>, Trap> {
        let ret = Short::from(self.uxn.pc);
        self.uxn.pc.jmp(val);
        Ok(Next(ret))
    }

    fn sth<V: Value>(&mut self, a: V) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a))
    }

    fn ldz<V: Value>(&mut self, addr: Byte) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(self.uxn.ram.peek(addr.into())))
    }

    fn stz<V: Value>(&mut self, (val, addr): (V, Byte)) -> Result<ControlFlow<()>, Trap> {
        self.uxn.ram.poke(addr.into(), val);
        Ok(Next(()))
    }

    fn ldr<V: Value>(&mut self, offs: Byte) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(
            self.uxn.ram.peek(Short::from(self.uxn.pc).offset(offs)),
        ))
    }

    fn str<V: Value>(&mut self, (val, offs): (V, Byte)) -> Result<ControlFlow<()>, Trap> {
        self.uxn
            .ram
            .poke(Short::from(self.uxn.pc).offset(offs), val);
        Ok(Next(()))
    }

    fn lda<V: Value>(&mut self, addr: Short) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(self.uxn.ram.peek(addr)))
    }

    fn sta<V: Value>(&mut self, (val, addr): (V, Short)) -> Result<ControlFlow<()>, Trap> {
        self.uxn.ram.poke(addr, val);
        Ok(Next(()))
    }

    fn dei<V: Value>(&mut self, addr: Byte) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(self.dev.dei(self.uxn, addr)))
    }

    fn deo<V: Value>(&mut self, (val, addr): (V, Byte)) -> Result<ControlFlow<()>, Trap> {
        self.dev.deo(self.uxn, addr, val);
        Ok(Next(()))
    }

    fn add<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a + b))
    }

    fn sub<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a - b))
    }

    fn mul<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a * b))
    }

    fn div<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        if b.is_zero() {
            return Err(Trap::DivisionByZero);
        }
        Ok(Next(a / b))
    }

    fn and<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a & b))
    }

    fn ora<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a | b))
    }

    fn eor<V: Value>(&mut self, (a, b): (V, V)) -> Result<ControlFlow<V>, Trap> {
        Ok(Next(a ^ b))
    }

    fn sft<V: Value>(&mut self, (a, b): (V, Byte)) -> Result<ControlFlow<V>, Trap> {
        let b = u8::from(b);
        let res = a >> (b & 0x0f) << ((b & 0xf0) >> 4);
        Ok(Next(res))
    }
}