use crate::{obj::{reserved::*, prelude::*}, vm::{consts::ConstPool, frame::*, inst::Inst, signal::*, Vm}};
use once_cell::sync::Lazy;
use shredder::{GcSafeWrapper, Scan};
use std::{fmt::{Debug, Formatter, self}, io::{self, Write}, sync::Arc};

pub type FunLocals = Vec<Sym>;

/// A function object which can create a new stack frame.
pub trait Fun {
    /// Creates a new `vm::Frame` from this function object.
    fn create_frame(&self, callee: ObjRef, vm: &Vm, _args: Vec<ObjRef>) -> Frame;
}

pub type MethodRef = ObjRef<Method>;

#[derive(Scan, Debug)]
pub struct Method {
    vtable: Vtable,
    attrs: Attrs,
}

impl Method {
    pub fn new_obj(owner: ObjRef, function: ObjRef) -> MethodRef {
        self_referring_obj! {
            Self {
                vtable: Default::default(),
                attrs: attrs! {
                    SELF_MEMBER_NAME.sym => owner,
                    FUNC_MEMBER_NAME.sym => function,
                }
            },
            vtable: |obj_ref: ObjRef| vtable! {
                CALL_MEMBER_NAME.sym => obj_ref.clone(),
            }
        }
    }
}

impl Obj for Method {
    fn vtable(&self) -> &Vtable {
        &self.vtable
    }

    fn attrs(&self) -> &Attrs {
        &self.attrs
    }

    fn attrs_mut(&mut self) -> Option<&mut Attrs> {
        Some(&mut self.attrs)
    }

    fn as_any(&self) -> &dyn std::any::Any { self }

    fn as_fun(&self) -> Option<&dyn Fun> { Some(self) }
}

impl Fun for Method {
    fn create_frame(&self, _callee: ObjRef, vm: &Vm, mut args: Vec<ObjRef>) -> Frame {
        // insert self argument
        args.insert(0, self.get_attr(SELF_MEMBER_NAME.sym).unwrap());
        // get function and use its create_frame method with the updated args
        let fun_ref = self.get_attr(FUNC_MEMBER_NAME.sym).unwrap();
        read_obj!(let fun_obj = fun_ref);
        let fun = fun_obj.as_fun().unwrap();
        fun.create_frame(fun_ref.clone(), vm, args)
    }
}

//
// struct UserFun
//

pub type UserFunRef = ObjRef<UserFun>;

#[derive(Scan)]
pub struct UserFun {
    vtable: Vtable,
    attrs: Attrs,
    // Safe because Vec<Inst> doesn't need to be scanned
    #[shredder(unsafe_skip)]
    code: Arc<Vec<Inst>>,
    // Safe because this is just an interner that points to symbols, which aren't GC'd
    #[shredder(unsafe_skip)]
    locals: FunLocals,
}

impl UserFun {
    pub fn new_obj(code: Vec<Inst>, locals: FunLocals) -> UserFunRef {
        self_referring_obj! (
            Self {
                vtable: Default::default(), // this is a placeholder for the real vtable
                attrs: Default::default(),
                code: Arc::new(code),
                locals,
            },
            vtable: |obj_ref: ObjRef| vtable! {
                TY_MEMBER_NAME.sym => USER_FUN_TY.clone(),
                CALL_MEMBER_NAME.sym => obj_ref.clone(),
            }
        )
    }

    pub fn code(&self) -> &Vec<Inst> {
        &self.code
    }

    pub fn locals(&self) -> &FunLocals {
        &self.locals
    }

