use indexmap::IndexMap;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::Receiver;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LambdaParam {
    Named(String),
    Ref(String),
    Variadic,
}

#[derive(Clone)]
pub enum FunctionValue {
    Named(String),
    InlineLambdaAST {
        params: Vec<LambdaParam>,
        body_expr: Box<crate::parser::ast::Expr>,
    },
    InlineLambdaASTPartial {
        params: Vec<LambdaParam>,
        body_expr: Box<crate::parser::ast::Expr>,
        bound_args: Vec<Value>,
    },
    CurriedMap(Box<FunctionValue>),
    CurriedMapPlaceholder {
        func: Option<Box<FunctionValue>>,
        data: Option<Value>,
    },
    CurriedFilter(Box<FunctionValue>),
    CurriedFilterPlaceholder {
        pred: Option<Box<FunctionValue>>,
        data: Option<Value>,
    },
    CurriedReducePartial(Box<FunctionValue>, Value),
    CurriedReducePlaceholder {
        func: Option<Box<FunctionValue>>,
        init_val: Option<Value>,
        data: Option<Value>,
    },
    Composition(Vec<Box<FunctionValue>>, ComposeMode),
    CompositionPartial {
        mode: ComposeMode,
        slots: Vec<Option<Box<FunctionValue>>>,
    },
    CurriedProp(String),
    CurriedAssoc(String, Value),
    RustClosure(
        String,
        Arc<Mutex<dyn for<'x> Fn(&'x mut crate::parser::interpreter::Interpreter, Vec<Value>) -> Result<Value, String> + Send + Sync + 'static>>,
        usize,
    ),
}

impl std::fmt::Debug for FunctionValue {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use FunctionValue::*;
        match self {
            Named(name) => write!(f, "Named({})", name),
            InlineLambdaAST { params, .. } => write!(f, "InlineLambdaAST({:?}, ...)", params),
            InlineLambdaASTPartial { params, .. } => write!(f, "InlineLambdaASTPartial({:?}, ...)", params),
            CurriedMap(_) => write!(f, "CurriedMap(...)"),
            CurriedMapPlaceholder { .. } => write!(f, "CurriedMapPlaceholder(...)"),
            CurriedFilter(_) => write!(f, "CurriedFilter(...)"),
            CurriedFilterPlaceholder { .. } => write!(f, "CurriedFilterPlaceholder(...)"),
            CurriedReducePartial(_, _) => write!(f, "CurriedReducePartial(...)"),
            CurriedReducePlaceholder { .. } => write!(f, "CurriedReducePlaceholder(...)"),
            Composition(_, mode) => write!(f, "Composition(..., {:?})", mode),
            CompositionPartial { mode, .. } => write!(f, "CompositionPartial(..., mode={:?})", mode),
            CurriedProp(s) => write!(f, "CurriedProp({})", s),
            CurriedAssoc(k, _) => write!(f, "CurriedAssoc({}, ...)", k),
            RustClosure(desc, _, ar) => write!(f, "RustClosure({}, arity={})", desc, ar),
        }
    }
}

