tsrun 0.1.23

//! Map built-in methods

use crate::error::JsError;
use crate::interpreter::Interpreter;
use crate::prelude::{Box, Vec, index_map_new, index_map_with_capacity, vec};
use crate::value::{CheapClone, ExoticObject, Guarded, JsMapKey, JsValue, PropertyKey};

/// Initialize Map.prototype with get, set, has, delete, clear, forEach methods
pub fn init_map_prototype(interp: &mut Interpreter) {
    let proto = interp.map_prototype.clone();

    interp.register_method(&proto, "get", map_get, 1);
    interp.register_method(&proto, "set", map_set, 2);
    interp.register_method(&proto, "has", map_has, 1);
    interp.register_method(&proto, "delete", map_delete, 1);
    interp.register_method(&proto, "clear", map_clear, 0);
    interp.register_method(&proto, "forEach", map_foreach, 1);
    interp.register_method(&proto, "keys", map_keys, 0);
    interp.register_method(&proto, "values", map_values, 0);
    interp.register_method(&proto, "entries", map_entries, 0);

    // Symbol.iterator = Map.prototype.entries (Map iterates over [key, value] pairs by default)
    let well_known = interp.well_known_symbols;
    let iterator_symbol =
        crate::value::JsSymbol::new(well_known.iterator, Some(interp.intern("Symbol.iterator")));
    let iterator_key = crate::value::PropertyKey::Symbol(Box::new(iterator_symbol));
    let entries_fn = interp.create_native_function("[Symbol.iterator]", map_entries, 0);
    proto
        .borrow_mut()
        .set_property(iterator_key, JsValue::Object(entries_fn));
}

/// Create Map constructor and register it globally
pub fn init_map(interp: &mut Interpreter) {
    init_map_prototype(interp);

    let constructor = interp.create_native_function("Map", map_constructor, 0);
    interp.root_guard.guard(constructor.clone());

    // Add static methods to constructor
    interp.register_method(&constructor, "groupBy", map_group_by, 2);

    // Set constructor.prototype = Map.prototype
    let proto_key = PropertyKey::String(interp.intern("prototype"));
    constructor
        .borrow_mut()
        .set_property(proto_key, JsValue::Object(interp.map_prototype.clone()));

    // Set Map.prototype.constructor = Map
    let constructor_key = PropertyKey::String(interp.intern("constructor"));
    interp
        .map_prototype
        .borrow_mut()
        .set_property(constructor_key, JsValue::Object(constructor.clone()));

    // Add Symbol.species getter
    interp.register_species_getter(&constructor);

    // Register globally
    let map_key = PropertyKey::String(interp.intern("Map"));
    interp
        .global
        .borrow_mut()
        .set_property(map_key, JsValue::Object(constructor));
}

pub fn map_constructor(
    interp: &mut Interpreter,
    _this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let size_key = PropertyKey::String(interp.intern("size"));

    let guard = interp.heap.create_guard();
    let map_obj = interp.create_object(&guard);
    {
        let mut obj = map_obj.borrow_mut();
        obj.exotic = ExoticObject::Map {
            entries: index_map_new(),
        };
        obj.prototype = Some(interp.map_prototype.clone());
        obj.set_property(size_key, JsValue::Number(0.0));
    }

    // If an iterable is passed, add its entries
    // First collect all pairs from the array, then add them to the map
    if let Some(JsValue::Object(arr)) = args.first() {
        let pairs: Vec<(JsValue, JsValue)> = {
            let arr_ref = arr.borrow();
            let mut result = Vec::new();
            if let Some(elements) = arr_ref.array_elements() {
                for elem in elements {
                    if let JsValue::Object(pair_arr) = elem {
                        let pair_ref = pair_arr.borrow();
                        if pair_ref.is_array() {
                            let key = pair_ref
                                .get_property(&PropertyKey::Index(0))
                                .unwrap_or(JsValue::Undefined);
                            let value = pair_ref
                                .get_property(&PropertyKey::Index(1))
                                .unwrap_or(JsValue::Undefined);
                            result.push((key, value));
                        }
                    }
                }
            }
            result
        };

        // Now add all pairs to the map
        let size_key = PropertyKey::String(interp.intern("size"));
        let mut map = map_obj.borrow_mut();
        if let ExoticObject::Map { ref mut entries } = map.exotic {
            for (key, value) in pairs {
                entries.insert(JsMapKey(key), value);
            }
            let len = entries.len();
            map.set_property(size_key, JsValue::Number(len as f64));
        }
    }

    Ok(Guarded::with_guard(JsValue::Object(map_obj), guard))
}

