use crate::heap::{Heap, Pointer};
use std::io::{BufWriter, Read, StdinLock, Write};

// The cute scheme virtual machine has two stacks:
//   - the data stack and
//   - the locals stack.
// There are four pointers:
//   - the instruction pointer,
//   - the argument pointer,
//   - and the (data) stack pointer.
// These pointers cannot be manipulated directly, but are referred to in
// the comments below on the opcodes.

#[derive(Debug, Eq, PartialEq)]
pub enum Op {
    // Signed arithmetic
    // =================

    // ( -- n )
    // Pushes a constant.
    Const(i64),
    // ( n1 n2 -- n3 )
    // Adds two integers.
    Add,
    // ( n1 n2 -- n3 )
    // Subtracts two integers.
    Sub,
    // ( n1 n2 -- n3 )
    // Multiplies two integers.
    Mul,
    // ( n1 n2 -- n3 )
    // Divides two integers and truncates the result.
    Div,
    // ( n1 n2 -- n3 )
    // Remainder from Div.
    Mod,

    // Heap
    // ======

    // ( n -- a )
    // Allocates n bytes and returns the address.
    Alloc,
    // ( a -- u )
    // Fetches a 64 bit word from the specified address plus the
    // given offset.
    Peek(u8),
    // ( u a )
    // Stores a 64 bit word at the specified address plus the
    // given offset.
    Poke(u8),
    // ( a -- n )
    // Fetches a byte from the specified address.
    PeekByte,
    // ( n a )
    // Stores a byte at the specified address.
    PokeByte,

    // Stack
    // =====

    // ( n -- )
    // Deletes an element from the stack.
    Pop,
    // ( -- n )
    // Copies an argument from the locals stack to the data stack. The argument is an index above the argument pointer.
    Local(u8),

    // Control flow
    // ============

    // ( f -- )
    // Jumps to the specified offset if the argument is non-zero.
    If(i64),
    // ( i-addr a1 a2 ... an  )
    // ( Locals stack: -- instruction-ptr a1 a2 ... an argument-ptr )
    // Pushes the instruction pointer, moves n arguments to the locals
    // stack, pushes the old argument pointer, and jumps to the
    // specified location.
    Call(u8),
    // ( Locals stack: instruction-ptr a1 a2 ... an argument-ptr )
    // Pops the argument pointer (effectively popping n more arguments which were added by Call),
    // and then pops the instruction pointer.
    Ret,
    // ( n )
    // Terminates the interpreter with the given status code.
    Exit,

    // Input and output
    // ================

    // ( b file -- f )
    // Writes a byte to a file. Returns true on success, false on error.
    PutC,
    // ( file -- b )
    // Reads one byte from a file. Returns -1 on EOF, 0 on error.
    GetC,
}

struct Stack {
    v: Vec<u64>,
}

impl Stack {
    fn pop(&mut self) -> Result<u64, String> {
        self.v.pop().ok_or(String::from("stack underflow"))
    }

    fn pop_int(&mut self) -> Result<i64, String> {
        Ok(self.pop()? as i64 >> 1)
    }

    fn push_int(&mut self, i: i64) {
        self.v.push((i as u64) << 1 | 1);
    }

    fn pop_pointer(&mut self) -> Result<Pointer, String> {
        Ok(Pointer::from_bytes(self.pop()?))
    }

    fn push_pointer(&mut self, p: Pointer) {
        self.v.push(p.bytes());
    }

    fn pop_usize(&mut self) -> Result<usize, String> {
        Ok(usize::try_from(self.pop()?).unwrap() >> 1)
    }

    fn push_usize(&mut self, p: usize) {
        self.v.push(u64::try_from(p).unwrap() << 1 | 1);
    }
}

fn fn_const(stack: &mut Stack, n: i64) {
    stack.push_int(n);
}

