jmonkey-test/src/main/scala/wow/doge/mygame/implicits/package.scala

package wow.doge.mygame

import scala.jdk.CollectionConverters._
import scala.reflect.ClassTag

import akka.actor.typed.ActorRef
import akka.actor.typed.Scheduler
import akka.util.Timeout
import com.jayfella.jme.jfx.JavaFxUI
import com.jme3.app.Application
import com.jme3.app.SimpleApplication
import com.jme3.app.state.AppState
import com.jme3.app.state.AppStateManager
import com.jme3.asset.AssetLocator
import com.jme3.asset.AssetManager
import com.jme3.bullet.BulletAppState
import com.jme3.bullet.PhysicsSpace
import com.jme3.bullet.PhysicsTickListener
import com.jme3.bullet.collision.PhysicsCollisionEvent
import com.jme3.bullet.collision.PhysicsCollisionListener
import com.jme3.bullet.control.BetterCharacterControl
import com.jme3.input.Action
import com.jme3.input.InputManager
import com.jme3.input.controls.ActionListener
import com.jme3.input.controls.AnalogListener
import com.jme3.input.controls.InputListener
import com.jme3.input.controls.Trigger
import com.jme3.light.Light
import com.jme3.math.Vector3f
import com.jme3.scene.CameraNode
import com.jme3.scene.Geometry
import com.jme3.scene.Node
import com.jme3.scene.SceneGraphVisitor
import com.jme3.scene.Spatial
import com.jme3.scene.control.CameraControl.ControlDirection
import com.jme3.scene.control.Control
import com.simsilica.es.EntityComponent
import com.simsilica.es.EntityData
import com.simsilica.es.EntityId
import enumeratum._
import monix.bio.Task
import monix.bio.UIO
import monix.execution.Ack
import monix.execution.Ack.Continue
import monix.execution.Ack.Stop
import monix.execution.Cancelable
import monix.execution.cancelables.SingleAssignCancelable
import monix.reactive.Observable
import monix.reactive.OverflowStrategy
import monix.reactive.observers.Subscriber
import wow.doge.mygame.math.ImVector3f
import wow.doge.mygame.state.MyBaseState
import com.jme3.material.Material

final case class ActionEvent(binding: Action, value: Boolean, tpf: Float)
final case class EnumActionEvent[T <: EnumEntry](
    binding: T,
    value: Boolean,
    tpf: Float
)

final case class AnalogEvent(binding: Action, value: Float, tpf: Float)
final case class EnumAnalogEvent[T <: EnumEntry](
    binding: T,
    value: Float,
    tpf: Float
)
final class PrePhysicsTickEvent(val space: PhysicsSpace) extends AnyVal
final class PhysicsTickEvent(val space: PhysicsSpace) extends AnyVal

