use chrono::{TimeZone, Utc};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt;
use std::rc::Rc;
use std::sync::Mutex; // Still needed for global keyword registry
use std::sync::OnceLock;

/// Simple source location
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceLocation {
    pub line: u32,
    pub col: u32,
}

/// Shared identity for nodes. The identity is defined by the object instance (unique ID).
/// SourceLocation is kept as optional metadata.
#[derive(Debug, Clone)]
pub struct NodeIdentity {
    pub id: u64,
    pub location: Option<SourceLocation>,
}

impl PartialEq for NodeIdentity {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for NodeIdentity {}

impl std::hash::Hash for NodeIdentity {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.id.hash(state);
    }
}

pub type Identity = Rc<NodeIdentity>;

static NODE_ID_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);

impl NodeIdentity {
    pub fn next_id() -> u64 {
        NODE_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
    }

    /// Creates a new unique identity at the given location.
    pub fn new(location: SourceLocation) -> Identity {
        Rc::new(NodeIdentity {
            id: Self::next_id(),
            location: Some(location),
        })
    }

    /// Creates a new unique identity without location metadata (for synthetic nodes).
    pub fn anonymous() -> Identity {
        Rc::new(NodeIdentity {
            id: Self::next_id(),
            location: None,
        })
    }
}

/// Interned string identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Keyword(pub u32);

static KEYWORD_REGISTRY: OnceLock<Mutex<HashMap<String, u32>>> = OnceLock::new();
// We use String here instead of Arc<str> because benchmarks (e.g. record_optimizations.myc)
// showed a 4% performance regression with Arc due to atomic overhead during formatting.
// Since name() is mostly used for diagnostics and UI, the deep clone is acceptable.
static KEYWORD_REVERSE: OnceLock<Mutex<Vec<String>>> = OnceLock::new();

impl Keyword {
    pub fn intern(name: &str) -> Self {
        let mut reg = KEYWORD_REGISTRY
            .get_or_init(|| Mutex::new(HashMap::new()))
            .lock()
            .unwrap();
        if let Some(&id) = reg.get(name) {
            Keyword(id)
        } else {
            let mut rev = KEYWORD_REVERSE
                .get_or_init(|| Mutex::new(Vec::new()))
                .lock()
                .unwrap();
            let id = rev.len() as u32;
            reg.insert(name.to_string(), id);
            rev.push(name.to_string());
            Keyword(id)
        }
    }

    pub fn name(&self) -> String {
        let rev = KEYWORD_REVERSE
            .get_or_init(|| Mutex::new(Vec::new()))
            .lock()
            .unwrap();
        rev[self.0 as usize].clone()
    }
}

/// A shared schema for Records, providing O(1) field lookup via FMap optimization.
#[derive(Debug, Clone)]
pub struct RecordLayout {
    pub fields: Vec<(Keyword, StaticType)>,
    /// Optimization: maps (Keyword.idx - min_idx) to index in fields.
    fmap: Vec<i32>,
    min_idx: u32,
}

impl PartialEq for RecordLayout {
    fn eq(&self, other: &Self) -> bool {
        self.fields == other.fields
    }
}

impl Eq for RecordLayout {}

impl std::hash::Hash for RecordLayout {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.fields.hash(state);
    }
}

type LayoutMap = HashMap<Vec<(Keyword, StaticType)>, std::sync::Arc<RecordLayout>>;

static LAYOUT_REGISTRY: OnceLock<Mutex<LayoutMap>> = OnceLock::new();

impl RecordLayout {
    pub fn get_or_create(fields: Vec<(Keyword, StaticType)>) -> std::sync::Arc<Self> {
        let mut reg = LAYOUT_REGISTRY
            .get_or_init(|| Mutex::new(HashMap::new()))
            .lock()
            .unwrap();

        if let Some(layout) = reg.get(&fields) {
            return layout.clone();
        }

        if fields.is_empty() {
            let layout = std::sync::Arc::new(RecordLayout {
                fields: vec![],
                fmap: vec![],
                min_idx: 0,
            });
            reg.insert(fields, layout.clone());
            return layout;
        }

        let mut min_idx = u32::MAX;
        let mut max_idx = 0;
        for (k, _) in &fields {
            min_idx = min_idx.min(k.0);
            max_idx = max_idx.max(k.0);
        }

        let mut fmap = vec![-1; (max_idx - min_idx + 1) as usize];
        for (i, (k, _)) in fields.iter().enumerate() {
            fmap[(k.0 - min_idx) as usize] = i as i32;
        }

        let layout = std::sync::Arc::new(RecordLayout {
            fields: fields.clone(),
            fmap,
            min_idx,
        });

        reg.insert(fields, layout.clone());
        layout
    }