fn add(stack: &mut Stack) -> Result<(), String> {
    let n2 = stack.pop_int()?;
    let n1 = stack.pop_int()?;
    stack.push_int(n1.wrapping_add(n2));
    Ok(())
}

fn sub(stack: &mut Stack) -> Result<(), String> {
    let n2 = stack.pop_int()?;
    let n1 = stack.pop_int()?;
    stack.push_int(n1.wrapping_sub(n2));
    Ok(())
}

fn mul(stack: &mut Stack) -> Result<(), String> {
    let n2 = stack.pop_int()?;
    let n1 = stack.pop_int()?;
    stack.push_int(n1.wrapping_mul(n2));
    Ok(())
}

fn div(stack: &mut Stack) -> Result<(), String> {
    let n2 = stack.pop_int()?;
    let n1 = stack.pop_int()?;
    stack.push_int(n1.wrapping_div(n2));
    Ok(())
}

fn fn_mod(stack: &mut Stack) -> Result<(), String> {
    let n2 = stack.pop_int()?;
    let n1 = stack.pop_int()?;
    stack.push_int(n1.wrapping_rem(n2));
    Ok(())
}

fn alloc(stack: &mut Stack, heap: &mut Heap) -> Result<(), String> {
    let n = stack.pop_int()?;
    let n_usize = match n.try_into() {
        Ok(x) => x,
        Err(_) => {
            return Err(String::from("invalid size"));
        }
    };
    stack.push_pointer(heap.alloc(n_usize));
    Ok(())
}

fn peek(stack: &mut Stack, heap: &Heap, n: u8) -> Result<(), String> {
    let a = stack.pop_pointer()?;
    stack.v.push(heap.peek(a.offset(n))?);
    Ok(())
}

fn poke(stack: &mut Stack, heap: &mut Heap, n: u8) -> Result<(), String> {
    let a = stack.pop_pointer()?;
    let u = stack.pop()?;
    heap.poke(u, a.offset(n))?;
    Ok(())
}

fn peek_byte(stack: &mut Stack, heap: &Heap) -> Result<(), String> {
    let a = stack.pop_pointer()?;
    stack.push_int(i64::from(heap.peek_byte(a)));
    Ok(())
}

fn poke_byte(stack: &mut Stack, heap: &mut Heap) -> Result<(), String> {
    let a = stack.pop_pointer()?;
    let n = stack.pop_int()?;
    heap.poke_byte((n & 0xff).try_into().unwrap(), a);
    Ok(())
}

fn pop(stack: &mut Stack) -> Result<(), String> {
    stack.pop()?;
    Ok(())
}

fn local(stack: &mut Stack, locals: &Stack, n: u8) -> Result<(), String> {
    let i = locals
        .v
        .len()
        .checked_sub(usize::from(n) + 1)
        .ok_or("out of bounds")?;
    stack.v.push(locals.v[i]);
    Ok(())
}

fn fn_if(stack: &mut Stack, ip: &mut usize, n: i64) -> Result<(), String> {
    let f = stack.pop_int()?;
    if f != 0 {
        if n > 0 {
            *ip = ip
                .checked_add(n.try_into().unwrap())
                .ok_or("invalid offset")?;
        } else {
            *ip = ip
                .checked_sub((-n).try_into().unwrap())
                .ok_or("invalid offset")?;
        }
    }
    Ok(())
}

fn call(stack: &mut Stack, locals: &mut Stack, ip: &mut usize, n: u8) -> Result<(), String> {
    // Push the instruction pointer.
    locals.push_usize(*ip);
    let old_ap = locals.v.len();
    // Move n arguments to the locals stack.
    let a1_idx = stack.v.len() - usize::from(n);
    for i in a1_idx..stack.v.len() {
        locals.v.push(stack.v[i]);
    }
    stack.v.truncate(a1_idx);
    // Push the old argument pointer.
    locals.push_usize(old_ap);
    // Jump to the specified location.
    let i_addr = stack.pop_int()?;
    let i_addr_usize = match usize::try_from(i_addr) {
        Ok(x) => x,
        Err(_) => {
            return Err(String::from("invalid address"));
        }
    };
    // Subtract 1 because the interpreter will also increment the
    // instruction pointer.
    *ip = i_addr_usize - 1;
    Ok(())
}

