use crate::compiler::{Chunk, Op};
use crate::parser::Ast;
use crate::spans::Spanned;
use crate::value::Value;
use chumsky::prelude::SimpleSpan;
use num_traits::FromPrimitive;
use std::mem::swap;

#[derive(Debug, Clone, PartialEq)]
pub struct Panic {
    pub input: &'static str,
    pub src: &'static str,
    pub msg: String,
    pub span: SimpleSpan,
    pub trace: Vec<Trace>,
    pub extra: String,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Trace {
    pub callee: Spanned<Ast>,
    pub caller: Spanned<Ast>,
    pub function: Value,
    pub arguments: Value,
    pub input: &'static str,
    pub src: &'static str,
}

pub struct Vm<'a> {
    pub stack: Vec<Value>,
    pub chunk: &'a Chunk<'a>,
    pub ip: usize,
    pub return_register: Value,
}

impl<'a> Vm<'a> {
    pub fn new(chunk: &'a Chunk) -> Vm<'a> {
        Vm {
            chunk,
            stack: vec![],
            ip: 0,
            return_register: Value::Nil,
        }
    }

    pub fn push(&mut self, value: Value) {
        println!("{:04} pushing {value:?}", self.ip);
        self.stack.push(value);
    }

    pub fn pop(&mut self) -> Value {
        let value = self.stack.pop().unwrap();
        println!("{:04} popping {value:?}", self.ip);
        value
    }

    pub fn interpret(&mut self) -> Result<Value, Panic> {
        let Some(byte) = self.chunk.bytecode.get(self.ip) else {
            return Ok(self.stack.pop().unwrap());
        };
        let op = Op::from_u8(*byte).unwrap();
        use Op::*;
        match op {
            Constant => {
                let const_idx = self.chunk.bytecode[self.ip + 1];
                let value = self.chunk.constants[const_idx as usize].clone();
                self.push(value);
                self.ip += 2;
                self.interpret()
            }
            Jump => {
                let jump_len = self.chunk.bytecode[self.ip + 1];
                self.ip += jump_len as usize;
                self.interpret()
            }
            JumpIfFalse => {
                let jump_len = self.chunk.bytecode[self.ip + 1];
                let cond = self.pop();
                match cond {
                    Value::Nil | Value::False => {
                        self.ip += jump_len as usize + 2;
                        self.interpret()
                    }
                    _ => {
                        self.ip += 2;
                        self.interpret()
                    }
                }
            }
            Pop => {
                self.pop();
                self.ip += 1;
                self.interpret()
            }
            PushBinding => {
                let binding_idx = self.chunk.bytecode[self.ip + 1] as usize;
                let binding_value = self.stack[binding_idx].clone();
                self.push(binding_value);
                self.ip += 2;
                self.interpret()
            }
            Store => {
                self.return_register = self.pop();
                self.push(Value::Nil);
                self.ip += 1;
                self.interpret()
            }
            Load => {
                let mut value = Value::Nil;
                swap(&mut self.return_register, &mut value);
                self.push(value);
                self.ip += 1;
                self.interpret()
            }
        }
    }
}