use std::collections::HashMap;

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

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

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

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

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
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: Block, 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 BL_MAIN: Block = Block { index: 0x0 };
const TY_NULL: Type = Type { id: 0x0 };
const NULL: Object = Object { index: 0 };

const SYSCALL_DBG: u32 = 0x1000;


fn main() {
    _main().unwrap()
}
fn _main() -> Result<(), String> {
    let mut program = Program::new();
    let ty_locals = program.type_prototype("main.locals");
    program.type_instantiate(ty_locals, 3)?;

    let bl_main_0 = program.block_prototype("main");
    let bl_main_1 = program.block_prototype("main.1");
    let bl_main_2 = program.block_prototype("main.2");

    program.block_instantiate(
        bl_main_0,
        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: bl_main_1, n: 1 },
            ]
        }
    )?;

    program.block_instantiate(
        bl_main_1,
        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: bl_main_2, n: 1 }
            ]
        }
    )?;

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

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

                Instruction::Exit,
            ]
        }
    )?;

    host_program(&program)
}

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: BL_MAIN,
            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
                .block_get(self.block)?
                .instructions.get(self.ip as usize).ok_or("instruction must exist")?;
        self.ip += 1;
        println!("instruction: {:?}", instruction);

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

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

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

            Instruction::PushLocals => {
                self.stack.push(Value::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: Value) -> Result<Value, &'static str> {
        match self {
            &Op11::Get(ix) => {
                let t0 = v0.object().ok_or("must be a object")?;
                heap.object_get_arg(t0, ix)
            }
       }
    }

}

impl Op21 {
    fn perform(&self, heap: &mut Heap, v0: Value, v1: Value) -> Result<Value, &'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(Value::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: Prototyper::<BlockData>::new(),
            types: Prototyper::<TypeData>::new(),
        };

        let ty_null = p.type_prototype("null");
        assert_eq!(TY_NULL, ty_null);
        p.type_instantiate(ty_null, 0).expect("cannot fail to create null instance");

        let bl_main = p.block_prototype("main");
        assert_eq!(BL_MAIN, bl_main);


        return p
    }

    fn type_prototype(&mut self, name: &str) -> Type {
        Type { id: self.types.prototype(name) }
    }

    fn type_instantiate(&mut self, ty: Type, len: u32) -> Result<(), &'static str> {
        self.types.instantiate(ty.id, TypeData { len })
    }

    fn type_get(&self, ty: Type) -> Result<&TypeData, &'static str> {
        self.types.get(ty.id)
    }

    fn block_prototype(&mut self, name: &str) -> Block {
        Block { index: self.blocks.prototype(name) }
    }

    fn block_instantiate(&mut self, bl: Block, data: BlockData) -> Result<(), &'static str> {
        self.blocks.instantiate(bl.index, data)
    }

    fn block_get(&self, bl: Block) -> Result<&BlockData, &'static str> {
        self.blocks.get(bl.index)
    }
}

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

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

#[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).expect("shouldn't fail during Heap construction");
        assert_eq!(NULL.index, null.index);
        h
    }

    fn object_create(&mut self, ty: Type) -> Result<Object, &'static str> {
        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::Value(Value::Object(NULL)));
        }
        Ok(Object { index })
    }

    fn object_get_type(&self, object: Object) -> Result<Type, &'static str> {
        let header = self.data.get(object.index as usize).ok_or("object should have been valid")?;
        if let &HeapItem::Type(ty) = header {
            Ok(ty)
        } else {
            Err("header is invalid")
        }
    }

    fn object_get_arg(&self, object: Object, ix: u32) -> Result<Value, &'static str> {
        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::Value(si) => Ok(si),
                _ => return Err("should have been a Value"),
            }
        } else {
            return Err("invalid field index")
        }
    }

    fn object_set_arg(&mut self, object: Object, ix: u32, arg: Value) -> Result<(), &'static str> {
        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::Value(arg);
            Ok(())
        } else {
            return Err("invalid field index")
        }
    }
}

struct Prototyper<T> {
    by_name: HashMap<String, u32>,
    names: Vec<String>,
    values: Vec<Option<T>>
}

impl<T> Prototyper<T> {
    pub fn new() -> Self {
        Self {
            by_name: HashMap::new(),
            names: vec![],
            values: vec![],
        }
    }

    pub fn prototype(&mut self, name: &str) -> u32 {
        if let Some(v) = self.by_name.get(name) { return *v; }

        let ix = self.values.len() as u32;
        self.by_name.insert(name.to_string(), ix);
        self.names.push(name.to_string());
        self.values.push(None);
        ix
    }

    pub fn instantiate(&mut self, handle: u32, value: T) -> Result<(), &'static str> {
        let slot = self.values
            .get_mut(handle as usize)
            .ok_or("handle was not prototyped on this prototyper")?;
        if slot.is_some() {
            return Err("can't instantiate handle twice")
        }
        *slot = Some(value);
        Ok(())
    }

    pub fn get(&self, handle: u32) -> Result<&T, &'static str> {
        match self.values.get(handle as usize) {
            Some(Some(h)) => Ok(h),
            Some(_) => Err("handle was prototyped but not instantiated"),
            _ => Err("handle was not prototyped on this prototyper"),
        }
    }

    pub fn get_mut(&mut self, handle: u32) -> Result<&mut T, &'static str> {
        match self.values.get_mut(handle as usize) {
            Some(Some(h)) => Ok(h),
            Some(_) => Err("handle was prototyped but not instantiated"),
            _ => Err("handle was not prototyped on this prototyper"),
        }
    }
}