fn ret(locals: &mut Stack, ip: &mut usize) -> Result<(), String> {
    let ap = locals.pop_usize()?;
    locals.v.truncate(ap);
    let return_address = locals.pop_usize()?;
    *ip = return_address;
    Ok(())
}

trait File: Write + Read {}

struct FileTable<'a> {
    files: Vec<Box<dyn File + 'a>>,
}

impl<'a> FileTable<'a> {
    fn file(&mut self, f: i64) -> Result<&mut (dyn File + 'a), String> {
        if !(0 <= f && usize::try_from(f).unwrap() < self.files.len()) {
            return Err(String::from("invalid file"));
        }
        Ok(&mut *self.files[usize::try_from(f).unwrap()])
    }
}

struct Stdin<'a>(StdinLock<'a>);

impl<'a> Write for Stdin<'a> {
    fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
        Err(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "can't write to stdin",
        ))
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

impl<'a> Read for Stdin<'a> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let Stdin(ref mut handle) = self;
        handle.read(buf)
    }
}

impl<'a> File for Stdin<'a> {}

struct Out<T>(T);

impl<T: Write> Write for Out<T> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let Out(ref mut handle) = self;
        handle.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        let Out(ref mut handle) = self;
        handle.flush()
    }
}

impl<T> Read for Out<T> {
    fn read(&mut self, _: &mut [u8]) -> std::io::Result<usize> {
        Err(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "invalid file for reading",
        ))
    }
}

impl<T: Write> File for Out<T> {}

fn putc(stack: &mut Stack, files: &mut FileTable) -> Result<(), String> {
    let file = stack.pop_int()?;
    let byte = stack.pop_int()?;
    match files
        .file(file)?
        .write(&vec![(byte & 0xff).try_into().unwrap()])
    {
        Ok(_) => {
            stack.push_int(-1);
        }
        Err(_) => {
            stack.push_int(0);
        }
    }
    Ok(())
}

fn getc(stack: &mut Stack, files: &mut FileTable) -> Result<(), String> {
    let file = stack.pop_int()?;
    let mut buf = vec![0];
    match files.file(file)?.read(&mut buf) {
        Ok(1) => {
            stack.push_int(buf[0].into());
        }
        Ok(0) => {
            stack.push_int(-1);
        }
        _ => {
            stack.push_int(0);
        }
    }
    Ok(())
}