    #[inline(always)]
    pub fn index_of(&self, key: Keyword) -> Option<usize> {
        if key.0 < self.min_idx {
            return None;
        }
        let offset = (key.0 - self.min_idx) as usize;
        if offset >= self.fmap.len() {
            return None;
        }
        let idx = self.fmap[offset];
        if idx < 0 { None } else { Some(idx as usize) }
    }
}

/// Interface for custom objects (Closures, Series, Streams)
pub trait Object: fmt::Debug {
    fn type_name(&self) -> &'static str;
    fn as_any(&self) -> &dyn Any;

    /// Optional optimization: If this object behaves like a time series (indexable lookbacks),
    /// it can return a reference to its Series trait implementation.
    fn as_series(&self) -> Option<&dyn Series> {
        None
    }
}

/// Unified interface for all series types (Local RecordSeries, SeriesViews, and future SharedSeries).
pub trait Series {
    /// Gets the value at the specified lookback index (0 = newest).
    fn get_item(&self, index: usize) -> Option<Value>;

    /// Pushes a new value into the series.
    fn push_value(&self, value: Value);

    /// Returns the current number of elements in the buffer.
    fn len(&self) -> usize;

    /// Returns the total number of elements that have passed through this series.
    fn total_count(&self) -> u64;

    /// Returns true if the series is empty.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// A shared sequence of values, used by both Tuples and Records.
pub type ValueList = Rc<Vec<Value>>;

/// Purity level of an expression or function.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Purity {
    /// Has side effects (e.g. Set, Print)
    Impure,
    /// No side effects, but not deterministic (e.g. now(), random())
    SideEffectFree,
    /// No side effects and deterministic (e.g. sin(x), 1 + 2)
    Pure,
}

pub type NativeFn = dyn Fn(&[Value]) -> Value;
pub type PipeFn = dyn FnMut(&[Value]) -> Value;

/// A native host function with metadata.
pub struct NativeFunction {
    pub func: Rc<NativeFn>,
    pub purity: Purity,
}

impl fmt::Debug for NativeFunction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "<native fn, {:?}>", self.purity)
    }
}

/// Core data value in Myc Script (similar to TDataValue)
#[derive(Clone)]
pub enum Value {
    Void,
    Bool(bool),
    Int(i64),
    Float(f64),
    DateTime(i64),
    Text(Rc<str>),
    Keyword(Keyword),
    Tuple(ValueList),
    Record(std::sync::Arc<RecordLayout>, ValueList),
    FieldAccessor(Keyword),
    Function(Rc<NativeFunction>),
    Object(Rc<dyn Object>),   // For compiled Closures and other opaque types
    Cell(Rc<RefCell<Value>>), // Boxed value for captures
    TailCallRequest(Rc<dyn Object>, usize), // Internal: For TCO (target object + argc on stack)
}

impl PartialEq for Value {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Value::Void, Value::Void) => true,
            (Value::Bool(a), Value::Bool(b)) => a == b,
            (Value::Int(a), Value::Int(b)) => a == b,
            (Value::Float(a), Value::Float(b)) => a == b,
            (Value::DateTime(a), Value::DateTime(b)) => a == b,
            (Value::Text(a), Value::Text(b)) => a == b,
            (Value::Keyword(a), Value::Keyword(b)) => a == b,
            (Value::Tuple(a), Value::Tuple(b)) => a == b,
            (Value::Record(la, va), Value::Record(lb, vb)) => {
                std::sync::Arc::ptr_eq(la, lb) && va == vb
            }
            (Value::FieldAccessor(a), Value::FieldAccessor(b)) => a == b,
            (Value::Function(a), Value::Function(b)) => Rc::ptr_eq(a, b),
            (Value::Object(a), Value::Object(b)) => Rc::ptr_eq(a, b),
            (Value::Cell(a), Value::Cell(b)) => Rc::ptr_eq(a, b),
            (Value::TailCallRequest(oa, ca), Value::TailCallRequest(ob, cb)) => {
                Rc::ptr_eq(oa, ob) && ca == cb
            }
            _ => false,
        }
    }
}

