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 cats.effect.concurrent.Ref
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.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

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

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

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

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

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

  }
  implicit 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 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 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 class BulletAppStateExt(private val bas: BulletAppState)
      extends AnyVal {
    def physicsSpace = bas.getPhysicsSpace()
    def speed = bas.getSpeed()
  }

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

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

  implicit 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
    }

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

    /**
      * A copy of the list of children of this node as a lazy list
      *
      * @return
      */
    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 -=(spatial: Spatial) = n.detachChild(spatial)

    def :-(spatial: Spatial) = {
      n -= spatial
      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 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 class EntityDataExt(private val ed: EntityData) extends AnyVal {

    def query = new EntityQuery(ed)

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

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

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

    /**
      * @param replyTo
      * @param timeout
      * @param scheduler
      * @return
      */
    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))
    def !!(msg: Req) = tellL(msg)

  }
  // def ?[Res](replyTo: ActorRef[Res] => Req)(implicit timeout: Timeout, scheduler: Scheduler): Future[Res] = {
  //     ask(replyTo)(timeout, scheduler)
  //   }

  implicit 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 class PhysicsSpaceExt(private val space: PhysicsSpace)
      extends AnyVal {

    def collisionObservable(): Observable[PhysicsCollisionEvent] = {

      Observable.create(OverflowStrategy.Unbounded) { 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 physicsTickObservable(): PhysicsTickObservable = {

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

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

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

        }

        space.addTickListener(cl)

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

    //TODO Create 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
    }

  }

  implicit 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 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
  // }

}