mod encoding;
mod heap;
mod value;

use encoding::{Arg, Op};
use heap::Heap;
use std::ffi::OsString;
use std::io::{stderr, Write};
use value::Value;

struct State {
    i: usize,
    heap: Heap,
}

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

    fn op(&mut self, op: Op) -> Result<(), String> {
        match op {
            Op::Alloc(op) => {
                let Some(i) = self.read_arg(op.size).to_int() else {
                    return Err(String::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(String::from("call needs an int"));
                };
                let Ok(u) = usize::try_from(i) else {
                    return Err(String::from("call needs a positive int"));
                };
                self.i = u;
            }
            Op::Poke(op) => {
                let Some(p) = self.heap.locals[usize::from(op.ptr)].to_pointer() else {
                    return Err(String::from("poke needs a pointer"));
                };
                self.heap.poke(p + op.offset, self.read_arg(op.val))?;
            }
            Op::Peek(op) => {
                let Some(p) = self.read_arg(op.val).to_pointer() else {
                    return Err(String::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.read_arg(op.val);
            }
            Op::Exit(op) => {
                let Some(i) = self.read_arg(op.val).to_int() else {
                    std::process::exit(255);
                };
                std::process::exit(i as i32);
            }
            Op::Add(op) => {
                let Some(v1) = self.read_arg(op.val1).to_int() else {
                    return Err(String::from("add needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::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.read_arg(op.val1).to_int() else {
                    return Err(String::from("sub needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::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.read_arg(op.val1).to_int() else {
                    return Err(String::from("mul needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::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.read_arg(op.val1).to_int() else {
                    return Err(String::from("div needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::from("div needs an int"));
                };
                let Some(res) = v1.checked_div(v2) else {
                    return Err(String::from("division by zero"));
                };
                self.heap.locals[usize::from(op.out)] = Value::from_int(res);
            }
            Op::Less(op) => {
                let Some(v1) = self.read_arg(op.val1).to_int() else {
                    return Err(String::from("less needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::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.read_arg(op.val1).to_int() else {
                    return Err(String::from("eq needs an int"));
                };
                let Some(v2) = self.read_arg(op.val2).to_int() else {
                    return Err(String::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.read_arg(op.test).to_int() else {
                    return Err(String::from("if needs an int"));
                };
                if t != 0 {
                    self.i = op.target;
                }
            }
        }
        Ok(())
    }
}

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

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