impl PartialEq for FunctionValue {
    fn eq(&self, other: &Self) -> bool {
        use FunctionValue::*;
        match (self, other) {
            (Named(a), Named(b)) => a == b,
            (
                InlineLambdaAST { params: p1, body_expr: b1 },
                InlineLambdaAST { params: p2, body_expr: b2 },
            ) => p1 == p2 && b1 == b2,
            (
                InlineLambdaASTPartial { params: p1, bound_args: ba1, body_expr: be1 },
                InlineLambdaASTPartial { params: p2, bound_args: ba2, body_expr: be2 },
            ) => p1 == p2 && ba1 == ba2 && be1 == be2,
            _ => false,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ComposeMode {
    Compose,
    Pipe,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ElemType {
    Float32,
    Float64,
    Int32,
    Int64,
}

#[derive(Clone)]
pub struct TensorValue {
    pub elem_type: ElemType,
    pub shape: Vec<usize>,
    pub data: Vec<u8>,
}

impl std::fmt::Debug for TensorValue {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "TensorValue {{ elem_type: {:?}, shape: {:?}, data_len: {} }}",
            self.elem_type,
            self.shape,
            self.data.len()
        )
    }
}

impl PartialEq for TensorValue {
    fn eq(&self, other: &Self) -> bool {
        self.elem_type == other.elem_type && self.shape == other.shape && self.data == other.data
    }
}

pub struct StreamHandle {
    pub stream_id: usize,
    pub label: String,
    pub chunk_receiver: Arc<Mutex<Receiver<Vec<u8>>>>,
}

pub enum InkIteratorKind {
    Core(Arc<Mutex<InkIteratorState>>),
    Plugin(Arc<Mutex<Box<dyn LavaStream>>>),
}

pub struct InkIteratorHandle {
    pub kind: InkIteratorKind,
}

impl Clone for StreamHandle {
    fn clone(&self) -> Self {
        StreamHandle {
            stream_id: self.stream_id,
            label: self.label.clone(),
            chunk_receiver: Arc::clone(&self.chunk_receiver),
        }
    }
}

impl Clone for InkIteratorHandle {
    fn clone(&self) -> Self {
        match &self.kind {
            InkIteratorKind::Core(state_arc) => InkIteratorHandle {
                kind: InkIteratorKind::Core(Arc::clone(state_arc)),
            },
            InkIteratorKind::Plugin(ptr) => InkIteratorHandle {
                kind: InkIteratorKind::Plugin(Arc::clone(ptr)),
            },
        }
    }
}

pub struct InkIteratorState {
    pub current: i32,
    pub end: i32,
    pub done: bool,
}

pub trait LavaStream: Send + Sync {
    fn next_value(&mut self) -> Result<Value, String>;
}

#[derive(Clone)]
pub enum Value {
    Int(i32),
    IntArray(Vec<i32>),
    Int2DArray(Vec<Vec<i32>>),
    Float(f64),
    FloatArray(Vec<f64>),
    Float2DArray(Vec<Vec<f64>>),
    Long(i64),
    Bool(bool),
    BoolArray(Vec<bool>),
    Placeholder,
    SingleString(String),
    StrArray(Vec<String>),
    KeyedArray(IndexMap<String, Value>),
    MixedArray(Vec<Value>),
    Function(Box<FunctionValue>),
    Stream(StreamHandle),
    InkIterator(InkIteratorHandle),
    InkTransform(Box<FunctionValue>),
    Tensor(TensorValue),
    Ref(Arc<Mutex<Value>>),
}

impl PartialEq for Value {
    fn eq(&self, other: &Self) -> bool {
        use Value::*;
        match (self, other) {
            (Int(a), Int(b)) => a == b,
            (Float(a), Float(b)) => a == b,
            (Long(a), Long(b)) => a == b,
            (Bool(a), Bool(b)) => a == b,
            (SingleString(a), SingleString(b)) => a == b,
            (StrArray(a), StrArray(b)) => a == b,
            (IntArray(a), IntArray(b)) => a == b,
            (FloatArray(a), FloatArray(b)) => a == b,
            (BoolArray(a), BoolArray(b)) => a == b,
            (KeyedArray(a), KeyedArray(b)) => a == b,
            (MixedArray(a), MixedArray(b)) => a == b,
            (Function(_), Function(_)) => false,
            (Stream(_), Stream(_)) => false,
            (InkIterator(_), InkIterator(_)) => false,
            (InkTransform(_), InkTransform(_)) => false,
            (Tensor(a), Tensor(b)) => a == b,
            (Ref(a), Ref(b)) => Arc::ptr_eq(a, b) || *a.lock().unwrap() == *b.lock().unwrap(),
            (Placeholder, Placeholder) => true,
            _ => false,
        }
    }
}

impl std::fmt::Debug for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use Value::*;
        match self {
            Int(n) => write!(f, "Int({})", n),
            Float(ff) => write!(f, "Float({})", ff),
            Long(l) => write!(f, "Long({})", l),
            Bool(b) => write!(f, "Bool({})", b),
            Placeholder => write!(f, "Placeholder"),
            SingleString(s) => write!(f, "SingleString({:?})", s),
            StrArray(ss) => write!(f, "StrArray({:?})", ss),
            IntArray(xs) => write!(f, "IntArray({:?})", xs),
            FloatArray(xs) => write!(f, "FloatArray({:?})", xs),
            BoolArray(xs) => write!(f, "BoolArray({:?})", xs),
            Int2DArray(rows) => write!(f, "Int2DArray({:?})", rows),
            Float2DArray(rows) => write!(f, "Float2DArray({:?})", rows),
            KeyedArray(map) => write!(f, "KeyedArray({:?})", map),
            MixedArray(items) => write!(f, "MixedArray({:?})", items),
            Function(_) => write!(f, "Function(...)"),
            Stream(_) => write!(f, "Stream(...)"),
            InkIterator(_) => write!(f, "InkIterator(...)"),
            InkTransform(_) => write!(f, "InkTransform(...)"),
            Tensor(tv) => write!(f, "Tensor({:?})", tv),
            Ref(arc_mutex) => write!(f, "Ref({:?})", arc_mutex.lock().unwrap()),
        }
    }
}
