use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind};
use crate::ast::compiler::tco::ExecNode;
use crate::ast::nodes::Node;
use crate::ast::types::{Object, Value};
use std::any::Any;
use std::cell::RefCell;
use std::rc::Rc;

#[derive(Debug, Clone)]
pub struct Closure {
    /// The analyzed parameter pattern.
    pub parameter_node: Rc<AnalyzedNode>,
    /// The analyzed body (before TCO).
    pub function_node: Rc<AnalyzedNode>,
    /// The executable node (after TCO).
    pub exec_node: Rc<ExecNode>,
    pub upvalues: Vec<Rc<RefCell<Value>>>,
    pub positional_count: Option<u32>,
}

impl Closure {
    #[inline]
    pub fn new(
        params: Rc<AnalyzedNode>,
        body: Rc<AnalyzedNode>,
        exec: Rc<ExecNode>,
        upvalues: Vec<Rc<RefCell<Value>>>,
        positional_count: Option<u32>,
    ) -> Self {
        Self {
            parameter_node: params,
            function_node: body,
            exec_node: exec,
            upvalues,
            positional_count,
        }
    }
}

impl Object for Closure {
    fn type_name(&self) -> &'static str {
        "closure"
    }
    fn as_any(&self) -> &dyn Any {
        self
    }
}

#[derive(Debug)]
struct CallFrame {
    stack_base: usize,
    closure: Option<Rc<Closure>>,
}

pub trait VMObserver {
    const ACTIVE: bool = false;
    fn before_eval(&mut self, _vm: &VM, _node: &ExecNode) {}
    fn after_eval(&mut self, _vm: &VM, _node: &ExecNode, _res: &Result<Value, String>) {}
}

pub struct NoOpObserver;
impl VMObserver for NoOpObserver {}

pub struct TracingObserver {
    pub logs: Vec<String>,
    indent: usize,
}

impl TracingObserver {
    pub fn new() -> Self {
        Self {
            logs: Vec::new(),
            indent: 0,
        }
    }
    fn pad(&self) -> String {
        "|  ".repeat(self.indent)
    }
}

impl Default for TracingObserver {
    fn default() -> Self {
        Self::new()
    }
}

impl VMObserver for TracingObserver {
    const ACTIVE: bool = true;
    fn before_eval(&mut self, _vm: &VM, node: &ExecNode) {
        let pad = self.pad();
        let metrics = &node.ty.original.ty;
        self.logs.push(format!(
            "{}{} [{} | P:{:?}{}]: {{",
            pad,
            node.kind.display_name(),
            node.ty.ty,
            metrics.purity,
            if metrics.is_recursive { " | REC" } else { "" }
        ));
        self.indent += 1;
    }
    fn after_eval(&mut self, vm: &VM, node: &ExecNode, res: &Result<Value, String>) {
        self.indent = self.indent.saturating_sub(1);
        let pad = self.pad();
        match &node.kind {
            BoundKind::Define { .. } | BoundKind::Set { .. } => {
                let s_pad = format!("{}|  ", pad);
                self.logs.push(format!("{}--- Scope Status ---", s_pad));
                self.logs.push(format!(
                    "{}Stack (top 5): {:?}",
                    s_pad,
                    vm.stack.iter().rev().take(5).collect::<Vec<_>>()
                ));
            }
            _ => {}
        }
        let res_str = match res {
            Ok(v) => format!("{}", v),
            Err(e) => format!("ERROR: {}", e),
        };
        self.logs.push(format!("{}}} -> {}", pad, res_str));
    }
}

pub struct VM {
    stack: Vec<Value>,
    globals: Rc<RefCell<Vec<Value>>>,
    frames: Vec<CallFrame>,
}

impl VM {
    pub fn new(globals: Rc<RefCell<Vec<Value>>>) -> Self {
        Self {
            stack: Vec::new(),
            globals,
            frames: Vec::new(),
        }
    }

