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

package wow.doge.mygame

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

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 sm: AppStateManager)
      extends AnyVal {
    def state[S <: AppState]()(implicit c: ClassTag[S]): S =
      sm.getState(c.runtimeClass.asInstanceOf[Class[S]])

  }

  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 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)

    def children = LazyList.from(n.getChildren().asScala)

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

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

    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 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 enumObservableAction[T <: enumeratum.EnumEntry](
    //     mappingNames: Enum[T]
    // ): Observable[ActionEvent] = {
    //   observableAction2(mappingNames).doOnNext()
    // }

    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 +=(spatial: Spatial) = space.addAll(spatial)

  }

  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 Vector3fExt(private val v: Vector3f) extends AnyVal {
    //TODO add more operations
    def +=(that: Vector3f) = v.addLocal(that)
    def +=(that: ImVector3f) = v.addLocal(that.x, that.y, that.z)
    def +=:(that: ImVector3f) = v += that
    def *=(that: Vector3f) = v.multLocal(that)
    def -=(that: Vector3f) = v.subtractLocal(that)
    def /=(that: Vector3f) = v.divideLocal(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 *(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 /(that: ImVector3f) = v.copy(v.x / that.x, v.y / that.y, v.z / that.z)
    def unary_- = v.copy(-v.x, -v.y, -v.z)
    // def unary_-(that: ImVector3f) = this.copy(this.x, this.y, this.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
  // }

}

// Observable.create(OverflowStrategy.Unbounded) { sub =>
//   // val c = SingleAssignCancelable()
//   val visitor = new SceneGraphVisitor {
//     override def visit(s: Spatial): Unit = {
//       sub.onNext(s)
//       // if (sub.onNext(s) == Ack.Stop)
//       // c.cancel()

//     }
//   }
//   n.depthFirstTraversal(visitor)
//   // c := Cancelable(() => ???)
//   // c
//   Cancelable.empty