mod encoding;
mod heap;
mod value;

use encoding::{Arg, Op};
use heap::Heap;
use std::error::Error;
use std::ffi::OsString;
use std::io::{stdin, BufRead, BufWriter, Write};
use value::Value;

trait ReadArg {
    fn read_arg(&self, a: Arg) -> Value;
}

impl ReadArg for Heap {
    fn read_arg(&self, a: Arg) -> Value {
        match a {
            Arg::Local(l) => self.locals[usize::from(l)],
            Arg::Const(c) => c,
        }
    }
}

fn write(
    w: &mut impl Write,
    heap: &Heap,
    ptr: u8,
    off: Arg,
    len: Arg,
) -> Result<(), Box<dyn Error>> {
    let Some(ptr) = heap.locals[usize::from(ptr)].to_pointer() else {
        return Err(Box::from("write needs a pointer"));
    };
    let Some(ioff) = heap.read_arg(off).to_int() else {
        return Err(Box::from("write needs an int"));
    };
    let Ok(off) = usize::try_from(ioff) else {
        return Err(Box::from("write: off is negative"));
    };
    let Some(ilen) = heap.read_arg(len).to_int() else {
        return Err(Box::from("write needs an int"));
    };
    let Ok(len) = usize::try_from(ilen) else {
        return Err(Box::from("write: len is negative"));
    };
    for i in 0..len {
        let Some(ibyte) = heap.peek(ptr + off + i)?.to_int() else {
            return Err(Box::from("write: buffer has non-integer entries"));
        };
        let Ok(byte) = u8::try_from(ibyte) else {
            return Err(Box::from("write: buffer entry is not between 0 and 255"));
        };
        w.write_all(&[byte])?;
    }
    w.flush()?;
    Ok(())
}

struct State {
    i: usize,
    heap: Heap,
    stdout: BufWriter<Box<dyn Write>>,
    stderr: BufWriter<Box<dyn Write>>,
}

impl State {
    fn op(&mut self, op: Op) -> Result<(), Box<dyn Error>> {
        match op {
            Op::Alloc(op) => {
                let Some(i) = self.heap.read_arg(op.size).to_int() else {
                    return Err(Box::from("alloc needs an int"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_pointer(self.heap.alloc(i)?);
            }
            Op::Call => {
                let Some(i) = self.heap.locals[0].to_int() else {
                    return Err(Box::from("call needs an int"));
                };
                self.i = usize::try_from(i)?;
            }
            Op::Poke(op) => {
                let Some(p) = self.heap.locals[usize::from(op.ptr)].to_pointer() else {
                    return Err(Box::from("poke needs a pointer"));
                };
                self.heap.poke(p, op.offset, self.heap.read_arg(op.val))?;
            }
            Op::Peek(op) => {
                let Some(p) = self.heap.read_arg(op.val).to_pointer() else {
                    return Err(Box::from("peek needs a pointer"));
                };
                self.heap.locals[usize::from(op.out)] = self.heap.peek(p + op.offset)?;
            }
            Op::Shuf(op) => {
                self.heap.locals[usize::from(op.out)] = self.heap.read_arg(op.val);
            }
            Op::Exit(op) => {
                let Some(i) = self.heap.read_arg(op.val).to_int() else {
                    std::process::exit(255);
                };
                std::process::exit(i as i32);
            }
            Op::Add(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("add needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("add needs an int"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_int(v1.wrapping_add(v2));
            }
            Op::Sub(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("sub needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("sub needs an int"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_int(v1.wrapping_sub(v2));
            }
            Op::Mul(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("mul needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("mul needs an int"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_int(v1.wrapping_mul(v2));
            }
            Op::Div(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("div needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("div needs an int"));
                };
                let Some(res) = v1.checked_div(v2) else {
                    return Err(Box::from("division by zero"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_int(res);
            }
            Op::Less(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("less needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("less needs an int"));
                };
                let res = if v1 < v2 { 1 } else { 0 };
                self.heap.locals[usize::from(op.out)] = Value::from_int(res);
            }
            Op::Eq(op) => {
                let Some(v1) = self.heap.read_arg(op.val1).to_int() else {
                    return Err(Box::from("eq needs an int"));
                };
                let Some(v2) = self.heap.read_arg(op.val2).to_int() else {
                    return Err(Box::from("eq needs an int"));
                };
                let res = if v1 == v2 { 1 } else { 0 };
                self.heap.locals[usize::from(op.out)] = Value::from_int(res);
            }
            Op::If(op) => {
                let Some(t) = self.heap.read_arg(op.test).to_int() else {
                    return Err(Box::from("if needs an int"));
                };
                if t != 0 {
                    self.i = op.target;
                }
            }
            Op::Read(op) => {
                let Some(ptr) = self.heap.locals[usize::from(op.ptr)].to_pointer() else {
                    return Err(Box::from("read needs a pointer"));
                };
                let Some(ioff) = self.heap.read_arg(op.off).to_int() else {
                    return Err(Box::from("read needs an int"));
                };
                let Ok(off) = usize::try_from(ioff) else {
                    return Err(Box::from("read: off is negative"));
                };
                let Some(ilen) = self.heap.read_arg(op.len).to_int() else {
                    return Err(Box::from("read needs an int"));
                };
                let Ok(len) = usize::try_from(ilen) else {
                    return Err(Box::from("read: len is negative"));
                };
                let mut stdin = stdin().lock();
                let buf = stdin.fill_buf()?;
                let n = std::cmp::min(len, buf.len());
                for (i, &byte) in buf[..n].iter().enumerate() {
                    self.heap
                        .poke(ptr, off + i, Value::from_int(i64::from(byte)))?;
                }
                stdin.consume(n);
                self.heap.locals[usize::from(op.out)] = Value::from_int(
                    i64::try_from(n)
                        .expect("we can never read more than len bytes, and len fit in an i64"),
                );
            }
            Op::Write(op) => {
                write(&mut self.stdout, &self.heap, op.ptr, op.off, op.len)?;
            }
            Op::WriteErr(op) => {
                write(&mut self.stderr, &self.heap, op.ptr, op.off, op.len)?;
            }
        }
        Ok(())
    }
}

fn run() -> Result<(), Box<dyn Error>> {
    let args: Vec<OsString> = std::env::args_os().collect();
    if args.len() != 2 {
        return Err(Box::from("usage: bytecode <file>"));
    }
    let prog = std::fs::read(&args[1])?;
    let mut st = State {
        i: 0,
        heap: Heap::new(),
        stdout: BufWriter::new(Box::new(std::io::stdout().lock())),
        stderr: BufWriter::new(Box::new(std::io::stderr().lock())),
    };
    loop {
        let (n, op) = Op::parse(&prog[st.i..])?;
        st.i += n;
        st.op(op)?;
    }
}

fn main() {
    if let Err(err) = run() {
        let _ = writeln!(std::io::stderr(), "FAIL: {}", err);
        std::process::exit(255);
    }
}