impl PartialOrd for Value {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        match (self, other) {
            (Value::Int(a), Value::Int(b)) => a.partial_cmp(b),
            (Value::Float(a), Value::Float(b)) => a.partial_cmp(b),
            (Value::Int(a), Value::Float(b)) => (*a as f64).partial_cmp(b),
            (Value::Float(a), Value::Int(b)) => a.partial_cmp(&(*b as f64)),
            (Value::DateTime(a), Value::DateTime(b)) => a.partial_cmp(b),
            (Value::Text(a), Value::Text(b)) => a.partial_cmp(b),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Signature {
    pub params: StaticType,
    pub ret: StaticType,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum StaticType {
    Any,
    Void,
    Bool,
    Int,
    Float,
    DateTime,
    Text,
    Keyword,
    Optional(Box<StaticType>), // Represents T | Void (e.g. for filter pipes)
    List(Box<StaticType>),     // Legacy / Dynamic list
    Series(Box<StaticType>),   // Time series of a specific type
    Stream(Box<StaticType>),   // Reactive stream of a specific type
    Tuple(Vec<StaticType>),    // Heterogeneous fixed-size
    Vector(Box<StaticType>, usize), // Homogeneous fixed-size
    Matrix(Box<StaticType>, Vec<usize>), // Multi-dimensional homogeneous
    Record(std::sync::Arc<RecordLayout>),
    FieldAccessor(Keyword),
    Function(Box<Signature>),
    FunctionOverloads(Vec<Signature>),
    Object(&'static str),
    /// A diagnostic poison type, allowing type-checking to continue after an error.
    Error,
}

impl fmt::Display for StaticType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StaticType::Any => write!(f, "any"),
            StaticType::Void => write!(f, "void"),
            StaticType::Bool => write!(f, "bool"),
            StaticType::Int => write!(f, "int"),
            StaticType::Float => write!(f, "float"),
            StaticType::DateTime => write!(f, "datetime"),
            StaticType::Text => write!(f, "text"),
            StaticType::Keyword => write!(f, "keyword"),
            StaticType::Optional(inner) => write!(f, "optional({})", inner),
            StaticType::List(inner) => write!(f, "list({})", inner),
            StaticType::Series(inner) => write!(f, "series<{}>", inner),
            StaticType::Stream(inner) => write!(f, "stream<{}>", inner),
            StaticType::Tuple(elements) => {
                write!(f, "[")?;
                for (i, el) in elements.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", el)?;
                }
                write!(f, "]")
            }
            StaticType::Vector(inner, len) => write!(f, "vector<{}, {}>", inner, len),
            StaticType::Matrix(inner, shape) => {
                write!(f, "matrix<{}, [", inner)?;
                for (i, s) in shape.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", s)?;
                }
                write!(f, "]>")
            }
            StaticType::Record(layout) => {
                write!(f, "{{")?;
                for (i, (k, v)) in layout.fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, ":{} {}", k.name(), v)?;
                }
                write!(f, "}}")
            }
            StaticType::FieldAccessor(k) => write!(f, ".{}", k.name()),
            StaticType::Function(sig) => {
                write!(f, "fn({}) -> {}", sig.params, sig.ret)
            }
            StaticType::FunctionOverloads(sigs) => {
                write!(f, "overloads({} variants)", sigs.len())
            }
            StaticType::Object(name) => write!(f, "{}", name),
            StaticType::Error => write!(f, "<error>"),
        }
    }
}

impl StaticType {
    /// Returns true if `other` can be assigned to a location of type `self`.
    pub fn is_assignable_from(&self, other: &StaticType) -> bool {
        if self == other
            || matches!(self, StaticType::Any)
            || matches!(other, StaticType::Any)
            || matches!(self, StaticType::Error)
            || matches!(other, StaticType::Error)
        {
            return true;
        }

        match (self, other) {
            // Implicit Coercions
            (StaticType::Float, StaticType::Int) => true,
            (StaticType::Int, StaticType::Bool) => true,
            (StaticType::Int, StaticType::DateTime) => true,
            (StaticType::DateTime, StaticType::Int) => true,
            
            // Optional(T) is assignable from T or Void
            (StaticType::Optional(inner), other) => {
                matches!(other, StaticType::Void) || inner.is_assignable_from(other)
            }
            // A Vector is a Tuple
            (StaticType::Tuple(elements), StaticType::Vector(inner, len)) => {
                if elements.len() != *len {
                    return false;
                }
                elements.iter().all(|e| e.is_assignable_from(inner))
            }
            // Tuple to Tuple (Element-wise)
            (StaticType::Tuple(elements), StaticType::Tuple(other_elements)) => {
                if elements.len() != other_elements.len() {
                    return false;
                }
                elements
                    .iter()
                    .zip(other_elements.iter())
                    .all(|(e, o)| e.is_assignable_from(o))
            }
            // A Matrix is a Vector (of Vectors/Matrices)
            (StaticType::Vector(inner, len), StaticType::Matrix(m_inner, shape)) => {
                if shape.is_empty() || shape[0] != *len {
                    return false;
                }
                if shape.len() == 1 {
                    inner.is_assignable_from(m_inner)
                } else {
                    // It's a matrix of higher dimension, so inner must be assignable from a sub-matrix
                    let sub_shape = shape[1..].to_vec();
                    inner.is_assignable_from(&StaticType::Matrix(m_inner.clone(), sub_shape))
                }
            }
            // Records are assignable if their layouts match (Structural identity via interning)
            (StaticType::Record(a), StaticType::Record(b)) => std::sync::Arc::ptr_eq(a, b),
            // Series are assignable if their inner types are assignable
            (StaticType::Series(inner_a), StaticType::Series(inner_b)) => {
                inner_a.is_assignable_from(inner_b)
            }
            // Streams are assignable if their inner types are assignable
            (StaticType::Stream(inner_a), StaticType::Stream(inner_b)) => {
                inner_a.is_assignable_from(inner_b)
            }
            _ => false,
        }
    }