package object implicits {
  type PrePhysicsTickEvent = PhysicsTickEvent
  type PhysicsTickObservable =
    Observable[Either[PrePhysicsTickEvent, PhysicsTickEvent]]

  implicit final class JMEAppExt(private val app: Application) extends AnyVal {

    def enqueueR(cb: () => Unit) =
      app.enqueue(new Runnable {
        override def run() = cb()
      })

  }
  implicit final class StateManagerExt(private val asm: AppStateManager)
      extends AnyVal {
    def state[S <: AppState]()(implicit c: ClassTag[S]): S =
      asm.getState(c.runtimeClass.asInstanceOf[Class[S]])

    // def appStates = asm.getStates()

  }

  implicit final class SimpleApplicationExt[T <: SimpleApplication](
      private val sa: T
  ) extends AnyVal {
    def stateManager: AppStateManager = sa.getStateManager()
    def inputManager: InputManager = sa.getInputManager()
    def assetManager: AssetManager = sa.getAssetManager()
    def guiNode = sa.getGuiNode()
    def flyCam = Option(sa.getFlyByCamera())
    def camera = sa.getCamera()
    def viewPort = sa.getViewPort()
    def rootNode = sa.getRootNode()

    def observableTick: Observable[Float] =
      Observable.create(OverflowStrategy.Unbounded) { sub =>
        val c = SingleAssignCancelable()
        val as = new MyBaseState {

          override def init(): Unit = {}

          override def update(tpf: Float) = {
            if (sub.onNext(tpf) == Ack.Stop)
              c.cancel()
          }

          override def stop(): Unit = {}

          override def onEnable() = {}
          override def onDisable() = {}
        }
        sa.stateManager.attach(as)
        c := Cancelable(() => sa.stateManager.detach(as))
        c
      }
  }

  implicit final class AssetManagerExt(private val am: AssetManager)
      extends AnyVal {
    def registerLocator(
        assetPath: os.RelPath,
        locator: Class[_ <: AssetLocator]
    ): Unit = {
      am.registerLocator(assetPath.toString(), locator)
    }

    def loadModel(assetPath: os.RelPath): Spatial = {
      am.loadModel(assetPath.toString())
    }
  }

  implicit final class BulletAppStateExt(private val bas: BulletAppState)
      extends AnyVal {
    def physicsSpace = bas.getPhysicsSpace()
    def speed = bas.getSpeed()
  }

  implicit final class BetterCharacterControlExt(
      private val bcc: BetterCharacterControl
  ) {
    def withJumpForce(force: ImVector3f) = {
      bcc.setJumpForce(force.mutable)
      bcc
    }
  }

  implicit final class SpatialExt[T <: Spatial](private val spat: T)
      extends AnyVal {
    // def asRef = Ref[Task].of(spat)
  }

  implicit final class NodeExt[T <: Node](private val n: T) extends AnyVal {

    /**
      * Attaches the given child
      *
      * @param s
      * @return
      */
    def withChild(s: Spatial): Node = {
      n.attachChild(s)
      n
    }

    /**
      * @return Gets the list of children as a monix observable
      */
    def observableChildren =
      Observable.fromIterable(n.getChildren().asScala)

    /**
      * @return A copy of the list of children of this node as a lazy list
      */
    def children = LazyList.from(n.getChildren().asScala)

    /**
      * Attach given children
      *
      * @param lst
      */
    def withChildren(lst: Spatial*): Node = {
      for (c <- lst) n.attachChild(c)
      n
    }

    def +=(spatial: Spatial) = n.attachChild(spatial)

    def :+(spatial: Spatial) = {
      n += spatial
      n
    }

    def :+(light: Light) = {
      n.addLight(light)
      n
    }

    def -=(spatial: Spatial) = n.detachChild(spatial)

    def :-(spatial: Spatial) = {
      n -= spatial
      n
    }

    def :-(light: Light) = {
      n.removeLight(light)
      n
    }

    def depthFirst(cb: Spatial => Unit) =
      n.depthFirstTraversal(new SceneGraphVisitor() {
        override def visit(s: Spatial) = cb(s)
      })

    def observableDepthFirst(): Observable[Spatial] = {
      def loop(
          subscriber: Subscriber[Spatial],
          spatials: LazyList[Spatial]
      ): Task[Unit] =
        spatials match {
          // spatial can be either a node or a geometry, but it's not a sealed trait
          case head #:: tail =>
            head match {
              case g: Geometry =>
                Task.deferFuture(subscriber.onNext(g)).flatMap {
                  case Continue =>
                    loop(subscriber, tail)
                  case Stop =>
                    Task(subscriber.onComplete())
                }
              case node: Node =>
                val children = node.children
                Task.deferFuture(subscriber.onNext(node)).flatMap {
                  case Continue =>
                    loop(subscriber, children #::: tail)
                  case Stop =>
                    Task(subscriber.onComplete())
                }
              // case _ => loop(subscriber, tail)
            }
          case LazyList() => Task.unit
        }

      Observable.create(OverflowStrategy.Unbounded) { sub =>
        implicit val sched = sub.scheduler
        loop(sub, LazyList(n)).runToFuture
      }
    }

    def breadthFirst(cb: Spatial => Unit) =
      n.breadthFirstTraversal(new SceneGraphVisitor() {
        override def visit(s: Spatial) = cb(s)
      })

    def observableBreadthFirst(): Observable[Spatial] = {
      def loop(
          subscriber: Subscriber[Spatial],
          spatials: LazyList[Spatial]
      ): Task[Unit] =
        spatials match {
          // spatial can be either a node or a geometry, but it's not a sealed trait
          case head #:: tail =>
            head match {
              case g: Geometry =>
                Task.deferFuture(subscriber.onNext(g)).flatMap {
                  case Continue =>
                    loop(subscriber, tail)
                  case Stop =>
                    Task(subscriber.onComplete())
                }
              case node: Node =>
                val children = node.children
                Task.deferFuture(subscriber.onNext(node)).flatMap {
                  case Continue =>
                    loop(subscriber, tail #::: children)
                  case Stop =>
                    Task(subscriber.onComplete())
                }
              case unknown =>
                Task.deferFuture(subscriber.onNext(unknown)).flatMap {
                  case Continue =>
                    loop(subscriber, tail)
                  case Stop =>
                    Task(subscriber.onComplete())
                }
            }
          case LazyList() => Task.unit
        }

      Observable.create(OverflowStrategy.Unbounded) { sub =>
        implicit val sched = sub.scheduler
        loop(sub, LazyList(n)).runToFuture
      }
    }

    def withControl[C <: Control](ctrl: C) = {
      n.addControl(ctrl)
      n
    }

    def withLocalTranslation(dir: ImVector3f) = {
      n.setLocalTranslation(dir.mutable)
      n
    }

    def withRotate(xAngle: Float, yAngle: Float, zAngle: Float) = {
      n.rotate(xAngle, yAngle, zAngle)
      n
    }

    def localRotation = n.getLocalRotation()

    def localTranslation = n.getLocalTranslation()

  }

  implicit final class CameraNodeExt(private val cn: CameraNode) {
    def withControlDir(controlDir: ControlDirection) = {
      cn.setControlDir(controlDir)
      cn
    }

    def withLookAt(position: ImVector3f, upVector: ImVector3f) = {
      cn.lookAt(position.mutable, upVector.mutable)
      cn
    }
  }

  implicit final class GeometryExt(private val geom: Geometry) {
    def withMaterial(mat: Material) = {
      geom.setMaterial(mat)
      geom
    }
  }

  implicit final class EntityDataExt(private val ed: EntityData)
      extends AnyVal {

    def query = new EntityQuery(ed)

    // def entities[T <: EntityComponent](entities: Seq[T])
  }

  implicit final class EntityExt(private val e: EntityId) extends AnyVal {
    def withComponents(classes: EntityComponent*)(implicit
        ed: EntityData
    ): EntityId = {
      ed.setComponents(e, classes: _*)
      e
    }
  }

  implicit final class ActorRefExt[Req](private val a: ActorRef[Req])
      extends AnyVal {
    import akka.actor.typed.scaladsl.AskPattern._

    /**
      *  Same as [[ask]] but returns a [[Task]]
      */
    def askL[Res](
        replyTo: ActorRef[Res] => Req
    )(implicit timeout: Timeout, scheduler: Scheduler): Task[Res] = {
      Task.deferFuture(a.ask(replyTo)(timeout, scheduler))
    }
    def ??[Res](
        replyTo: ActorRef[Res] => Req
    )(implicit timeout: Timeout, scheduler: Scheduler): Task[Res] =
      askL(replyTo)

    /**
      * Same as [[tell]], but wrapped in a Task
      *
      * @param msg
      * @return
      */
    def tellL(msg: Req) = UIO(a.tell(msg))

    /**
      * Same as [[tell]], but wrapped in a Task
      *
      * @param msg
      * @return
      */
    def !!(msg: Req) = tellL(msg)

  }

  implicit final class InputManagerExt(private val inputManager: InputManager)
      extends AnyVal {
    def withMapping(mapping: String, triggers: Trigger*): InputManager = {
      inputManager.addMapping(mapping, triggers: _*)
      inputManager
    }

    def withListener(listener: InputListener, mappings: String*) = {
      inputManager.addListener(listener, mappings: _*)
      inputManager
    }

    def withEnumMappings[T <: EnumEntry](
        mappingEnum: Enum[T]
    )(mappingFn: T => Seq[Trigger]): InputManager = {
      for (entry <- mappingEnum.values) {
        val mappings = mappingFn(entry)
        inputManager.addMapping(entry.entryName, mappings: _*)
      }
      inputManager
    }

    def observableAction(mappingNames: String*): Observable[ActionEvent] = {

      Observable.create(OverflowStrategy.DropOld(10)) { sub =>
        val c = SingleAssignCancelable()
        val al = new ActionListener {
          override def onAction(
              binding: String,
              value: Boolean,
              tpf: Float
          ): Unit = {
            if (
              sub.onNext(ActionEvent(Action(binding), value, tpf)) == Ack.Stop
            ) {
              sub.onComplete()
              c.cancel()
            }
          }
        }

        inputManager.addListener(al, mappingNames: _*)

        c := Cancelable(() => inputManager.removeListener(al))
        c
      }
    }

    def enumObservableAction[T <: EnumEntry](
        mappingEnum: Enum[T]
    ): Observable[EnumActionEvent[T]] = {

      Observable.create(OverflowStrategy.DropOld(10)) { sub =>
        val c = SingleAssignCancelable()
        val entryNames = mappingEnum.values.map(_.entryName)
        val al = new ActionListener {
          override def onAction(
              binding: String,
              value: Boolean,
              tpf: Float
          ): Unit = {
            mappingEnum.withNameOption(binding).foreach { b =>
              if (sub.onNext(EnumActionEvent(b, value, tpf)) == Ack.Stop) {
                sub.onComplete()
                c.cancel()
              }
            }
          }
        }

        inputManager.addListener(al, entryNames: _*)

        c := Cancelable(() => inputManager.removeListener(al))
        c
      }
    }

    def analogObservable(mappingNames: String*): Observable[AnalogEvent] = {

      Observable.create(OverflowStrategy.DropOld(50)) { sub =>
        val c = SingleAssignCancelable()
        val al = new AnalogListener {
          override def onAnalog(
              binding: String,
              value: Float,
              tpf: Float
          ): Unit = {
            if (
              sub.onNext(AnalogEvent(Action(binding), value, tpf)) == Ack.Stop
            ) {
              sub.onComplete()
              c.cancel()
            }
          }
        }

        inputManager.addListener(al, mappingNames: _*)

        c := Cancelable(() => inputManager.removeListener(al))
        c
      }
    }

    def enumAnalogObservable[T <: EnumEntry](
        mappingEnum: Enum[T]
    ): Observable[EnumAnalogEvent[T]] = {

      Observable.create(OverflowStrategy.DropOld(50)) { sub =>
        val c = SingleAssignCancelable()
        val entryNames = mappingEnum.values.map(_.entryName)
        val al = new AnalogListener {
          override def onAnalog(
              binding: String,
              value: Float,
              tpf: Float
          ): Unit = {
            mappingEnum.withNameOption(binding).foreach { b =>
              if (sub.onNext(EnumAnalogEvent(b, value, tpf)) == Ack.Stop) {
                sub.onComplete()
                c.cancel()
              }
            }
          }
        }

        inputManager.addListener(al, entryNames: _*)

        c := Cancelable(() => inputManager.removeListener(al))
        c
      }
    }
  }

  implicit final class PhysicsSpaceExt(private val space: PhysicsSpace)
      extends AnyVal {

    def collisionObservable(): Observable[PhysicsCollisionEvent] = {

      Observable.create(OverflowStrategy.DropOld(50)) { sub =>
        val c = SingleAssignCancelable()
        val cl = new PhysicsCollisionListener {
          override def collision(event: PhysicsCollisionEvent): Unit = {

            if (sub.onNext(event) == Ack.Stop) {
              sub.onComplete()
              c.cancel()
            }

          }
        }

        space.addCollisionListener(cl)

        c := Cancelable(() => space.removeCollisionListener(cl))
        c
      }
    }

    def prePhysicsTickObservable(): Observable[PrePhysicsTickEvent] = {

      Observable.create(OverflowStrategy.DropOld(50)) { sub =>
        val c = SingleAssignCancelable()
        val cl = new PhysicsTickListener {

          override def prePhysicsTick(space: PhysicsSpace, tpf: Float): Unit = {
            val event = new PrePhysicsTickEvent(space)
            if (sub.onNext(event) == Ack.Stop) {
              sub.onComplete()
              c.cancel()
            }
          }

          override def physicsTick(space: PhysicsSpace, tpf: Float): Unit = {}

        }

        space.addTickListener(cl)

        c := Cancelable(() => space.removeTickListener(cl))
        c
      }
    }

    def physicsTickObservable(): Observable[PhysicsTickEvent] = {

      Observable.create(OverflowStrategy.DropOld(50)) { sub =>
        val c = SingleAssignCancelable()
        val cl = new PhysicsTickListener {

          override def prePhysicsTick(
              space: PhysicsSpace,
              tpf: Float
          ): Unit = {}

          override def physicsTick(space: PhysicsSpace, tpf: Float): Unit = {
            val event = new PhysicsTickEvent(space)
            if (sub.onNext(event) == Ack.Stop) {
              sub.onComplete()
              c.cancel()
            }
          }

        }

        space.addTickListener(cl)

        c := Cancelable(() => space.removeTickListener(cl))
        c
      }
    }

    //TODO Consider creating a typeclass for this
    def +=(anyObject: Any) = space.add(anyObject)

    def :+(anyObject: Any) = {
      space.add(anyObject)
      space
    }
    def :-(anyObject: Any) = {
      space.remove(anyObject)
      space
    }

    def +=(spatial: Spatial) = space.addAll(spatial)

    def :+(spatial: Spatial) = {
      space.addAll(spatial)
      space
    }

    def :-(spatial: Spatial) = {
      space.removeAll(spatial)
      space
    }

    // def asRef = Ref[Task].of(space)

  }

  implicit final class Vector3fExt(private val v: Vector3f) extends AnyVal {
    //TODO add more operations
    def +=(that: Vector3f) = v.addLocal(that)
    def +=(f: Float) = v.addLocal(f, f, f)
    def +=(that: ImVector3f) = v.addLocal(that.x, that.y, that.z)
    def +=:(that: ImVector3f) = v += that
    def *=(that: Vector3f) = v.multLocal(that)
    def *=(that: ImVector3f) = v.multLocal(that.x, that.y, that.z)
    def *=:(that: ImVector3f) = v *= that
    def -=(that: Vector3f) = v.subtractLocal(that)
    def -=(that: ImVector3f) = v.subtractLocal(that.x, that.y, that.z)
    def -=:(that: ImVector3f) = v *= that
    def /=(that: Vector3f) = v.divideLocal(that)
    def /=(that: ImVector3f) = v.divideLocal(that.mutable)
    def /=:(that: ImVector3f) = v *= that
    def unary_- = v.negateLocal()
    def immutable = ImVector3f(v.x, v.y, v.z)
  }

  implicit final class ImVector3fExt(private val v: ImVector3f) extends AnyVal {
    def +(that: ImVector3f) = v.copy(v.x + that.x, v.y + that.y, v.z + that.z)
    def +(f: Float): ImVector3f = v.copy(v.x + f, v.y + f, v.z + f)
    def *(that: ImVector3f) = v.copy(v.x * that.x, v.y * that.y, v.z * that.z)
    def *(f: Float): ImVector3f =
      v.copy(v.x * f, v.y * f, v.z * f)
    // v * ImVector3f(f, f, f)
    def -(that: ImVector3f) = v.copy(v.x - that.x, v.y - that.y, v.z - that.z)
    def -(f: Float): ImVector3f = v.copy(v.x - f, v.y - f, v.z - f)
    def /(that: ImVector3f) = v.copy(v.x / that.x, v.y / that.y, v.z / that.z)
    def /(f: Float): ImVector3f = v.copy(v.x / f, v.y / f, v.z / f)
    def unary_- = v.copy(-v.x, -v.y, -v.z)
    def mutable = new Vector3f(v.x, v.y, v.z)
  }

  // val TasktoUIO = new FunctionK[Task, UIO] {
  //   def apply[T](f: Task[T]): UIO[T] =
  //     f.hideErrors
  // }

  implicit final class JavaFxUIExt(private val jfxui: JavaFxUI) extends AnyVal {
    def +=(node: scalafx.scene.Node) = jfxui.attachChild(node)
    def -=(node: scalafx.scene.Node) = jfxui.detachChild(node)
  }
}