import scala.reflect.runtime.universe._

package object ecs {

case class EntityRef(id: Long) extends AnyVal {
  def set[T : TypeTag](newValue: T): Set[T] =
    Set[T](typeTag[T], this, newValue)
}

object EntityRef {
  private var nextId: Long = 0

  def next(): EntityRef = {
    val ref = new EntityRef(nextId)
    nextId += 1
    ref
  }
}

sealed trait ComponentChange[+T]

case class Set[+T](
  tt: TypeTag[_],
  entityRef: EntityRef,
  newValue: T
) extends ComponentChange[T]

case class Delete[+T](
  tt: TypeTag[_],
  entityRef: EntityRef,
) extends ComponentChange[T]

class ComponentStore(
  private val components: Map[TypeTag[_], Map[EntityRef, Any]] = Map.empty
) {
  def add[T : TypeTag](ref: EntityRef, component: T): ComponentStore = {
    val tt = typeTag[T]
    val valuesByEntityRef = components.getOrElse(tt, Map.empty)
    new ComponentStore(
      components.updated(tt, valuesByEntityRef.updated(ref, component))
    )
  }

  def get[T : TypeTag]: Seq[(EntityRef, T)] =
    components.getOrElse(typeTag[T], Map.empty)
      .toSeq
      .map { case (ref, component) => (ref, component.asInstanceOf[T]) }

  def get[T : TypeTag, U: TypeTag]: Seq[(EntityRef, T, U)] = {
    val ts = components.getOrElse(typeTag[T], Map.empty)
    val us = components.getOrElse(typeTag[U], Map.empty)
    (ts.keys ++ us.keys)
      .map { ref =>
        for {
          t <- ts.get(ref).map(_.asInstanceOf[T])
          u <- us.get(ref).map(_.asInstanceOf[U])
        } yield (ref, t, u)
      }
      .collect { case Some(x) => x }
      .toSeq
  }

  def applyChange(change: ComponentChange[Any]): ComponentStore = {
    val newComponents = change match {
      case Set(tt, entityRef, newValue) =>
        val valuesByEntityRef = components.getOrElse(tt, Map.empty)
        components.updated(tt, valuesByEntityRef.updated(entityRef, newValue))
      case Delete(tt, entityRef) =>
        components.get(tt) match {
          case Some(valuesByEntityRef) =>
            components.updated(tt, valuesByEntityRef - entityRef) // TODO don't put it back if it's empty
          case None =>
            components
        }
    }
    new ComponentStore(newComponents)
  }

  def applySystem(system: System[Any]): ComponentStore = {
    val changes = system(this)
    changes.foldLeft(this)(_.applyChange(_))
  }

  override def toString: String = components.toString

}

type System[T] = ComponentStore => Seq[ComponentChange[T]]

object Test {

  case class Position(x: Int, y: Int)
  case object Player
  case object NextToSomeone

  val e1 = EntityRef.next()
  val e2 = EntityRef.next()

  val moveAllUp: System[Position] = { cs: ComponentStore =>
    for {
      (ref, Position(x, y)) <- cs.get[Position]
    } yield ref.set(Position(x, y + 1))
  }

  val movePlayerRight: System[Position] = { cs: ComponentStore =>
    for {
      (ref, Position(x, y), _) <- cs.get[Position, Player.type]
    } yield ref.set(Position(x + 1, y))
  }

  val markNextToSomeone: System[NextToSomeone.type] = { cs: ComponentStore =>
    for {
      // TODO everyone gets marked twice, it's not observable though
      (ref1, p1) <- cs.get[Position]
      (ref2, p2) <- cs.get[Position]
      if (Math.abs(p1.x - p2.x) == 1 || Math.abs(p1.y - p2.y) == 1)
      result <- Seq(ref1.set(NextToSomeone), ref2.set(NextToSomeone))
    } yield result
  }

  val cs0 = new ComponentStore()
    .add(e1, Position(0, 0))
    .add(e1, Player)
    .add(e2, Position(1, 1))

  val cs1 = cs0
    .applySystem(moveAllUp)
    .applySystem(movePlayerRight)
    .applySystem(markNextToSomeone)

  println(cs0)
  println(cs1)
}

}