use crate::compiler::{Chunk, Op};
use crate::parser::Ast;
use crate::spans::Spanned;
use crate::value::Value;
use chumsky::prelude::SimpleSpan;
use imbl::{HashMap, Vector};
use num_traits::FromPrimitive;
use std::cell::RefCell;
use std::mem::swap;
use std::rc::Rc;

#[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,
// }
pub struct Panic(&'static str);

#[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,
    pub matches: bool,
}

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

    pub fn push(&mut self, value: Value) {
        self.stack.push(value);
    }

    pub fn pop(&mut self) -> Value {
        self.stack.pop().unwrap()
    }

    fn print_stack(&self) {
        let inner = self
            .stack
            .iter()
            .map(|val| val.to_string())
            .collect::<Vec<_>>()
            .join("|");
        println!("{:04}: [{inner}] {}", self.ip, self.return_register);
    }

    fn print_debug(&self) {
        self.chunk.dissasemble_instr(self.ip);
        self.print_stack();
    }

    pub fn interpret(&mut self) -> Result<Value, Panic> {
        let Some(byte) = self.chunk.bytecode.get(self.ip) else {
            return Ok(self.stack.pop().unwrap());
        };
        if crate::DEBUG_RUN {
            self.print_debug();
        }
        let op = Op::from_u8(*byte).unwrap();
        use Op::*;
        match op {
            Nil => {
                self.push(Value::Nil);
                self.ip += 1;
                self.interpret()
            }
            True => {
                self.push(Value::True);
                self.ip += 1;
                self.interpret()
            }
            False => {
                self.push(Value::False);
                self.ip += 1;
                self.interpret()
            }
            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()
            }
            ResetMatch => {
                self.matches = false;
                self.ip += 1;
                self.interpret()
            }
            MatchWord => {
                self.matches = true;
                self.ip += 1;
                self.interpret()
            }
            MatchNil => {
                if *self.stack.last().unwrap() == Value::Nil {
                    self.matches = true;
                };
                self.ip += 1;
                self.interpret()
            }
            MatchTrue => {
                if *self.stack.last().unwrap() == Value::True {
                    self.matches = true;
                };
                self.ip += 1;
                self.interpret()
            }
            MatchFalse => {
                if *self.stack.last().unwrap() == Value::False {
                    self.matches = true;
                }
                self.ip += 1;
                self.interpret()
            }
            PanicIfNoMatch => {
                if !self.matches {
                    Err(Panic("no match"))
                } else {
                    self.ip += 1;
                    self.interpret()
                }
            }
            MatchConstant => {
                let const_idx = self.chunk.bytecode[self.ip + 1];
                let value = self.stack.last().unwrap();
                self.matches = *value == self.chunk.constants[const_idx as usize];
                self.ip += 2;
                self.interpret()
            }
            PushTuple => {
                let tuple_len = self.chunk.bytecode[self.ip + 1];
                let tuple_members = self.stack.split_off(self.stack.len() - tuple_len as usize);
                let tuple = Value::Tuple(Rc::new(tuple_members));
                self.stack.push(tuple);
                self.ip += 2;
                self.interpret()
            }
            PushList => {
                let list_len = self.chunk.bytecode[self.ip + 1];
                let list_members = self.stack.split_off(self.stack.len() - list_len as usize);
                let list = Value::List(Box::new(Vector::from(list_members)));
                self.stack.push(list);
                self.ip += 2;
                self.interpret()
            }
            PushDict => {
                let dict_len = self.chunk.bytecode[self.ip + 1] as usize * 2;
                let dict_members = self.stack.split_off(self.stack.len() - dict_len);
                let mut dict = HashMap::new();
                let mut dict_iter = dict_members.iter();
                while let Some(kw) = dict_iter.next() {
                    let Value::Keyword(key) = kw else {
                        unreachable!()
                    };
                    let value = dict_iter.next().unwrap();
                    dict.insert(*key, value.clone());
                }
                self.stack.push(Value::Dict(Box::new(dict)));
                self.ip += 2;
                self.interpret()
            }
            PushBox => {
                let val = self.pop();
                self.stack.push(Value::Box(Rc::new(RefCell::new(val))));
                self.ip += 1;
                self.interpret()
            }
            GetKey => {
                let key = self.pop();
                let Value::Keyword(idx) = key else {
                    unreachable!()
                };
                let dict = self.pop();
                let value = match dict {
                    Value::Dict(d) => d.as_ref().get(&idx).unwrap_or(&Value::Nil).clone(),
                    _ => Value::Nil,
                };
                self.push(value);
                self.ip += 1;
                self.interpret()
            }
            MatchTuple => {
                todo!()
            }
            JumpIfNoMatch => {
                let jump_len = self.chunk.bytecode[self.ip + 1] as usize;
                if !self.matches {
                    self.ip += jump_len + 2;
                } else {
                    self.ip += 2;
                }
                self.interpret()
            }
            PanicNoWhen | PanicNoMatch => Err(Panic("no match")),
        }
    }
}