pub fn map_get(
    _interp: &mut Interpreter,
    this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.get called on non-object",
        ));
    };

    let key = args.first().cloned().unwrap_or(JsValue::Undefined);
    let map = map_obj.borrow();

    if let ExoticObject::Map { ref entries } = map.exotic
        && let Some(value) = entries.get(&JsMapKey(key))
    {
        return Ok(Guarded::unguarded(value.clone()));
    }

    Ok(Guarded::unguarded(JsValue::Undefined))
}

pub fn map_set(
    interp: &mut Interpreter,
    this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this.clone() else {
        return Err(JsError::type_error(
            "Map.prototype.set called on non-object",
        ));
    };

    let size_key = PropertyKey::String(interp.intern("size"));

    let key = args.first().cloned().unwrap_or(JsValue::Undefined);
    let value = args.get(1).cloned().unwrap_or(JsValue::Undefined);

    let mut map = map_obj.borrow_mut();

    if let ExoticObject::Map { ref mut entries } = map.exotic {
        entries.insert(JsMapKey(key), value);
        let len = entries.len();
        map.set_property(size_key, JsValue::Number(len as f64));
    }

    drop(map);
    Ok(Guarded::unguarded(this)) // Return the map for chaining
}

pub fn map_has(
    _interp: &mut Interpreter,
    this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.has called on non-object",
        ));
    };

    let key = args.first().cloned().unwrap_or(JsValue::Undefined);
    let map = map_obj.borrow();

    if let ExoticObject::Map { ref entries } = map.exotic {
        return Ok(Guarded::unguarded(JsValue::Boolean(
            entries.contains_key(&JsMapKey(key)),
        )));
    }

    Ok(Guarded::unguarded(JsValue::Boolean(false)))
}

pub fn map_delete(
    interp: &mut Interpreter,
    this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.delete called on non-object",
        ));
    };

    let size_key = PropertyKey::String(interp.intern("size"));

    let key = args.first().cloned().unwrap_or(JsValue::Undefined);
    let mut map = map_obj.borrow_mut();

    if let ExoticObject::Map { ref mut entries } = map.exotic
        && entries.shift_remove(&JsMapKey(key)).is_some()
    {
        let len = entries.len();
        map.set_property(size_key, JsValue::Number(len as f64));
        return Ok(Guarded::unguarded(JsValue::Boolean(true)));
    }

    Ok(Guarded::unguarded(JsValue::Boolean(false)))
}

pub fn map_clear(
    interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.clear called on non-object",
        ));
    };

    let size_key = PropertyKey::String(interp.intern("size"));

    let mut map = map_obj.borrow_mut();

    if let ExoticObject::Map { ref mut entries } = map.exotic {
        entries.clear();
        map.set_property(size_key, JsValue::Number(0.0));
    }

    Ok(Guarded::unguarded(JsValue::Undefined))
}

pub fn map_foreach(
    interp: &mut Interpreter,
    this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this.clone() else {
        return Err(JsError::type_error(
            "Map.prototype.forEach called on non-object",
        ));
    };

    let callback = args.first().cloned().unwrap_or(JsValue::Undefined);
    let this_arg = args.get(1).cloned().unwrap_or(JsValue::Undefined);

    // Collect entries first to avoid borrow issues
    let entries: Vec<(JsValue, JsValue)>;
    {
        let map = map_obj.borrow();
        if let ExoticObject::Map { entries: ref e } = map.exotic {
            entries = e.iter().map(|(k, v)| (k.0.clone(), v.clone())).collect();
        } else {
            return Err(JsError::type_error(
                "Map.prototype.forEach called on non-Map",
            ));
        }
    }

    for (key, value) in entries {
        interp.call_function(
            callback.clone(),
            this_arg.clone(),
            &[value, key, this.clone()],
        )?;
    }

    Ok(Guarded::unguarded(JsValue::Undefined))
}

pub fn map_keys(
    interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.keys called on non-object",
        ));
    };

    let keys: Vec<JsValue>;
    {
        let map = map_obj.borrow();
        if let ExoticObject::Map { entries: ref e } = map.exotic {
            keys = e.keys().map(|k| k.0.clone()).collect();
        } else {
            return Err(JsError::type_error("Map.prototype.keys called on non-Map"));
        }
    }

    let guard = interp.heap.create_guard();
    let arr = interp.create_array_from(&guard, keys);
    Ok(Guarded::with_guard(JsValue::Object(arr), guard))
}