    pub fn run(&mut self, root: &ExecNode) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();
        self.frames.push(CallFrame {
            stack_base: 0,
            closure: None,
        });
        let mut result = self.eval(root);
        self.frames.pop();
        loop {
            match result {
                Ok(Value::TailCallRequest(payload)) => {
                    let (next_obj, next_args) = *payload;
                    if let Some(closure) = next_obj.as_any().downcast_ref::<Closure>() {
                        self.stack.clear();
                        self.frames.push(CallFrame {
                            stack_base: 0,
                            closure: Some(Rc::new(closure.clone())),
                        });
                        if let Some(count) = closure.positional_count
                            && next_args.len() == count as usize
                        {
                            self.stack.extend(next_args);
                        } else {
                            self.unpack(&closure.parameter_node, &next_args, &mut 0)?;
                        }
                        result = self.eval(&closure.exec_node);
                        self.frames.pop();
                    } else {
                        return Err(format!(
                            "Tail call target is not a closure: {}",
                            next_obj.type_name()
                        ));
                    }
                }
                _ => return result,
            }
        }
    }

    pub fn run_with_args(&mut self, closure: &Closure, args: Vec<Value>) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();
        self.frames.push(CallFrame {
            stack_base: 0,
            closure: Some(Rc::new(closure.clone())),
        });
        if let Some(count) = closure.positional_count
            && args.len() == count as usize
        {
            self.stack.extend(args);
        } else {
            self.unpack(&closure.parameter_node, &args, &mut 0)?;
        }
        self.eval(&closure.exec_node)
    }

    pub fn run_with_args_observed<O: VMObserver>(
        &mut self,
        observer: &mut O,
        closure: &Closure,
        args: Vec<Value>,
    ) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();
        self.frames.push(CallFrame {
            stack_base: 0,
            closure: Some(Rc::new(closure.clone())),
        });
        if let Some(count) = closure.positional_count
            && args.len() == count as usize
        {
            self.stack.extend(args);
        } else {
            self.unpack(&closure.parameter_node, &args, &mut 0)?;
        }
        self.eval_observed(observer, &closure.exec_node)
    }

    pub fn run_with_observer<O: VMObserver>(
        &mut self,
        observer: &mut O,
        root: &ExecNode,
    ) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();
        self.frames.push(CallFrame {
            stack_base: 0,
            closure: None,
        });
        let result = self.eval_observed(observer, root);
        self.frames.pop();
        result
    }

    #[inline(always)]
    fn eval(&mut self, node: &ExecNode) -> Result<Value, String> {
        self.eval_internal(&mut NoOpObserver, node)
    }

    fn eval_observed<O: VMObserver>(
        &mut self,
        observer: &mut O,
        node: &ExecNode,
    ) -> Result<Value, String> {
        self.eval_internal(observer, node)
    }

    #[inline(always)]
    fn eval_internal<O: VMObserver>(
        &mut self,
        obs: &mut O,
        node: &ExecNode,
    ) -> Result<Value, String> {
        if O::ACTIVE {
            obs.before_eval(self, node);
            let result = self.eval_core(obs, node);
            obs.after_eval(self, node, &result);
            result
        } else {
            self.eval_core(obs, node)
        }
    }

    #[inline(always)]
    fn eval_core<O: VMObserver>(&mut self, obs: &mut O, node: &ExecNode) -> Result<Value, String> {
        match &node.kind {
            BoundKind::Nop => Ok(Value::Void),
            BoundKind::Constant(v) => Ok(v.clone()),
            BoundKind::Define {
                addr,
                value,
                captured_by,
                ..
            } => {
                let val = self.eval_internal(obs, value)?;
                let final_val = if !captured_by.is_empty() && matches!(addr, Address::Local(_)) {
                    Value::Cell(Rc::new(RefCell::new(val)))
                } else {
                    val
                };
                self.set_value(*addr, final_val.clone())?;
                Ok(final_val)
            }
            BoundKind::Destructure { pattern, value } => {
                let val = self.eval_internal(obs, value)?;
                let mut offset = 0;

                // Destructuring works on tuples/vectors, or single values wrapped in a slice
                if let Some(vals) = val.as_slice() {
                    self.unpack(pattern, vals, &mut offset)?;
                } else {
                    self.unpack(pattern, std::slice::from_ref(&val), &mut offset)?;
                }
                Ok(val)
            }
            BoundKind::Get { addr, .. } => self.get_value(*addr),
            BoundKind::Set { addr, value } => {
                let val = self.eval_internal(obs, value)?;
                self.set_value(*addr, val.clone())?;
                Ok(val)
            }
            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let c = self.eval_internal(obs, cond)?;
                if c.is_truthy() {
                    self.eval_internal(obs, then_br)
                } else if let Some(e) = else_br {
                    self.eval_internal(obs, e)
                } else {
                    Ok(Value::Void)
                }
            }
            BoundKind::Block { exprs } => {
                let mut last = Value::Void;
                for e in exprs {
                    last = self.eval_internal(obs, e)?;
                }
                Ok(last)
            }
            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => {
                let mut captured = Vec::with_capacity(upvalues.len());
                for addr in upvalues {
                    captured.push(self.capture_upvalue(*addr)?);
                }
                let closure = Closure::new(
                    params.ty.original.clone(),
                    body.ty.original.clone(),
                    body.clone(),
                    captured,
                    *positional_count,
                );
                Ok(Value::Object(Rc::new(closure)))
            }
            BoundKind::Call { callee, args } => {
                let mut func_val = self.eval_internal(obs, callee)?;

                let mut arg_vals = match &args.kind {
                    BoundKind::Tuple { elements } => {
                        let mut vals = Vec::with_capacity(elements.len());
                        let mut is_complex = false;
                        for e in elements {
                            if matches!(e.kind, BoundKind::Tuple { .. }) {
                                is_complex = true;
                                break;
                            }
                            vals.push(self.eval_internal(obs, e)?);
                        }
                        if is_complex {
                            self.prepare_args_internal(obs, args)?
                        } else {
                            vals
                        }
                    }
                    _ => self.prepare_args_internal(obs, args)?,
                };

                if node.ty.is_tail {
                    match func_val {
                        Value::Object(obj) => {
                            return Ok(Value::TailCallRequest(Box::new((obj, arg_vals))));
                        }
                        Value::Function(f) => return Ok((f.func)(arg_vals)),
                        _ => {
                            return Err(format!(
                                "Tail call target is not a function: {}",
                                func_val
                            ));
                        }
                    }
                }

                loop {
                    match func_val {
                        Value::Function(f) => break Ok((f.func)(arg_vals)),
                        Value::Object(obj) => {
                            if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
                                let old_stack_top = self.stack.len();
                                let closure_rc = Rc::new(closure.clone());
                                self.frames.push(CallFrame {
                                    stack_base: old_stack_top,
                                    closure: Some(closure_rc.clone()),
                                });
                                if let Some(count) = closure.positional_count
                                    && arg_vals.len() == count as usize
                                {
                                    self.stack.extend(arg_vals);
                                } else {
                                    // Map each parameter pattern to the provided arguments
                                    if let BoundKind::Tuple { elements } =
                                        &closure.parameter_node.kind
                                    {
                                        let mut offset = 0;
                                        for el in elements {
                                            self.unpack(el, &arg_vals, &mut offset)?;
                                        }
                                    } else {
                                        self.unpack(&closure.parameter_node, &arg_vals, &mut 0)?;
                                    }
                                }
                                let result = self.eval_internal(obs, &closure.exec_node);
                                self.frames.pop();
                                self.stack.truncate(old_stack_top);
                                match result {
                                    Ok(Value::TailCallRequest(payload)) => {
                                        let (next_obj, next_args) = *payload;
                                        func_val = Value::Object(next_obj);
                                        arg_vals = next_args;
                                        continue;
                                    }
                                    res => break res,
                                }
                            } else {
                                break Err(format!("Object is not a closure: {}", obj.type_name()));
                            }
                        }
                        _ => break Err(format!("Attempt to call non-function: {}", func_val)),
                    }
                }
            }
            BoundKind::Again { args } => {
                let arg_vals = match &args.kind {
                    BoundKind::Tuple { elements } => {
                        let mut vals = Vec::with_capacity(elements.len());
                        let mut is_complex = false;
                        for e in elements {
                            if matches!(e.kind, BoundKind::Tuple { .. }) {
                                is_complex = true;
                                break;
                            }
                            vals.push(self.eval_internal(obs, e)?);
                        }
                        if is_complex {
                            self.prepare_args_internal(obs, args)?
                        } else {
                            vals
                        }
                    }
                    _ => self.prepare_args_internal(obs, args)?,
                };

                let frame = self.frames.last().ok_or("No call frame for 'again'")?;
                if let Some(closure) = &frame.closure {
                    Ok(Value::TailCallRequest(Box::new((
                        Rc::new(closure.as_ref().clone()) as Rc<dyn Object>,
                        arg_vals,
                    ))))
                } else {
                    Err("'again' called outside of a closure".to_string())
                }
            }
            BoundKind::Tuple { elements } => {
                let mut vals = Vec::with_capacity(elements.len());
                for e in elements {
                    vals.push(self.eval_internal(obs, e)?);
                }
                Ok(Value::make_tuple(vals))
            }
            BoundKind::Record { fields } => {
                let mut keys = Vec::with_capacity(fields.len());
                let mut values = Vec::with_capacity(fields.len());
                for (k, v) in fields {
                    let key = self.eval_internal(obs, k)?;
                    let val = self.eval_internal(obs, v)?;
                    if let Value::Keyword(kw) = key {
                        keys.push(kw);
                        values.push(val);
                    } else {
                        return Err(format!("Record key must be keyword, got {}", key));
                    }
                }
                Ok(Value::make_record(keys, values))
            }
            BoundKind::Expansion { bound_expanded, .. } => self.eval_internal(obs, bound_expanded),
            BoundKind::Extension(ext) => Err(format!(
                "Execution of extension '{}' not implemented yet",
                ext.display_name()
            )),
        }
    }

    pub fn resolve_tail_calls(&mut self, mut result: Value) -> Value {
        while let Value::TailCallRequest(payload) = result {
            let (next_obj, next_args) = *payload;
            if let Some(closure) = next_obj.as_any().downcast_ref::<Closure>() {
                result = match self.run_with_args(closure, next_args) {
                    Ok(v) => v,
                    Err(e) => panic!("Myc Runtime Error (TailCall): {}", e),
                };
            } else {
                panic!(
                    "Tail call target is not a closure: {}",
                    next_obj.type_name()
                );
            }
        }
        result
    }

    fn capture_upvalue(&mut self, addr: Address) -> Result<Rc<RefCell<Value>>, String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                    if let Value::Cell(cell) = &self.stack[abs_index] {
                        Ok(cell.clone())
                    } else {
                        let val = self.stack[abs_index].clone();
                        let cell = Rc::new(RefCell::new(val));
                        self.stack[abs_index] = Value::Cell(cell.clone());
                        Ok(cell)
                    }
                } else {
                    Err(format!("Stack underflow capture local {}", idx))
                }
            }
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        Ok(closure.upvalues[idx].clone())
                    } else {
                        Err(format!("Upvalue access out of bounds capture {}", idx))
                    }
                } else {
                    Err("Current frame has no closure".to_string())
                }
            }
            Address::Global(_) => Err("Cannot capture global directly".to_string()),
        }
    }

    fn get_value(&self, addr: Address) -> Result<Value, String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                    match &self.stack[abs_index] {
                        Value::Cell(cell) => Ok(cell.borrow().clone()),
                        val => Ok(val.clone()),
                    }
                } else {
                    Err(format!("Stack underflow access local {}", idx))
                }
            }
            Address::Global(idx) => {
                let idx = idx as usize;
                let globals = self.globals.borrow();
                if idx < globals.len() {
                    Ok(globals[idx].clone())
                } else {
                    Err(format!("Global access out of bounds {}", idx))
                }
            }
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        Ok(closure.upvalues[idx].borrow().clone())
                    } else {
                        Err(format!("Upvalue access out of bounds {}", idx))
                    }
                } else {
                    Err("Current frame has no closure (cannot access upvalues)".to_string())
                }
            }
        }
    }

    fn set_value(&mut self, addr: Address, value: Value) -> Result<(), String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                    if let Value::Cell(cell) = &self.stack[abs_index] {
                        *cell.borrow_mut() = value;
                    } else {
                        self.stack[abs_index] = value;
                    }
                } else if abs_index == self.stack.len() {
                    self.stack.push(value);
                } else {
                    return Err(format!("Stack gap write local {}", idx));
                }
                Ok(())
            }
            Address::Global(idx) => {
                let idx = idx as usize;
                let mut globals = self.globals.borrow_mut();
                if idx >= globals.len() {
                    globals.resize(idx + 1, Value::Void);
                }
                globals[idx] = value;
                Ok(())
            }
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        *closure.upvalues[idx].borrow_mut() = value;
                        Ok(())
                    } else {
                        Err(format!("Upvalue assignment out of bounds {}", idx))
                    }
                } else {
                    Err("Current frame has no closure".to_string())
                }
            }
        }
    }

    fn flatten_value(val: Value, into: &mut Vec<Value>) {
        if let Some(values) = val.as_slice() {
            into.extend_from_slice(values);
        } else {
            into.push(val);
        }
    }

    fn prepare_args_internal<O: VMObserver>(
        &mut self,
        obs: &mut O,
        args: &ExecNode,
    ) -> Result<Vec<Value>, String> {
        let mut arg_vals = Vec::new();
        match &args.kind {
            BoundKind::Tuple { elements } => {
                self.eval_and_flatten_internal(obs, elements, &mut arg_vals)?;
            }
            _ => {
                VM::flatten_value(self.eval_internal(obs, args)?, &mut arg_vals);
            }
        }
        Ok(arg_vals)
    }

    fn eval_and_flatten_internal<O: VMObserver>(
        &mut self,
        obs: &mut O,
        elements: &[ExecNode],
        into: &mut Vec<Value>,
    ) -> Result<(), String> {
        for e in elements {
            into.push(self.eval_internal(obs, e)?);
        }
        Ok(())
    }

    fn unpack<T>(
        &mut self,
        pattern: &Node<BoundKind<T>, T>,
        values: &[Value],
        offset: &mut usize,
    ) -> Result<(), String> {
        match &pattern.kind {
            BoundKind::Define { addr, .. } => {
                let val = values.get(*offset).cloned().unwrap_or(Value::Void);
                *offset += 1;
                self.set_value(*addr, val)
            }
            BoundKind::Set { addr, .. } => {
                let val = values.get(*offset).cloned().unwrap_or(Value::Void);
                *offset += 1;
                self.set_value(*addr, val)
            }
            BoundKind::Tuple { elements } => {
                if let Some(sub_values) = values.get(*offset).and_then(|v| v.as_slice()) {
                    *offset += 1;
                    let mut sub_offset = 0;
                    for el in elements {
                        self.unpack(el, sub_values, &mut sub_offset)?;
                    }
                    return Ok(());
                }
                for el in elements {
                    self.unpack(el, values, offset)?;
                }
                Ok(())
            }
            _ => Err("Invalid node in parameter pattern".to_string()),
        }
    }
}