pub fn eval(prog: &[Op]) -> Result<u8, String> {
    let mut stack = Stack { v: Vec::new() };
    let mut locals_stack = Stack { v: Vec::new() };
    let mut heap = Heap::new();
    let mut ip = 0;
    let stdin = std::io::stdin();
    let mut files = FileTable {
        files: vec![
            Box::new(Stdin(stdin.lock())),
            Box::new(Out(BufWriter::new(std::io::stdout()))),
            Box::new(Out(BufWriter::new(std::io::stderr()))),
        ],
    };
    loop {
        if ip >= prog.len() {
            return Err(String::from("invalid instruction pointer"));
        }
        match prog[ip] {
            Op::Const(n) => fn_const(&mut stack, n),
            Op::Add => add(&mut stack)?,
            Op::Sub => sub(&mut stack)?,
            Op::Mul => mul(&mut stack)?,
            Op::Div => div(&mut stack)?,
            Op::Mod => fn_mod(&mut stack)?,
            Op::Alloc => alloc(&mut stack, &mut heap)?,
            Op::Peek(n) => peek(&mut stack, &heap, n)?,
            Op::Poke(n) => poke(&mut stack, &mut heap, n)?,
            Op::PeekByte => peek_byte(&mut stack, &heap)?,
            Op::PokeByte => poke_byte(&mut stack, &mut heap)?,
            Op::Pop => pop(&mut stack)?,
            Op::Local(n) => local(&mut stack, &locals_stack, n)?,
            Op::If(n) => fn_if(&mut stack, &mut ip, n)?,
            Op::Call(n) => call(&mut stack, &mut locals_stack, &mut ip, n)?,
            Op::Ret => ret(&mut locals_stack, &mut ip)?,
            Op::PutC => putc(&mut stack, &mut files)?,
            Op::GetC => getc(&mut stack, &mut files)?,
            Op::Exit => {
                let n = stack.pop_int()?;
                return Ok(u8::try_from(n & 0xff).unwrap());
            }
        }
        ip += 1;
    }
}

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

    #[test]
    fn eval_const() {
        assert_eq!(Ok(5), eval(&vec![Const(5), Exit]));
    }

    #[test]
    fn eval_add() {
        assert_eq!(Ok(10), eval(&vec![Const(5), Const(5), Add, Exit]));
    }

    #[test]
    fn eval_sub() {
        assert_eq!(Ok(2), eval(&vec![Const(5), Const(3), Sub, Exit]));
    }

    #[test]
    fn eval_mul() {
        assert_eq!(Ok(25), eval(&vec![Const(5), Const(5), Mul, Exit]));
    }

    #[test]
    fn eval_div() {
        assert_eq!(Ok(2), eval(&vec![Const(5), Const(2), Div, Exit]));
    }

    #[test]
    fn eval_mod() {
        assert_eq!(Ok(1), eval(&vec![Const(5), Const(2), Mod, Exit]));
    }

    #[test]
    fn eval_alloc() {
        assert_eq!(Ok(0), eval(&vec![Const(10), Alloc, Pop, Const(0), Exit]));
    }

    #[test]
    fn eval_peek() {
        assert_eq!(Ok(0), eval(&vec![Const(8), Alloc, Peek(0), Exit]));
    }

    #[test]
    fn eval_poke() {
        assert_eq!(
            Ok(0),
            eval(&vec![Const(5), Const(8), Alloc, Poke(0), Const(0), Exit])
        );
    }

    #[test]
    fn eval_peek_byte() {
        assert_eq!(Ok(0), eval(&vec![Const(1), Alloc, PeekByte, Exit]));
    }

    #[test]
    fn eval_poke_byte() {
        assert_eq!(
            Ok(0),
            eval(&vec![Const(5), Const(1), Alloc, PokeByte, Const(0), Exit])
        );
    }

    #[test]
    fn eval_pop() {
        assert_eq!(Ok(5), eval(&vec![Const(5), Const(10), Pop, Exit]));
    }

    #[test]
    fn eval_if_true() {
        assert_eq!(
            Ok(5),
            eval(&vec![Const(5), Const(1), If(2), Const(5), Add, Exit])
        );
    }

    #[test]
    fn eval_if_false() {
        assert_eq!(
            Ok(10),
            eval(&vec![Const(5), Const(0), If(2), Const(5), Add, Exit])
        );
    }

    #[test]
    fn eval_call() {
        assert_eq!(
            Ok(5),
            eval(&vec![Const(3), Const(5), Call(1), Local(1), Exit])
        );
    }

    #[test]
    fn eval_ret() {
        assert_eq!(
            Ok(5),
            eval(&vec![Const(4), Const(5), Call(1), Exit, Local(1), Ret])
        );
    }

    #[test]
    fn eval_putc() {
        assert_eq!(
            Ok(5),
            eval(&vec![
                Const(5),
                Const(88),
                Const(1),
                PutC,
                If(2),
                Const(5),
                Add,
                Exit
            ])
        );
    }

    #[test]
    fn eval_getc_err() {
        assert_eq!(Ok(0), eval(&vec![Const(1), GetC, Exit]));
    }
}