    /// Tries to resolve a call with the given argument type (usually a Tuple) and returns the return type.
    pub fn resolve_call(&self, args_ty: &StaticType) -> Option<StaticType> {
        match self {
            StaticType::Any => Some(StaticType::Any),
            StaticType::FieldAccessor(k) => {
                // (.name { :name val })
                // The args_ty should be a Tuple of [Record]
                let rec_ty = if let StaticType::Tuple(t) = args_ty {
                    t.first()?
                } else {
                    args_ty
                };

                match rec_ty {
                    StaticType::Record(layout) => {
                        if let Some(idx) = layout.index_of(*k) {
                            return Some(layout.fields[idx].1.clone());
                        }
                    }
                    StaticType::Series(inner) => {
                        if let StaticType::Record(layout) = inner.as_ref()
                            && let Some(idx) = layout.index_of(*k)
                        {
                            return Some(StaticType::Series(Box::new(layout.fields[idx].1.clone())));
                        }
                    }
                    StaticType::Stream(inner) => {
                        if let StaticType::Record(layout) = inner.as_ref()
                            && let Some(idx) = layout.index_of(*k)
                        {
                            return Some(StaticType::Stream(Box::new(layout.fields[idx].1.clone())));
                        }
                    }
                    StaticType::Any => return Some(StaticType::Any),
                    _ => {}
                }
                None
            }
            StaticType::Function(sig) => {
                if sig.params.is_assignable_from(args_ty) {
                    Some(sig.ret.clone())
                } else {
                    None
                }
            }
            StaticType::FunctionOverloads(sigs) => sigs
                .iter()
                .find(|sig| sig.params == *args_ty) // 1. Try exact match first
                .or_else(|| sigs.iter().find(|sig| sig.params.is_assignable_from(args_ty))) // 2. Fallback to implicit coercion
                .map(|sig| sig.ret.clone()),
            _ => None,
        }
    }

    /// Returns true if this type and all its recursive elements are scalars (Int, Float, etc.)
    pub fn is_scalar_pure(&self) -> bool {
        match self {
            StaticType::Int | StaticType::Float | StaticType::Bool | StaticType::DateTime => true,
            StaticType::Tuple(elements) => elements.iter().all(|e| e.is_scalar_pure()),
            StaticType::Vector(inner, _) => inner.is_scalar_pure(),
            StaticType::Matrix(inner, _) => inner.is_scalar_pure(),
            _ => false,
        }
    }
}

impl Value {
    pub fn get_field(&self, key: Keyword) -> Option<Value> {
        match self {
            Value::Record(layout, values) => {
                layout.index_of(key).map(|idx| values[idx].clone())
            }
            _ => None
        }
    }

    pub fn as_int(&self) -> Option<i64> {
        match self {
            Value::Int(i) => Some(*i),
            Value::DateTime(d) => Some(*d),
            Value::Bool(b) => Some(if *b { 1 } else { 0 }),
            _ => None,
        }
    }

    pub fn as_float(&self) -> Option<f64> {
        match self {
            Value::Float(f) => Some(*f),
            Value::Int(i) => Some(*i as f64),
            _ => None,
        }
    }

    pub fn is_truthy(&self) -> bool {
        match self {
            Value::Void => false,
            Value::Bool(b) => *b,
            Value::Cell(c) => c.borrow().is_truthy(),
            _ => true,
        }
    }

