use std::{
    borrow::Borrow,
    cell::RefCell,
    collections::HashMap,
    fmt,
    hash::Hash,
    rc::Rc,
    sync::atomic::{AtomicU32, Ordering},
};

use crate::{
    error::MachineError,
    vm::{
        machine::Machine,
        value::{
            BytecodeFunction, IdentifierValue, NativeFunction, NativeObject, StringValue, Value,
            convert::{AnyFunction, TryFromValue},
        },
    },
};

#[derive(Clone, Debug)]
pub enum Macro {
    Native(NativeFunction),
    Bytecode(Rc<BytecodeFunction>),
}

pub trait EnvironmentAccessHook {
    fn read_variable(&mut self, name: &str) -> Option<Value>;
    fn write_variable(
        &mut self,
        name: &IdentifierValue,
        value: &Value,
    ) -> Result<bool, MachineError>;
}

#[derive(Default)]
pub struct Environment {
    globals: RefCell<HashMap<IdentifierValue, Value>>,
    macros: RefCell<HashMap<IdentifierValue, Macro>>,
    parent: Option<Rc<Environment>>,
    gensym_index: AtomicU32,

    access_hook: RefCell<Option<Box<dyn EnvironmentAccessHook>>>,
}

impl Environment {
    pub fn new(parent: Option<Rc<Self>>) -> Self {
        Self {
            parent,
            ..Default::default()
        }
    }

    pub fn gensym(&self) -> IdentifierValue {
        // TODO do this in a root environment?
        let index = self.gensym_index.fetch_add(1, Ordering::SeqCst);
        format!("___gensym{index}").into()
    }

    pub fn set_access_hook(
        &self,
        hook: Option<Box<dyn EnvironmentAccessHook>>,
    ) -> Option<Box<dyn EnvironmentAccessHook>> {
        self.access_hook.replace(hook)
    }

    pub fn defun_native<S, D, F>(&self, identifier: S, docstring: D, function: F) -> Value
    where
        S: Into<IdentifierValue>,
        D: Into<StringValue>,
        F: Fn(&mut Machine, &Rc<Environment>, &[Value]) -> Result<Value, MachineError> + 'static,
    {
        let identifier = identifier.into();
        let native = NativeFunction::new(identifier.clone(), docstring, function);
        let value = Value::NativeFunction(native);
        self.set_global_value(identifier, value.clone())
            .expect("defun_native failed");
        value
    }

    pub fn defmacro_native<S, D, F>(&self, identifier: S, docstring: D, function: F)
    where
        S: Into<IdentifierValue>,
        D: Into<StringValue>,
        F: Fn(&mut Machine, &Rc<Environment>, &[Value]) -> Result<Value, MachineError> + 'static,
    {
        let identifier = identifier.into();
        let native = NativeFunction::new(identifier.clone(), docstring, function);
        self.macros
            .borrow_mut()
            .insert(identifier, Macro::Native(native));
    }

    pub fn defmacro_bytecode<S: Into<IdentifierValue>>(
        &self,
        identifier: S,
        value: Rc<BytecodeFunction>,
    ) {
        let identifier = identifier.into();
        // println!("Export macro: {identifier}: {value}");
        self.macros
            .borrow_mut()
            .insert(identifier, Macro::Bytecode(value));
    }

    pub fn global_value<Q>(&self, identifier: &Q) -> Option<Value>
    where
        IdentifierValue: Borrow<Q>,
        Q: Hash + Eq + ?Sized + AsRef<str>,
    {
        if let Some(hook) = &mut *self.access_hook.borrow_mut()
            && let Some(value) = hook.read_variable(identifier.as_ref())
        {
            return Some(value);
        }

        self.globals.borrow().get(identifier).cloned().or_else(|| {
            self.parent
                .as_ref()
                .and_then(|parent| parent.global_value(identifier))
        })
    }

    pub fn set_global_value<S: Into<IdentifierValue>, V: Into<Value>>(
        &self,
        identifier: S,
        value: V,
    ) -> Result<(), MachineError> {
        let identifier = identifier.into();
        let value = value.into();

        if let Some(hook) = &mut *self.access_hook.borrow_mut() {
            match hook.write_variable(&identifier, &value) {
                Ok(true) => return Ok(()),
                Ok(false) => {}
                Err(error) => return Err(error),
            }
        }

        self.globals.borrow_mut().insert(identifier, value);
        Ok(())
    }

    pub fn global_macro<Q>(&self, identifier: &Q) -> Option<Macro>
    where
        IdentifierValue: Borrow<Q>,
        Q: Hash + Eq,
    {
        self.macros.borrow().get(identifier).cloned().or_else(|| {
            self.parent
                .as_ref()
                .and_then(|parent| parent.global_macro(identifier))
        })
    }

    pub fn evaluate_global_function<Q>(
        self: &Rc<Self>,
        vm: &mut Machine,
        identifier: &Q,
        arguments: &[Value],
    ) -> Result<Value, MachineError>
    where
        IdentifierValue: Borrow<Q>,
        Q: Hash + Eq + ?Sized + AsRef<str>,
        for<'a> &'a Q: Into<IdentifierValue>,
    {
        let value = self
            .global_value(identifier)
            .ok_or_else(|| MachineError::UnboundIdentifier(identifier.into()))?;
        let function = AnyFunction::try_from_value(&value)?;
        function.invoke(vm, self, arguments)
    }
}

impl fmt::Display for Environment {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "environment")
    }
}

impl fmt::Debug for Environment {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Environment")
            .field("globals", &self.globals)
            .field("macros", &self.macros)
            .finish_non_exhaustive()
    }
}

impl NativeObject for Environment {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}