pub fn map_values(
    interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.values called on non-object",
        ));
    };

    let values: Vec<JsValue>;
    {
        let map = map_obj.borrow();
        if let ExoticObject::Map { entries: ref e } = map.exotic {
            values = e.iter().map(|(_, v)| v.clone()).collect();
        } else {
            return Err(JsError::type_error(
                "Map.prototype.values called on non-Map",
            ));
        }
    }

    let guard = interp.heap.create_guard();
    let arr = interp.create_array_from(&guard, values);
    Ok(Guarded::with_guard(JsValue::Object(arr), guard))
}

pub fn map_entries(
    interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(map_obj) = this else {
        return Err(JsError::type_error(
            "Map.prototype.entries called on non-object",
        ));
    };

    // Create an iterator object
    let guard = interp.heap.create_guard();
    guard.guard(map_obj.clone());

    // Collect entries upfront to avoid borrow issues during iteration
    let raw_entries: Vec<(JsValue, JsValue)>;
    {
        let map = map_obj.borrow();
        if let ExoticObject::Map { entries: ref e } = map.exotic {
            raw_entries = e.iter().map(|(k, v)| (k.0.clone(), v.clone())).collect();
        } else {
            return Err(JsError::type_error(
                "Map.prototype.entries called on non-Map",
            ));
        }
    }

    // Store entries as an array of [key, value] pairs in the iterator
    let mut entry_arrays = Vec::with_capacity(raw_entries.len());
    for (k, v) in raw_entries {
        let arr = interp.create_array_from(&guard, vec![k, v]);
        entry_arrays.push(JsValue::Object(arr));
    }
    let entries_arr = interp.create_array_from(&guard, entry_arrays);

    // Create iterator object
    let iter_obj = interp.create_object_raw(&guard);
    let entries_key = interp.property_key("__entries__");
    let index_key = interp.property_key("__index__");
    let next_key = interp.property_key("next");

    iter_obj
        .borrow_mut()
        .set_property(entries_key, JsValue::Object(entries_arr));
    iter_obj
        .borrow_mut()
        .set_property(index_key, JsValue::Number(0.0));

    // Add next() method
    let next_fn = interp.create_native_function("next", map_iterator_next, 0);
    guard.guard(next_fn.cheap_clone());
    iter_obj
        .borrow_mut()
        .set_property(next_key, JsValue::Object(next_fn));

    // Add Symbol.iterator that returns the iterator itself (iterator protocol)
    let well_known = interp.well_known_symbols;
    let iterator_symbol =
        crate::value::JsSymbol::new(well_known.iterator, Some(interp.intern("Symbol.iterator")));
    let iterator_key = crate::value::PropertyKey::Symbol(Box::new(iterator_symbol));
    let self_iterator_fn = interp.create_native_function("[Symbol.iterator]", map_iterator_self, 0);
    guard.guard(self_iterator_fn.cheap_clone());
    iter_obj
        .borrow_mut()
        .set_property(iterator_key, JsValue::Object(self_iterator_fn));

    Ok(Guarded::with_guard(JsValue::Object(iter_obj), guard))
}

/// Iterator Symbol.iterator function - returns the iterator itself
fn map_iterator_self(
    _interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    // Iterator's Symbol.iterator returns itself
    Ok(Guarded::unguarded(this))
}

