// Copyright 2025 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
enum SlotEntry[V] {
  Vacant(Int?, Int)
  Detached(Int)
  Occupied(Int, V)
}

///|
pub struct SlotMap[K, V] {
  slots : Array[SlotEntry[V]]
  mut free_head : Int?
  mut size : Int
  _key : K?
}

///|
pub fn[K, V] SlotMap::new(capacity? : Int = 0) -> SlotMap[K, V] {
  let slots : Array[SlotEntry[V]] = []
  if capacity > 0 {
    slots.reserve_capacity(capacity)
  }
  { slots, free_head: None, size: 0, _key: None }
}

///|
pub fn[K, V] SlotMap::default() -> SlotMap[K, V] {
  SlotMap::new()
}

///|
pub fn[K, V] SlotMap::with_capacity(capacity : Int) -> SlotMap[K, V] {
  SlotMap::new(capacity~)
}

///|
pub fn[K, V] SlotMap::length(self : SlotMap[K, V]) -> Int {
  self.size
}

///|
pub fn[K, V] SlotMap::len(self : SlotMap[K, V]) -> Int {
  self.length()
}

///|
pub fn[K, V] SlotMap::is_empty(self : SlotMap[K, V]) -> Bool {
  self.size == 0
}

///|
pub fn[K, V] SlotMap::capacity(self : SlotMap[K, V]) -> Int {
  self.slots.capacity()
}

///|
pub fn[K, V] SlotMap::reserve(self : SlotMap[K, V], additional : Int) -> Unit {
  if additional > 0 {
    self.slots.reserve_capacity(additional)
  }
}

///|
pub fn[K : Key, V] SlotMap::contains_key(self : SlotMap[K, V], key : K) -> Bool {
  if key.is_null() {
    return false
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    return false
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, _) => version == key.version()
    _ => false
  }
}

///|
pub fn[K : Key, V] SlotMap::contains(self : SlotMap[K, V], key : K) -> Bool {
  self.contains_key(key)
}

///|
pub fn[K : Key, V] SlotMap::insert(self : SlotMap[K, V], value : V) -> K {
  match self.free_head {
    Some(index) =>
      match self.slots.unsafe_get(index) {
        SlotEntry::Vacant(next_free, version) => {
          let occupied_version = version + 1
          let key = K::from_raw_parts(index, occupied_version)
          self.slots.unsafe_set(
            index,
            SlotEntry::Occupied(occupied_version, value),
          )
          self.free_head = next_free
          self.size += 1
          key
        }
        _ => panic()
      }
    None => {
      let index = self.slots.length()
      let key = K::from_raw_parts(index, 1)
      self.slots.push(SlotEntry::Occupied(1, value))
      self.size += 1
      key
    }
  }
}

///|
pub fn[K : Key, V] SlotMap::insert_with_key(
  self : SlotMap[K, V],
  create : (K) -> V,
) -> K {
  match self.free_head {
    Some(index) =>
      match self.slots.unsafe_get(index) {
        SlotEntry::Vacant(next_free, version) => {
          let occupied_version = version + 1
          let key = K::from_raw_parts(index, occupied_version)
          let value = create(key)
          self.slots.unsafe_set(
            index,
            SlotEntry::Occupied(occupied_version, value),
          )
          self.free_head = next_free
          self.size += 1
          key
        }
        _ => panic()
      }
    None => {
      let index = self.slots.length()
      let key = K::from_raw_parts(index, 1)
      let value = create(key)
      self.slots.push(SlotEntry::Occupied(1, value))
      self.size += 1
      key
    }
  }
}

///|
pub fn[K : Key, V] SlotMap::remove(self : SlotMap[K, V], key : K) -> V? {
  if key.is_null() {
    return None
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    return None
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, value) if version == key.version() => {
      self.slots.unsafe_set(
        index,
        SlotEntry::Vacant(self.free_head, version + 1),
      )
      self.free_head = Some(index)
      self.size -= 1
      Some(value)
    }
    _ => None
  }
}

///|
pub fn[K : Key, V] SlotMap::detach(self : SlotMap[K, V], key : K) -> V? {
  if key.is_null() {
    return None
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    return None
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, value) if version == key.version() => {
      self.slots.unsafe_set(index, SlotEntry::Detached(version + 1))
      self.size -= 1
      Some(value)
    }
    _ => None
  }
}

///|
pub fn[K : Key, V] SlotMap::reattach(
  self : SlotMap[K, V],
  detached_key : K,
  value : V,
) -> Unit {
  let index = detached_key.index()
  if index < 0 || index >= self.slots.length() {
    abort("key is not detached")
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Detached(version) if version == detached_key.version() + 1 => {
      self.slots.unsafe_set(
        index,
        SlotEntry::Occupied(detached_key.version(), value),
      )
      self.size += 1
    }
    _ => abort("key is not detached")
  }
}

///|
pub fn[K : Key, V] SlotMap::get(self : SlotMap[K, V], key : K) -> V? {
  if key.is_null() {
    return None
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    return None
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, value) if version == key.version() =>
      Some(value)
    _ => None
  }
}