    pub fn dump(&self, writer: &mut dyn Write, const_pool: &ConstPool, globals: &FunLocals) -> io::Result<()> {
        fn num_digits(n: usize, radix: usize) -> usize {
            ((n as f64) + 1.0).log(radix as f64).ceil() as usize
        }

        // column widths
        let addr_w = num_digits(self.code().len(), 16).max(4);
        let inst_col = 16 - addr_w;
        let inst_w = 16;

        for (addr, inst) in self.code().iter().enumerate() {
            let (param_val, mut param_name) = match inst {
                Inst::PushSym(sym) | Inst::GetAttr(sym) | Inst::SetAttr(sym) => (
                    sym.index().to_string(),
                    global_sym_lookup(*sym).unwrap().to_string(),
                ),
                Inst::PushConst(hdl) => (
                    hdl.index().to_string(),
                    {
                        let obj_ref = const_pool.get(*hdl);
                        read_obj!(let obj = obj_ref);
                        format!("{:?}", obj)
                    },
                ),
                Inst::LoadLocal(local) => (
                    local.index().to_string(),
                    global_sym_lookup(self.locals()[local.index()]).unwrap().to_string(),
                ),
                Inst::LoadGlobal(global) => (
                    global.index().to_string(),
                    global_sym_lookup(globals[global.index()]).unwrap().to_string(),
                ),
                Inst::PopLocal(local) => {
                    if let Some(local) = local {
                        let index = local.index();
                        let sym = self.locals()[index];
                        let name = global_sym_lookup(sym).unwrap().to_string();
                        (index.to_string(), name)
                    } else {
                        ("(discarded)".to_string(), String::new())
                    }
                }
                Inst::PopGlobal(global) => {
                    if let Some(global) = global{
                        let index = global.index();
                        let sym = globals[index];
                        let name = global_sym_lookup(sym).unwrap().to_string();
                        (index.to_string(), name)
                    } else {
                        ("(discarded)".to_string(), String::new())
                    }
                }
                Inst::Jump(addr) | Inst::JumpTrue(addr) => (
                    addr.to_string(),
                    String::new(),
                ),
                Inst::Call(argc) => (
                    argc.to_string(),
                    String::new(),
                ),
                _ => (String::new(), String::new()),
            };

            if !param_name.is_empty() {
                param_name = format!("({})", param_name);
            }

            writeln!(
                writer,
                "{:0addr_w$x}{space: >inst_col$}{: <inst_w$}{: <4} {}",
                addr,
                inst.name(),
                param_val,
                param_name,
                space = "",
                addr_w = addr_w,
                inst_w = inst_w,
                inst_col = inst_col,
            )?;
        }
        Ok(())
    }
}

impl Obj for UserFun {
    fn vtable(&self) -> &Vtable {
        &self.vtable
    }

    fn attrs(&self) -> &Attrs {
        &self.attrs
    }

    fn attrs_mut(&mut self) -> Option<&mut Attrs> {
        Some(&mut self.attrs)
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn as_fun(&self) -> Option<&dyn Fun> {
        Some(self)
    }
}

impl Fun for UserFun {
    fn create_frame(&self, callee: ObjRef, vm: &Vm, args: Vec<ObjRef>) -> Frame {
        let bindings: FrameBindings = args.into_iter()
            .enumerate()
            .collect();
        Frame::User(
            UserFrame::new(callee, bindings, vm.pc(), Arc::clone(&self.code))
        )
    }
}

impl Debug for UserFun {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("UserFun")
            .field("attrs", &self.attrs)
            .field("code", &self.code)
            .field("locals", &self.locals)
            .finish()
    }
}

pub static USER_FUN_TY: Lazy<ObjRef<Ty>> = Lazy::new(|| Ty::new_obj(USER_FUN_NAME.sym_ref()));

//
// struct NativeFun
//

pub type NativeFunPtr = fn(ObjRef, &mut Vm, Vec<ObjRef>) -> Signal;
pub type NativeFunRef = ObjRef<NativeFun>;

#[derive(Scan)]
pub struct NativeFun {
    vtable: Vtable,
    attrs: Attrs,
    #[shredder(skip)]
    fun: GcSafeWrapper<NativeFunPtr>,
}

//
// impl NativeFun
//

impl NativeFun {
    pub fn new_obj(fun: NativeFunPtr) -> ObjRef<Self> {
        self_referring_obj! (
            Self {
                vtable: Default::default(),
                attrs: Default::default(),
                fun: GcSafeWrapper::new(fun),
            },
            vtable: |obj_ref: ObjRef| vtable! {
                TY_MEMBER_NAME.sym => NATIVE_FUN_TY.clone(),
                CALL_MEMBER_NAME.sym => obj_ref.clone(),
            }
        )
    }

    /// Gets the native function that can be invoked.
    pub fn fun(&self) -> &NativeFunPtr {
        &self.fun
    }
}

//
// impl Debug for NativeFun
//

impl Debug for NativeFun {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("NativeFun")
            .field("fun", &format!("(function at {:x})", &self.fun as *const _ as usize))
            .field("attrs", &self.attrs)
            .finish()
    }
}

impl Obj for NativeFun {
    fn vtable(&self) -> &Vtable {
        &self.vtable
    }

    fn attrs(&self) -> &Attrs {
        &self.attrs
    }

    fn attrs_mut(&mut self) -> Option<&mut Attrs> {
        None
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn as_fun(&self) -> Option<&dyn Fun> {
        Some(self)
    }
}

impl Fun for NativeFun {
    fn create_frame(&self, callee: ObjRef, _vm: &Vm, args: Vec<ObjRef>) -> Frame {
        Frame::Native(
            NativeFrame::new(callee, *self.fun, args)
        )
    }
}

pub static NATIVE_FUN_TY: Lazy<ObjRef<Ty>> = Lazy::new(|| Ty::new_obj(NATIVE_FUN_NAME.sym_ref()));