import gleam/list
import gleam/option

type Entry(a) {
  Vacant(option.Option(Int))
  Occupied(a)
}

pub opaque type Slab(a) {
  Slab(inner: List(Entry(a)), next_id: option.Option(Int), length: Int)
}

pub type SlabError {
  BadIndex(index: Int)
}

pub fn new() -> Slab(a) {
  Slab(inner: [], next_id: option.None, length: 0)
}

pub fn insert(slab: Slab(a), value: a) -> #(Slab(a), Int) {
  case slab.next_id {
    option.Some(slot_index) -> {
      let index = slab.length - 1 - slot_index

      let #(left, right) =
        slab.inner
        |> list.split(index)

      let assert Ok(Vacant(next_id)) =
        right
        |> list.first

      let right = case right {
        [_, ..rest] -> rest
        [] -> []
      }

      let inner = list.append(left, [Occupied(value), ..right])

      #(Slab(..slab, inner:, next_id:), slot_index)
    }

    option.None -> {
      let index = list.length(slab.inner)
      #(
        Slab(
          ..slab,
          inner: [Occupied(value), ..slab.inner],
          length: slab.length + 1,
        ),
        index,
      )
    }
  }
}

fn insert_accum(
  slab: Slab(a),
  values: List(a),
  acc: List(Int),
) -> #(Slab(a), List(Int)) {
  case values {
    [v, ..rest] -> {
      let #(slab, index) = insert(slab, v)
      insert_accum(slab, rest, [index, ..acc])
    }
    [] -> #(slab, acc)
  }
}

pub fn insert_many(slab: Slab(a), values: List(a)) -> #(Slab(a), List(Int)) {
  insert_accum(slab, values, [])
}

pub fn remove(slab: Slab(a), index: Int) -> Result(#(Slab(a), a), SlabError) {
  case index >= 0 && index < slab.length {
    True -> {
      let #(left, right) =
        slab.inner
        |> list.split(slab.length - 1 - index)

      let assert Ok(entry) =
        right
        |> list.first

      let right = case right {
        [_, ..rest] -> rest
        [] -> []
      }

      let inner = list.append(left, [Vacant(slab.next_id), ..right])

      let assert Occupied(value) = entry
      Ok(#(Slab(..slab, inner:, next_id: option.Some(index)), value))
    }

    False -> Error(BadIndex(index))
  }
}

fn get_items(items: List(Entry(a)), acc: Int, index: Int) -> option.Option(a) {
  case items {
    [v, _] | [v] if acc == index -> {
      case v {
        Vacant(_) -> option.None
        Occupied(v) -> option.Some(v)
      }
    }

    [_, ..rest] -> get_items(rest, acc - 1, index)

    [] -> option.None
  }
}

pub fn get(slab: Slab(a), index: Int) -> option.Option(a) {
  get_items(slab.inner, slab.length - 1, index)
}