// 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.

///|
pub struct SecondaryMap[K, V] {
  slots : Array[(Int, V)?]
  mut size : Int
  _key : K?
}

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

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

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

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

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

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

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

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

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

///|
pub fn[K : Key, V] SecondaryMap::insert(
  self : SecondaryMap[K, V],
  key : K,
  value : V,
) -> V? {
  if key.is_null() {
    return None
  }
  let index = key.index()
  if index < 0 {
    return None
  }
  while self.slots.length() <= index {
    self.slots.push(None)
  }
  match self.slots.unsafe_get(index) {
    Some((version, current)) if version == key.version() => {
      self.slots.unsafe_set(index, Some((version, value)))
      Some(current)
    }
    Some((version, _)) if version_is_older(key.version(), version) => None
    Some(_) => {
      self.slots.unsafe_set(index, Some((key.version(), value)))
      None
    }
    None => {
      self.slots.unsafe_set(index, Some((key.version(), value)))
      self.size += 1
      None
    }
  }
}

///|
#alias("_[_]=_")
pub fn[K : Key, V] SecondaryMap::set(
  self : SecondaryMap[K, V],
  key : K,
  value : V,
) -> Unit {
  ignore(self.insert(key, value))
}

///|
pub fn[K : Key, V] SecondaryMap::remove(
  self : SecondaryMap[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) {
    Some((version, value)) if version == key.version() => {
      self.slots.unsafe_set(index, None)
      self.size -= 1
      Some(value)
    }
    _ => None
  }
}

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

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

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

///|
pub fn[K : Key, V] SecondaryMap::get_or_default(
  self : SecondaryMap[K, V],
  key : K,
  default : V,
) -> V? {
  if key.is_null() {
    return None
  }
  let index = key.index()
  if index < 0 {
    return None
  }
  while self.slots.length() <= index {
    self.slots.push(None)
  }
  match self.slots.unsafe_get(index) {
    Some((version, value)) if version == key.version() => Some(value)
    Some((version, _)) if version_is_older(key.version(), version) => None
    Some(_) => {
      self.slots.unsafe_set(index, Some((key.version(), default)))
      Some(default)
    }
    None => {
      self.slots.unsafe_set(index, Some((key.version(), default)))
      self.size += 1
      Some(default)
    }
  }
}

///|
pub fn[K : Key, V] SecondaryMap::get_or_init(
  self : SecondaryMap[K, V],
  key : K,
  init : () -> V,
) -> V? {
  if key.is_null() {
    ignore(init())
    return None
  }
  let index = key.index()
  if index < 0 {
    ignore(init())
    return None
  }
  if index >= self.slots.length() {
    let value = init()
    while self.slots.length() <= index {
      self.slots.push(None)
    }
    self.slots.unsafe_set(index, Some((key.version(), value)))
    self.size += 1
    return Some(value)
  }
  match self.slots.unsafe_get(index) {
    Some((version, value)) if version == key.version() => Some(value)
    Some((version, _)) if version_is_older(key.version(), version) => {
      ignore(init())
      None
    }
    Some(_) => {
      let value = init()
      self.slots.unsafe_set(index, Some((key.version(), value)))
      Some(value)
    }
    None => {
      let value = init()
      self.slots.unsafe_set(index, Some((key.version(), value)))
      self.size += 1
      Some(value)
    }
  }
}

///|
pub fn[K : Key, V] SecondaryMap::retain(
  self : SecondaryMap[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, None)
          self.size -= 1
        }
      }
      None => ()
    }
  }
}

///|
pub fn[K, V] SecondaryMap::clear(self : SecondaryMap[K, V]) -> Unit {
  let len = self.slots.length()
  for i in 0.. 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, None)
      }
      None => ()
    }
  }
  self.size = 0
  items
}

///|
pub fn[K : Key, V] SecondaryMap::to_array(
  self : SecondaryMap[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))
      None => ()
    }
  }
  items
}

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

///|
#alias(iterator2, deprecated)
pub fn[K : Key, V] SecondaryMap::iter2(
  self : SecondaryMap[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) {
        Some((version, value)) =>
          next_item = Some((K::from_raw_parts(current, version), value))
        None => ()
      }
    }
    next_item
  })
}

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

///|
pub fn[K : Key, V] SecondaryMap::keys(self : SecondaryMap[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))
      None => ()
    }
  }
  keys
}

///|
pub fn[K, V] SecondaryMap::values(self : SecondaryMap[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)
      None => ()
    }
  }
  values
}