Functional State Management

// Simple State monad implementation
case class State[S, A](run: S => (A, S)):
  def map[B](f: A => B): State[S, B] =
    State { s =>
      val (a, s1) = run(s)
      (f(a), s1)
    }

  def flatMap[B](f: A => State[S, B]): State[S, B] =
    State { s =>
      val (a, s1) = run(s)
      f(a).run(s1)
    }

object State:
  def get[S]: State[S, S] = State(s => (s, s))
  def set[S](s: S): State[S, Unit] = State(_ => ((), s))
  def modify[S](f: S => S): State[S, Unit] = State(s => ((), f(s)))
  def pure[S, A](a: A): State[S, A] = State(s => (a, s))

// Stack operations
type Stack[A] = List[A]

def push[A](item: A): State[Stack[A], Unit] =
  State.modify(item :: _)

def pop[A]: State[Stack[A], Option[A]] =
  State { stack =>
    stack match
      case head :: tail => (Some(head), tail)
      case Nil          => (None, Nil)
  }

def peek[A]: State[Stack[A], Option[A]] =
  State(stack => (stack.headOption, stack))

def stackSize[A]: State[Stack[A], Int] =
  State(stack => (stack.size, stack))

// Random number generator
case class RNG(seed: Long):
  def next: (Int, RNG) =
    val newSeed = (seed * 6364136223846793005L + 1442695040888963407L) & Long.MaxValue
    val n = (newSeed >>> 16).toInt
    (n, RNG(newSeed))

def nextInt: State[RNG, Int] = State(_.next)

def nextIntBound(bound: Int): State[RNG, Int] =
  nextInt.map(n => (n & Int.MaxValue) % bound)

def nextBool: State[RNG, Boolean] =
  nextInt.map(_ % 2 == 0)

def rollDice: State[RNG, Int] =
  nextIntBound(6).map(_ + 1)

// Game state
case class GameState(score: Int, lives: Int, level: Int)

def addScore(points: Int): State[GameState, Unit] =
  State.modify(g => g.copy(score = g.score + points))

def loseLife: State[GameState, Boolean] =
  for
    _ <- State.modify[GameState](g => g.copy(lives = g.lives - 1))
    g <- State.get[GameState]
  yield g.lives > 0

def nextLevel: State[GameState, Int] =
  for
    _ <- State.modify[GameState](g => g.copy(level = g.level + 1))
    g <- State.get[GameState]
  yield g.level

@main def run(): Unit =
  // Stack
  val stackProgram = for
    _ <- push(1)
    _ <- push(2)
    _ <- push(3)
    top <- pop[Int]
    size <- stackSize[Int]
  yield (top, size)

  val (result, finalStack) = stackProgram.run(Nil)
  println(s"Popped: ${result._1}, Size: ${result._2}, Stack: $finalStack")

  // Random
  val diceRolls = for
    d1 <- rollDice
    d2 <- rollDice
    d3 <- rollDice
  yield List(d1, d2, d3)

  val (rolls, _) = diceRolls.run(RNG(42))
  println(s"Dice rolls: $rolls")

  // Game
  val game = for
    _ <- addScore(100)
    _ <- addScore(50)
    alive <- loseLife
    lvl <- nextLevel
    g <- State.get[GameState]
  yield (alive, lvl, g)

  val (gameResult, finalState) = game.run(GameState(0, 3, 1))
  println(s"Alive: ${gameResult._1}, Level: ${gameResult._2}")
  println(s"Final: $finalState")