///|
#alias("_[_]")
pub fn[K : Key, V] SlotMap::at(self : SlotMap[K, V], key : K) -> V {
  match self.get(key) {
    Some(value) => value
    None => abort("invalid slotmap key")
  }
}

///|
pub fn[K : Key, V] SlotMap::replace(
  self : SlotMap[K, V],
  key : K,
  value : V,
) -> Bool {
  if key.is_null() {
    return false
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    return false
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, _) if version == key.version() => {
      self.slots.unsafe_set(index, SlotEntry::Occupied(version, value))
      true
    }
    _ => false
  }
}

///|
#alias("_[_]=_")
pub fn[K : Key, V] SlotMap::set(
  self : SlotMap[K, V],
  key : K,
  value : V,
) -> Unit {
  if !self.replace(key, value) {
    abort("invalid slotmap key")
  }
}

///|
pub fn[K : Key, V] SlotMap::update(
  self : SlotMap[K, V],
  key : K,
  updater : (V?) -> V?,
) -> Unit {
  if key.is_null() {
    ignore(updater(None))
    return
  }
  let index = key.index()
  if index < 0 || index >= self.slots.length() {
    ignore(updater(None))
    return
  }
  match self.slots.unsafe_get(index) {
    SlotEntry::Occupied(version, current) if version == key.version() =>
      match updater(Some(current)) {
        Some(value) =>
          self.slots.unsafe_set(index, SlotEntry::Occupied(version, value))
        None => {
          self.slots.unsafe_set(
            index,
            SlotEntry::Vacant(self.free_head, version + 1),
          )
          self.free_head = Some(index)
          self.size -= 1
        }
      }
    SlotEntry::Detached(version) if version == key.version() + 1 =>
      match updater(None) {
        Some(value) => {
          self.slots.unsafe_set(
            index,
            SlotEntry::Occupied(key.version(), value),
          )
          self.size += 1
        }
        None => ()
      }
    _ => ignore(updater(None))
  }
}

///|
pub fn[K : Key, V] SlotMap::retain(
  self : SlotMap[K, V],
  predicate : (K, V) -> Bool,
) -> Unit {
  let len = self.slots.length()
  for i in 0.. {
        let key = K::from_raw_parts(i, version)
        if !predicate(key, value) {
          self.slots.unsafe_set(
            i,
            SlotEntry::Vacant(self.free_head, version + 1),
          )
          self.free_head = Some(i)
          self.size -= 1
        }
      }
      _ => ()
    }
  }
}

///|
pub fn[K, V] SlotMap::clear(self : SlotMap[K, V]) -> Unit {
  let len = self.slots.length()
  for i in 0.. {
        self.slots.unsafe_set(i, SlotEntry::Vacant(self.free_head, version + 1))
        self.free_head = Some(i)
      }
      _ => ()
    }
  }
  self.size = 0
}

///|
pub fn[K : Key, V] SlotMap::drain(self : SlotMap[K, V]) -> Array[(K, V)] {
  let items : Array[(K, V)] = Array::new(capacity=self.size)
  let len = self.slots.length()
  for i in 0.. {
        items.push((K::from_raw_parts(i, version), value))
        self.slots.unsafe_set(i, SlotEntry::Vacant(self.free_head, version + 1))
        self.free_head = Some(i)
      }
      _ => ()
    }
  }
  self.size = 0
  items
}

///|
pub fn[K : Key, V] SlotMap::to_array(self : SlotMap[K, V]) -> Array[(K, V)] {
  let items : Array[(K, V)] = Array::new(capacity=self.size)
  let len = self.slots.length()
  for i in 0..
        items.push((K::from_raw_parts(i, version), value))
      _ => ()
    }
  }
  items
}

///|
pub fn[K : Key, V] SlotMap::iter(self : SlotMap[K, V]) -> Array[(K, V)] {
  self.to_array()
}

///|
#alias(iterator2, deprecated)
pub fn[K : Key, V] SlotMap::iter2(self : SlotMap[K, V]) -> Iter2[K, V] {
  let mut index = 0
  let len = self.slots.length()
  Iter2::new(fn() {
    let mut next_item : (K, V)? = None
    while next_item is None && index < len {
      let current = index
      index += 1
      match self.slots.unsafe_get(current) {
        SlotEntry::Occupied(version, value) =>
          next_item = Some((K::from_raw_parts(current, version), value))
        _ => ()
      }
    }
    next_item
  })
}

///|
pub fn[K : Key, V] SlotMap::each(
  self : SlotMap[K, V],
  visit : (K, V) -> Unit,
) -> Unit {
  let len = self.slots.length()
  for i in 0..
        visit(K::from_raw_parts(i, version), value)
      _ => ()
    }
  }
}

///|
pub fn[K : Key, V] SlotMap::keys(self : SlotMap[K, V]) -> Array[K] {
  let keys : Array[K] = Array::new(capacity=self.size)
  let len = self.slots.length()
  for i in 0..
        keys.push(K::from_raw_parts(i, version))
      _ => ()
    }
  }
  keys
}

///|
pub fn[K, V] SlotMap::values(self : SlotMap[K, V]) -> Array[V] {
  let values : Array[V] = Array::new(capacity=self.size)
  let len = self.slots.length()
  for i in 0.. values.push(value)
      _ => ()
    }
  }
  values
}