/// Iterator next() function for Map entries/keys/values iterators
fn map_iterator_next(
    interp: &mut Interpreter,
    this: JsValue,
    _args: &[JsValue],
) -> Result<Guarded, JsError> {
    let JsValue::Object(iter_obj) = this else {
        return Err(JsError::type_error("next called on non-object"));
    };

    let entries_key = interp.property_key("__entries__");
    let index_key = interp.property_key("__index__");
    let value_key = interp.property_key("value");
    let done_key = interp.property_key("done");

    // Get the entries array and current index
    let entries_val = iter_obj.borrow().get_property(&entries_key);
    let index_val = iter_obj.borrow().get_property(&index_key);

    let Some(JsValue::Object(entries_arr)) = entries_val else {
        return Err(JsError::type_error("Invalid map iterator"));
    };

    let index = match index_val {
        Some(JsValue::Number(n)) => n as u32,
        _ => 0,
    };

    // Get length of entries array
    let length = entries_arr.borrow().array_length().unwrap_or(0);

    if index >= length {
        // Done
        let guard = interp.heap.create_guard();
        let result = interp.create_object_raw(&guard);
        result
            .borrow_mut()
            .set_property(value_key, JsValue::Undefined);
        result
            .borrow_mut()
            .set_property(done_key, JsValue::Boolean(true));
        Ok(Guarded::with_guard(JsValue::Object(result), guard))
    } else {
        // Get the entry at current index
        let value = entries_arr
            .borrow()
            .get_property(&PropertyKey::Index(index))
            .unwrap_or(JsValue::Undefined);

        // Increment index
        let index_key = interp.property_key("__index__");
        iter_obj
            .borrow_mut()
            .set_property(index_key, JsValue::Number((index + 1) as f64));

        let guard = interp.heap.create_guard();
        let result = interp.create_object_raw(&guard);
        let value_key = interp.property_key("value");
        let done_key = interp.property_key("done");
        result.borrow_mut().set_property(value_key, value);
        result
            .borrow_mut()
            .set_property(done_key, JsValue::Boolean(false));
        Ok(Guarded::with_guard(JsValue::Object(result), guard))
    }
}

/// Map.groupBy(items, callbackFn)
/// Groups elements of an iterable using a callback function.
/// Returns a Map where keys are group values and values are arrays.
/// Unlike Object.groupBy, keys can be any value (not just strings).
pub fn map_group_by(
    interp: &mut Interpreter,
    _this: JsValue,
    args: &[JsValue],
) -> Result<Guarded, JsError> {
    let items = args.first().cloned().unwrap_or(JsValue::Undefined);
    let callback = args.get(1).cloned().unwrap_or(JsValue::Undefined);

    // Items must be iterable - for now we support arrays
    let JsValue::Object(items_ref) = items else {
        return Err(JsError::type_error("Map.groupBy requires an iterable"));
    };

    // Guard the inputs
    let guard = interp.heap.create_guard();
    guard.guard(items_ref.clone());
    if let JsValue::Object(cb_obj) = &callback {
        guard.guard(cb_obj.clone());
    }

    // Get array elements
    let elements: Vec<JsValue> = {
        let items_borrowed = items_ref.borrow();
        if let Some(elems) = items_borrowed.array_elements() {
            elems.to_vec()
        } else {
            return Err(JsError::type_error(
                "Map.groupBy requires an array-like object",
            ));
        }
    };

    // Create a new Map for the result
    let size_key = PropertyKey::String(interp.intern("size"));
    let map_obj = interp.create_object(&guard);
    {
        let mut obj = map_obj.borrow_mut();
        obj.exotic = ExoticObject::Map {
            entries: index_map_new(),
        };
        obj.prototype = Some(interp.map_prototype.clone());
        obj.set_property(size_key.clone(), JsValue::Number(0.0));
    }

    // Track groups using IndexMap for O(1) lookup with SameValueZero semantics
    let mut groups = index_map_new::<JsMapKey, Vec<JsValue>>();

    // Iterate and group
    for (index, item) in elements.into_iter().enumerate() {
        // Guard the item in case callback triggers GC
        if let JsValue::Object(item_obj) = &item {
            guard.guard(item_obj.clone());
        }

        // Call the callback with (item, index)
        let key_result = interp.call_function(
            callback.clone(),
            JsValue::Undefined,
            &[item.clone(), JsValue::Number(index as f64)],
        )?;

        let key = key_result.value;

        // Add to existing group or create new one
        groups.entry(JsMapKey(key)).or_default().push(item);
    }

    // Now build the Map from the groups
    // First, create all the arrays (which may trigger GC)
    let mut built_entries = index_map_with_capacity::<JsMapKey, JsValue>(groups.len());
    for (key, items) in groups {
        let arr = interp.create_array_from(&guard, items);
        built_entries.insert(key, JsValue::Object(arr));
    }

    // Then add entries to the map
    {
        let mut map = map_obj.borrow_mut();
        if let ExoticObject::Map { ref mut entries } = map.exotic {
            *entries = built_entries;
            let len = entries.len();
            map.set_property(size_key, JsValue::Number(len as f64));
        }
    }

    Ok(Guarded::with_guard(JsValue::Object(map_obj), guard))
}