// 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 DenseSlotEntry {
  Vacant(Int?, Int)
  Detached(Int)
  Occupied(Int, Int)
}

///|
pub struct DenseSlotMap[K, V] {
  slots : Array[DenseSlotEntry]
  dense_keys : Array[K]
  dense_values : Array[V]
  mut free_head : Int?
}

///|
pub fn[K, V] DenseSlotMap::new(capacity? : Int = 0) -> DenseSlotMap[K, V] {
  let slots : Array[DenseSlotEntry] = []
  let dense_keys : Array[K] = []
  let dense_values : Array[V] = []
  if capacity > 0 {
    slots.reserve_capacity(capacity)
    dense_keys.reserve_capacity(capacity)
    dense_values.reserve_capacity(capacity)
  }
  { slots, dense_keys, dense_values, free_head: None }
}

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

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

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

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

///|
pub fn[K, V] DenseSlotMap::is_empty(self : DenseSlotMap[K, V]) -> Bool {
  self.dense_values.is_empty()
}

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

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

///|
pub fn[K : Key, V] DenseSlotMap::contains_key(
  self : DenseSlotMap[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) {
    DenseSlotEntry::Occupied(version, _) => version == key.version()
    _ => false
  }
}

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

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

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

///|
fn[K : Key, V] DenseSlotMap::remove_dense_index(
  self : DenseSlotMap[K, V],
  dense_index : Int,
) -> V {
  let last_index = self.dense_values.length() - 1
  let removed_value = self.dense_values.unsafe_get(dense_index)
  if dense_index != last_index {
    let moved_key = self.dense_keys.unsafe_get(last_index)
    let moved_value = self.dense_values.unsafe_get(last_index)
    self.dense_keys.unsafe_set(dense_index, moved_key)
    self.dense_values.unsafe_set(dense_index, moved_value)
    let moved_slot = moved_key.index()
    match self.slots.unsafe_get(moved_slot) {
      DenseSlotEntry::Occupied(version, _) =>
        self.slots.unsafe_set(
          moved_slot,
          DenseSlotEntry::Occupied(version, dense_index),
        )
      _ => panic()
    }
  }
  ignore(self.dense_keys.pop())
  ignore(self.dense_values.pop())
  removed_value
}

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

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

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

///|
pub fn[K : Key, V] DenseSlotMap::get(self : DenseSlotMap[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) {
    DenseSlotEntry::Occupied(version, dense_index) if version == key.version() =>
      Some(self.dense_values.unsafe_get(dense_index))
    _ => None
  }
}

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

///|
pub fn[K : Key, V] DenseSlotMap::replace(
  self : DenseSlotMap[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) {
    DenseSlotEntry::Occupied(version, dense_index) if version == key.version() => {
      self.dense_values.unsafe_set(dense_index, value)
      true
    }
    _ => false
  }
}

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

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

///|
pub fn[K : Key, V] DenseSlotMap::retain(
  self : DenseSlotMap[K, V],
  predicate : (K, V) -> Bool,
) -> Unit {
  let len = self.dense_values.length()
  let mut write = 0
  for read in 0..
        if predicate(key, value) {
          if write != read {
            self.dense_keys.unsafe_set(write, key)
            self.dense_values.unsafe_set(write, value)
            self.slots.unsafe_set(
              slot_index,
              DenseSlotEntry::Occupied(version, write),
            )
          }
          write += 1
        } else {
          self.slots.unsafe_set(
            slot_index,
            DenseSlotEntry::Vacant(self.free_head, version + 1),
          )
          self.free_head = Some(slot_index)
        }
      _ => panic()
    }
  }
  self.dense_keys.truncate(write)
  self.dense_values.truncate(write)
}

///|
pub fn[K : Key, V] DenseSlotMap::clear(self : DenseSlotMap[K, V]) -> Unit {
  let len = self.dense_keys.length()
  for i in 0.. {
        self.slots.unsafe_set(
          slot_index,
          DenseSlotEntry::Vacant(self.free_head, version + 1),
        )
        self.free_head = Some(slot_index)
      }
      _ => panic()
    }
  }
  self.dense_keys.clear()
  self.dense_values.clear()
}

///|
pub fn[K : Key, V] DenseSlotMap::drain(
  self : DenseSlotMap[K, V],
) -> Array[(K, V)] {
  let len = self.length()
  let items : Array[(K, V)] = Array::new(capacity=len)
  for i in 0.. {
        self.slots.unsafe_set(
          slot_index,
          DenseSlotEntry::Vacant(self.free_head, version + 1),
        )
        self.free_head = Some(slot_index)
      }
      _ => panic()
    }
  }
  self.dense_keys.clear()
  self.dense_values.clear()
  items
}

///|
pub fn[K, V] DenseSlotMap::to_array(self : DenseSlotMap[K, V]) -> Array[(K, V)] {
  Array::makei(self.length(), fn(i) {
    (self.dense_keys.unsafe_get(i), self.dense_values.unsafe_get(i))
  })
}

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

///|
#alias(iterator2, deprecated)
pub fn[K, V] DenseSlotMap::iter2(self : DenseSlotMap[K, V]) -> Iter2[K, V] {
  let mut index = 0
  let len = self.length()
  Iter2::new(fn() {
    guard index < len else { None }
    let item = (
      self.dense_keys.unsafe_get(index),
      self.dense_values.unsafe_get(index),
    )
    index += 1
    Some(item)
  })
}

///|
pub fn[K, V] DenseSlotMap::each(
  self : DenseSlotMap[K, V],
  visit : (K, V) -> Unit,
) -> Unit {
  let len = self.length()
  for i in 0.. Array[K] {
  self.dense_keys.copy()
}

///|
pub fn[K, V] DenseSlotMap::values(self : DenseSlotMap[K, V]) -> Array[V] {
  self.dense_values.copy()
}

///|
pub fn[K, V] DenseSlotMap::keys_as_slice(
  self : DenseSlotMap[K, V],
) -> ArrayView[K] {
  self.dense_keys[:]
}

///|
pub fn[K, V] DenseSlotMap::values_as_slice(
  self : DenseSlotMap[K, V],
) -> ArrayView[V] {
  self.dense_values[:]
}

///|
pub fn[K, V] DenseSlotMap::as_slices(
  self : DenseSlotMap[K, V],
) -> (ArrayView[K], ArrayView[V]) {
  (self.keys_as_slice(), self.values_as_slice())
}