///| A transaction-style model of KV-cache ownership. It does not store tensor

///|
/// values: it records the slot lifecycle an inference adapter must follow.
pub enum SlotState {
  Free
  Reserved(Int)
  Committed(Int)
} derive(Eq, Debug)

///|
pub enum SlotError {
  InvalidSlot(Int)
  NotFree(Int)
  InvalidTransaction(Int)
} derive(Eq, Debug)

///|
pub struct KvSlots {
  slots : Array[SlotState]
  mut next_transaction : Int
}

///|
pub fn KvSlots::new(capacity : Int) -> KvSlots {
  let slots : Array[SlotState] = []
  for _ in 0.. Int {
  self.slots.length()
}

///|
pub fn KvSlots::free_count(self : KvSlots) -> Int {
  let mut count = 0
  for state in self.slots {
    if state == Free {
      count = count + 1
    }
  }
  count
}

///|
/// Return the lifecycle state of one cache slot.
pub fn KvSlots::state_at(
  self : KvSlots,
  index : Int,
) -> Result[SlotState, SlotError] {
  if index < 0 || index >= self.slots.length() {
    Err(InvalidSlot(index))
  } else {
    Ok(self.slots[index])
  }
}

///|
pub fn KvSlots::reserve(
  self : KvSlots,
  count : Int,
) -> Result[(Int, Array[Int]), SlotError] {
  if count <= 0 || count > self.free_count() {
    return Err(InvalidTransaction(count))
  }
  let transaction = self.next_transaction
  self.next_transaction = self.next_transaction + 1
  let indices : Array[Int] = []
  for index in 0.. Result[Unit, SlotError] {
  if keep < 0 {
    return Err(InvalidTransaction(transaction))
  }
  let mut reserved = 0
  for state in self.slots {
    match state {
      Reserved(id) if id == transaction => reserved = reserved + 1
      _ => ()
    }
  }
  if reserved == 0 || keep > reserved {
    return Err(InvalidTransaction(transaction))
  }
  let mut encountered = 0
  for index in 0.. {
        if encountered < keep {
          self.slots[index] = Committed(transaction)
        } else {
          self.slots[index] = Free
        }
        encountered = encountered + 1
      }
      _ => ()
    }
  }
  Ok(())
}

///|
pub fn KvSlots::rollback(
  self : KvSlots,
  transaction : Int,
) -> Result[Unit, SlotError] {
  let mut found = false
  for index in 0.. {
        self.slots[index] = Free
        found = true
      }
      _ => ()
    }
  }
  if found {
    Ok(())
  } else {
    Err(InvalidTransaction(transaction))
  }
}

///| Release every committed slot owned by a completed decoding transaction.

///| A transaction has to be committed before it can be released; releasing a

///|
/// merely reserved or unknown transaction is a lifecycle error.
pub fn KvSlots::release_committed(
  self : KvSlots,
  transaction : Int,
) -> Result[Unit, SlotError] {
  let mut found = false
  for index in 0.. {
        self.slots[index] = Free
        found = true
      }
      _ => ()
    }
  }
  if found {
    Ok(())
  } else {
    Err(NotFree(transaction))
  }
}