struct VM<'p> { 
    program: &'p Program, 
    block: u32, ip: u32, 
    heap: Heap<'p>,
    locals: Object,
    stack: Vec<StackItem>, 
    status: Status 
}

#[derive(Clone, Copy)]
enum Status {
    Running,
    AwaitSyscall(u32),
    Crash(&'static str),
    Exit,
}

#[derive(Clone, Copy, Debug)]
enum StackItem {
    Object(Object),
    Integer(u32),
}
impl StackItem {
    fn object(self) -> Option<Object> {
        if let StackItem::Object(t) = self { return Some(t) }
        return None
    }
    
    fn integer(self) -> Option<u32> {
        if let StackItem::Integer(u) = self { return Some(u) }
        return None
    }
}

#[derive(Clone, Copy, Debug)]
struct Object { index: u32 }  
struct Program { 
    blocks: Vec<BlockData>,
    types: Vec<TypeData>,
}
struct BlockData { instructions: Vec<Instruction> }

struct Block { index: u32 }

#[derive(Debug)]
enum Instruction {
    // push NULL to stack
    PushNull, 
    // push an integer to the stack
    PushInteger(u32),

    // push a new object to the stack with the given type
    PushNew(Type),

    // get the locals register
    PushLocals,
    // pop the locals register
    PopLocals,

    Op11(Op11), // [t1] -> [t2]
    Op21(Op21), // [t1, t2] -> [t3]

    // assert that the stack has n items and jump to the given block
    Call { block: u32, n: u32, },  
    // assert that the stack has n items, clear it, and execute a syscall
    Syscall { syscall: u32, n: u32 },

    // kill the program
    Crash(&'static str),
    // end the program, popping the top stack item
    Exit
}


#[derive(Debug)]
enum Op11 { Get(u32) }

#[derive(Debug)]
enum Op21 { Add, Set(u32) }

const TY_NULL: Type = Type { id: 0x0 };
const NULL: Object = Object { index: 0 };

const SYSCALL_DBG: u32 = 0x1000;


fn main() {
    let mut program = Program::new();
    let ty_locals = program.type_create(3);
    program.block_create(
        BlockData { 
            instructions: vec![ 
                Instruction::PushNew(ty_locals), Instruction::PopLocals, 
                
                Instruction::PushLocals, 
                Instruction::PushInteger(0xdead0000), Instruction::Op21(Op21::Set(0)),
                Instruction::PushInteger(0x0000ba75), Instruction::Op21(Op21::Set(1)),
                // stack: [locals]

                Instruction::Call { block: 1, n: 1 },
            ]
        }
    );

    program.block_create(
        BlockData {
            instructions: vec![
                Instruction::PopLocals,

                Instruction::PushLocals,
                    Instruction::PushLocals, Instruction::Op11(Op11::Get(0)),
                    Instruction::PushLocals, Instruction::Op11(Op11::Get(1)),
                Instruction::Op21(Op21::Add),
                Instruction::Op21(Op21::Set(2)),

                Instruction::Call { block: 2, n: 1 }
            ]
        }
    );

    program.block_create(
        BlockData { 
            instructions: vec![
                Instruction::PopLocals,

                Instruction::PushLocals,
                Instruction::Op11(Op11::Get(0)),
                Instruction::Syscall { syscall: SYSCALL_DBG, n: 1 },

                Instruction::Exit,
            ]
        }
    );

    host_program(&program).unwrap()
}

fn host_program(p: &Program) -> Result<(), String> {
    let mut vm = VM::start(p);
    loop {
        let status = vm.step();
        match status {
            Status::Running => {} 
            Status::AwaitSyscall(ix) => {
                match ix {
                    SYSCALL_DBG => { 
                        let top = vm.stack.pop().ok_or("should be an item on the stack")?;
                        println!("debug print: {:08x?}", top);
                        vm.complete_syscall()
                    }
                    _ => { return Err(format!("unknown syscall: {}", ix)) }
                }
            }
            Status::Crash(err) => { return Err(err.to_string()); }
            Status::Exit => { return Ok(()); }
        }
    }
}


impl<'p> VM<'p> {
    pub fn start(program: &'p Program) -> VM<'p> {
        let heap = Heap::new(program);

        VM {
            program,
            block: 0,
            ip: 0,
            heap,
            locals: NULL,
            stack: vec![],
            status: Status::Running,
        }
    }

    pub fn step(&mut self) -> Status {
        let error = match self.internal_step() {
            Ok(_) => { return self.status; }
            Err(e) => e
        };
        self.status = Status::Crash(error);
        self.status
    }
    fn internal_step(&mut self) -> Result<(), &'static str> {
        if let Status::Running = self.status { } else { return Ok(()); }

        let instruction = 
            self.program
                .blocks.get(self.block as usize).ok_or("block must exist")?
                .instructions.get(self.ip as usize).ok_or("instruction must exist")?;
        self.ip += 1;
        println!("instruction: {:?}", instruction);

        match instruction {
            &Instruction::PushNull => {
                self.stack.push(StackItem::Object(NULL))
            }

            &Instruction::PushInteger(int) => {
                self.stack.push(StackItem::Integer(int))
            }

            &Instruction::PushNew(ty) => {
                let t = self.heap.object_create(ty);
                self.stack.push(StackItem::Object(t))
            }

            Instruction::PushLocals => {
                self.stack.push(StackItem::Object(self.locals))
            }
            
            Instruction::PopLocals => {
                let t0 = self.stack.pop().ok_or("t0 must be present")?.object().ok_or("t0 must be a object")?;
                self.locals = t0;
            }

            Instruction::Op11(op) => {
                let t0 = self.stack.pop().ok_or("t0 must be present")?;
                self.stack.push(op.perform(&mut self.heap, t0)?);
            }

            Instruction::Op21(op) => {
                let t1 = self.stack.pop().ok_or("t1 must be present")?;
                let t0 = self.stack.pop().ok_or("t0 must be present")?;
                self.stack.push(op.perform(&mut self.heap, t0, t1)?);
            }

            &Instruction::Call { block, n } => {
                if self.stack.len() != n as usize {
                    return Err("stack has wrong number of elements");
                }
                self.locals = NULL;
                self.block = block;
                self.ip = 0;

            }
            &Instruction::Syscall { syscall, n } => {
                if self.stack.len() != n as usize {
                    return Err("stack has wrong number of elements");
                }
                self.status = Status::AwaitSyscall(syscall)
            }
            Instruction::Crash(c) => {
                panic!("crashed: {}", c)
            }
            Instruction::Exit => {
                self.status = Status::Exit
            }
        }

        Ok(())
    }