    /// Returns the underlying values as a slice if this is a Tuple.
    pub fn as_slice(&self) -> Option<&[Value]> {
        match self {
            Value::Tuple(v) => Some(v),
            _ => None,
        }
    }

    pub fn make_tuple(values: Vec<Value>) -> Self {
        Value::Tuple(Rc::new(values))
    }

    pub fn make_record(keys: Vec<Keyword>, values: Vec<Value>) -> Self {
        let fields: Vec<_> = keys
            .into_iter()
            .zip(values.iter().map(|v| v.static_type()))
            .collect();
        let layout = RecordLayout::get_or_create(fields);
        Value::Record(layout, Rc::new(values))
    }

    pub fn make_function(purity: Purity, func: impl Fn(&[Value]) -> Value + 'static) -> Self {
        Value::Function(Rc::new(NativeFunction {
            func: Rc::new(func),
            purity,
        }))
    }

    pub fn static_type(&self) -> StaticType {
        match self {
            Value::Void => StaticType::Void,
            Value::Bool(_) => StaticType::Bool,
            Value::Int(_) => StaticType::Int,
            Value::Float(_) => StaticType::Float,
            Value::DateTime(_) => StaticType::DateTime,
            Value::Text(_) => StaticType::Text,
            Value::Keyword(_) => StaticType::Keyword,
            Value::Tuple(values) => {
                if values.is_empty() {
                    return StaticType::Vector(Box::new(StaticType::Any), 0);
                }

                let element_types: Vec<_> = values.iter().map(|v| v.static_type()).collect();

                // Check for Homogeneity (Vector)
                let first_ty = &element_types[0];
                let all_same = element_types.iter().all(|t| t == first_ty);

                if all_same {
                    match first_ty {
                        StaticType::Vector(inner, len) => {
                            // Possible Matrix
                            StaticType::Matrix(inner.clone(), vec![values.len(), *len])
                        }
                        StaticType::Matrix(inner, shape) => {
                            let mut new_shape = vec![values.len()];
                            new_shape.extend(shape);
                            StaticType::Matrix(inner.clone(), new_shape)
                        }
                        _ => StaticType::Vector(Box::new(first_ty.clone()), values.len()),
                    }
                } else {
                    StaticType::Tuple(element_types)
                }
            }
            Value::Record(layout, _) => StaticType::Record(layout.clone()),
            Value::FieldAccessor(k) => StaticType::FieldAccessor(*k),
            Value::Function(_) => StaticType::Any, // Dynamic function
            Value::Object(o) => StaticType::Object(o.type_name()),
            Value::Cell(c) => c.borrow().static_type(),
            Value::TailCallRequest(_, _) => StaticType::Any, // Internal state, but typable as Any
        }
    }
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Value::Void => write!(f, "void"),
            Value::Bool(b) => write!(f, "{}", b),
            Value::Int(i) => write!(f, "{}", i),
            Value::Float(fl) => write!(f, "{}", fl),
            Value::DateTime(ts) => match Utc.timestamp_millis_opt(*ts) {
                chrono::LocalResult::Single(dt) => {
                    write!(f, "#{}#", dt.format("%Y-%m-%d %H:%M:%S"))
                }
                _ => write!(f, "#timestamp({})#", ts),
            },
            Value::Text(t) => write!(f, "\"{}\"", t),
            Value::Keyword(k) => write!(f, ":{}", k.name()),
            Value::Tuple(values) => {
                write!(f, "[")?;
                for (i, val) in values.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", val)?;
                }
                write!(f, "]")
            }
            Value::Record(layout, values) => {
                write!(f, "{{")?;
                for i in 0..values.len() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, ":{} {}", layout.fields[i].0.name(), values[i])?;
                }
                write!(f, "}}")
            }
            Value::FieldAccessor(k) => write!(f, ".{}", k.name()),
            Value::Function(f_meta) => write!(f, "<native fn, {:?}>", f_meta.purity),
            Value::Object(o) => {
                if let Some(val_series) =
                    o.as_any()
                        .downcast_ref::<crate::ast::rtl::series::SharedValueSeries>()
                {
                    write!(
                        f,
                        "SharedValueSeries[len: {}]",
                        val_series.buffer.borrow().len()
                    )
                } else {
                    write!(f, "<{}>", o.type_name())
                }
            }
            Value::Cell(c) => write!(f, "{}", c.borrow()),
            Value::TailCallRequest(_, _) => write!(f, "<tail call request>"),
        }
    }
}

impl fmt::Debug for Value {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self, f)
    }
}