use std::rc::Rc;
use crate::ast::types::{Value, StaticType, Identity};
use crate::ast::nodes::{Node, Symbol};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Address {
    Local(u32),   // Stack-Slot index (relative to frame base)
    Upvalue(u32), // Index in the closure's upvalue array
    Global(u32),  // Index in the global environment vector
}

/// Trait for DSL-specific AST extensions (like Pipe, Series-Access, etc.)
pub trait BoundExtension<T>: std::fmt::Debug {
    fn clone_box(&self) -> Box<dyn BoundExtension<T>>;
    fn display_name(&self) -> String;
}

impl<T> Clone for Box<dyn BoundExtension<T>> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

/// Type alias for a node that has been bound. Defaults to untyped (T = ()).
pub type BoundNode<T = ()> = Node<BoundKind<T>, T>;

/// Type alias for a node that has been fully type-checked.
pub type TypedNode = BoundNode<StaticType>;

#[derive(Debug, Clone)]
pub enum BoundKind<T = ()> {
    Nop,
    Constant(Value),
    
    /// A positional parameter in a Lambda's parameter list.
    Parameter {
        name: Symbol,
        slot: u32,
    },
    
    // Variable Access (Resolved)
    Get {
        addr: Address,
        name: Symbol,
    },
    
    // Variable Update (Assignment)
    Set {
        addr: Address,
        value: Box<BoundNode<T>>,
    },

    // Variable Declaration (Local)
    DefLocal {
        name: Symbol,
        slot: u32,
        value: Box<BoundNode<T>>,
        captured_by: Vec<Identity>,
    },

    If {
        cond: Box<BoundNode<T>>,
        then_br: Box<BoundNode<T>>,
        else_br: Option<Box<BoundNode<T>>>,
    },
    
    // Global Definition (Locals are implicit by stack position)
    DefGlobal {
        name: Symbol,
        global_index: u32,
        value: Box<BoundNode<T>>,
    },

    Lambda {
        params: Rc<BoundNode<T>>,
        // The list of variables captured from enclosing scopes
        upvalues: Vec<Address>, 
        body: Rc<BoundNode<T>>,
        /// Static optimization: number of positional parameters if the pattern is flat.
        positional_count: Option<u32>,
    },

    Call {
        callee: Box<BoundNode<T>>,
        args: Box<BoundNode<T>>,
    },

    Block {
        exprs: Vec<BoundNode<T>>,
    },

    Tuple {
        elements: Vec<BoundNode<T>>,
    },

    Record {
        fields: Vec<RecordField<T>>,
    },

    /// An expanded macro call, preserving the original call for debugging and UI.
    Expansion {
        /// The original call from the untyped AST.
        original_call: Rc<Node<crate::ast::nodes::UntypedKind>>,
        /// The result of binding the expanded AST.
        bound_expanded: Box<BoundNode<T>>,
    },

    /// DSL-specific extension slot
    Extension(Box<dyn BoundExtension<T>>),
}

impl<T> PartialEq for BoundKind<T> 
where T: PartialEq {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (BoundKind::Nop, BoundKind::Nop) => true,
            (BoundKind::Constant(a), BoundKind::Constant(b)) => a == b,
            (BoundKind::Parameter { name: na, slot: sa }, BoundKind::Parameter { name: nb, slot: sb }) => {
                na == nb && sa == sb
            }
            (BoundKind::Get { addr: aa, name: na }, BoundKind::Get { addr: ab, name: nb }) => {
                aa == ab && na == nb
            }
            (BoundKind::Set { addr: aa, value: va }, BoundKind::Set { addr: ab, value: vb }) => {
                aa == ab && va == vb
            }
            (BoundKind::DefLocal { name: na, slot: sa, value: va, captured_by: ca }, 
             BoundKind::DefLocal { name: nb, slot: sb, value: vb, captured_by: cb }) => {
                na == nb && sa == sb && va == vb && ca == cb
            }
            (BoundKind::If { cond: ca, then_br: ta, else_br: ea }, 
             BoundKind::If { cond: cb, then_br: tb, else_br: eb }) => {
                ca == cb && ta == tb && ea == eb
            }
            (BoundKind::DefGlobal { name: na, global_index: ga, value: va }, 
             BoundKind::DefGlobal { name: nb, global_index: gb, value: vb }) => {
                na == nb && ga == gb && va == vb
            }
            (BoundKind::Lambda { params: pa, upvalues: ua, body: ba, positional_count: pca }, 
             BoundKind::Lambda { params: pb, upvalues: ub, body: bb, positional_count: pcb }) => {
                pa == pb && ua == ub && ba == bb && pca == pcb
            }
            (BoundKind::Call { callee: ca, args: aa }, BoundKind::Call { callee: cb, args: ab }) => {
                ca == cb && aa == ab
            }
            (BoundKind::Block { exprs: ea }, BoundKind::Block { exprs: eb }) => {
                ea == eb
            }
            (BoundKind::Tuple { elements: ea }, BoundKind::Tuple { elements: eb }) => {
                ea == eb
            }
            (BoundKind::Record { fields: fa }, BoundKind::Record { fields: fb }) => {
                fa == fb
            }
            (BoundKind::Expansion { original_call: oa, bound_expanded: ba }, 
             BoundKind::Expansion { original_call: ob, bound_expanded: bb }) => {
                Rc::ptr_eq(oa, ob) && ba == bb
            }
            (BoundKind::Extension(_), BoundKind::Extension(_)) => false, // Currently no equality for extensions
            _ => false,
        }
    }
}

/// A single field in a Record literal (Key-Value pair)
pub type RecordField<T> = (BoundNode<T>, BoundNode<T>);

impl<T> BoundKind<T> {
    pub fn display_name(&self) -> String {
        match self {
            BoundKind::Nop => "NOP".to_string(),
            BoundKind::Constant(v) => format!("CONST({})", v),
            BoundKind::Parameter { name, slot } => format!("PARAM({}, Slot:{})", name.name, slot),
            BoundKind::Get { addr, name } => format!("GET({}, {:?})", name.name, addr),
            BoundKind::Set { addr, .. } => format!("SET({:?})", addr),
            BoundKind::DefLocal { name, slot, .. } => format!("DEF_LOCAL({}, Slot:{})", name.name, slot),
            BoundKind::If { .. } => "IF".to_string(),
            BoundKind::DefGlobal { name, global_index, .. } => format!("DEF_GLOBAL({}, Idx:{})", name.name, global_index),
            BoundKind::Lambda { params, upvalues, .. } => {
                let p_str = match &params.kind {
                    BoundKind::Tuple { elements } => format!("p:{}", elements.len()),
                    _ => "p:1".to_string(),
                };
                format!("LAMBDA({}, Captures:{})", p_str, upvalues.len())
            },
            BoundKind::Call { .. } => "CALL".to_string(),
            BoundKind::Block { .. } => "BLOCK".to_string(),
            BoundKind::Tuple { elements } => format!("TUPLE({})", elements.len()),
            BoundKind::Record { fields } => format!("RECORD({})", fields.len()),
            BoundKind::Expansion { .. } => "EXPANSION".to_string(),
            BoundKind::Extension(ext) => ext.display_name(),
        }
    }
}