    pub fn complete_syscall(&mut self) {
        match &self.status {
            Status::AwaitSyscall(_) => {
                self.status = Status::Running;
            }, // continue
            _ => { 
                self.status = Status::Crash("tried to complete a syscall while not awaiting one")
            }
        }
    }
}

impl Op11 {
    fn perform(&self, heap: &mut Heap, v0: StackItem) -> Result<StackItem, &'static str> {
        match self {
            &Op11::Get(ix) => {
                let t0 = v0.object().ok_or("must be a object")?;
                Ok(heap.object_get_arg(t0, ix))
            }
       }
    }

}

impl Op21 {
    fn perform(&self, heap: &mut Heap, v0: StackItem, v1: StackItem) -> Result<StackItem, &'static str> {
        match self {
            Op21::Add => {
                let i0 = v0.integer().ok_or("must be an integer")?;
                let i1 = v1.integer().ok_or("must be an integer")?;
                Ok(StackItem::Integer(i0.wrapping_add(i1)))
            }
            &Op21::Set(ix) => {
                let t0 = v0.object().ok_or("must be a object")?;
                heap.object_set_arg(t0, ix, v1);
                Ok(v0)
            }
        }
    }
}


impl Program {
    fn new() -> Program {
        let mut p = Self {
            blocks: vec![],
            types: vec![],
        };
        let ty_null = p.type_create(0);
        assert_eq!(TY_NULL, ty_null);
        return p
    }

    fn type_create(&mut self, len: u32) -> Type {
        let index = self.types.len();
        self.types.push(TypeData { len });
        Type { id: index as u32 }
    }

    fn type_get(&self, ty: Type) -> &TypeData {
        return self.types.get(ty.id as usize).expect("type should have existed")
    }

    fn block_create(&mut self, block: BlockData) -> Block {
        let index = self.blocks.len();
        self.blocks.push(block);
        Block { index: index as u32 }
    }
}

struct Heap<'p> { 
    program: &'p Program,
    data: Vec<HeapItem> 
}

#[derive(Debug)]
enum HeapItem {
    Type(Type),
    StackItem(StackItem),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Type {
    id: u32
}

struct TypeData {
    len: u32
}

impl<'p> Heap<'p> {
    fn new(p: &'p Program) -> Self {
        let mut h = Heap {
            program: p,
            data: vec![],
        };
        let null = h.object_create(TY_NULL);
        assert_eq!(NULL.index, null.index);
        h
    }

    fn object_create(&mut self, ty: Type) -> Object {
        let td = self.program.type_get(ty);
        let index = self.data.len() as u32;
        self.data.push(HeapItem::Type(ty));
        for _ in 0..td.len {
            self.data.push(HeapItem::StackItem(StackItem::Object(NULL)));
        }
        Object { index }
    }

    fn object_get_type(&self, object: Object) -> Type {
        let header = self.data.get(object.index as usize).expect("object should have been valid");
        if let &HeapItem::Type(ty) = header {
            ty
        } else {
            panic!("header is invalid: {:?}", header)
        }
    }

    fn object_get_arg(&self, object: Object, ix: u32) -> StackItem {
        let ty = self.object_get_type(object);
        let td = self.program.type_get(ty);

        if ix < td.len {
            let value = self.data.get((object.index + ix + 1) as usize).expect("should have been present");
            match value {
                &HeapItem::StackItem(si) => si,
                _ => panic!("should have been a StackItem")
            }
        } else {
            panic!("invalid ix: {}", ix)
        }
    }

    fn object_set_arg(&mut self, object: Object, ix: u32, arg: StackItem) {
        let ty = self.object_get_type(object);
        let td = self.program.type_get(ty);

        if ix < td.len {
            let slot = self.data.get_mut((object.index + ix + 1) as usize).expect("should have been present");
            *slot = HeapItem::StackItem(arg)
        } else {
            panic!("invalid ix: {}", ix)
        }
    }
}