package nelson

import cats.Eval
import cats.effect.{Effect, IO, Timer}
import cats.free.Cofree
import cats.syntax.functor._
import cats.syntax.monadError._

import fs2.{Pipe, Sink, Stream}

import quiver.{Context, Decomp, Graph}

import java.util.concurrent.TimeoutException

import scala.concurrent.ExecutionContext
import scala.concurrent.duration.FiniteDuration
import scala.collection.immutable.{Stream => SStream}

object CatsHelpers {
  implicit class NelsonEnrichedIO[A](val io: IO[A]) extends AnyVal {
    /** Run `other` if this IO fails */
    def or(other: IO[A]): IO[A] = io.attempt.flatMap {
      case Right(a) => IO.pure(a)
      case Left(_)  => other
    }

    /** Fail with error if the result of the IO does not satsify the predicate
      *
      * Taken from https://github.com/scalaz/scalaz/blob/series/7.3.x/concurrent/src/main/scala/scalaz/concurrent/Task.scala
      */
    def ensure(failure: => Throwable)(f: A => Boolean): IO[A] =
      io.flatMap(a => if (f(a)) IO.pure(a) else IO.raiseError(failure))

    def timed(timeout: FiniteDuration)(implicit ec: ExecutionContext): IO[A] =
      IO.race(
        Timer[IO].sleep(timeout).as(new TimeoutException(s"Timed out after ${timeout.toMillis} milliseconds"): Throwable),
        io
      ).rethrow
  }

  private def sinkW[F[_], W, O](actualSink: Sink[F, W]): Sink[F, Either[W, O]] =
    stream => actualSink(stream.collect { case Left(e) => e })

  private def pipeO[F[_], W, O, O2](actualPipe: Pipe[F, O, O2]): Pipe[F, Either[W, O], Either[W, O2]] =
    _.flatMap {
      case Left(a)  => Stream.emit(Left(a))
      case Right(b) => actualPipe(Stream.emit(b)).map(Right(_))
    }

  implicit class NelsonEnrichedWriterStream[F[_], W, O](val stream: Stream[F, Either[W, O]]) {
    def observeW(sink: Sink[F, W])(implicit F: Effect[F], ec: ExecutionContext): Stream[F, Either[W, O]] =
      stream.observe(sinkW(sink))

    def stripW: Stream[F, O] = stream.collect { case Right(o) => o }

    def throughO[O2](pipe: Pipe[F, O, O2]): Stream[F, Either[W, O2]] =
      stream.through(pipeO(pipe))
  }


  /** This is pending release of https://github.com/Verizon/quiver/pull/31 */
  private type Tree[A] = Cofree[SStream, A]

  private def flattenTree[A](tree: Tree[A]): SStream[A] = {
    def go(tree: Tree[A], xs: SStream[A]): SStream[A] =
      SStream.cons(tree.head, tree.tail.value.foldRight(xs)(go(_, _)))
    go(tree, SStream.Empty)
  }

  private def Node[A](root: A, forest: => SStream[Tree[A]]): Tree[A] =
    Cofree[SStream, A](root, Eval.later(forest))

  implicit class NelsonEnrichedGraph[N, A, B](val graph: Graph[N, A, B]) extends AnyVal {
    def reachable(v: N): Vector[N] =
      xdfWith(Seq(v), _.successors, _.vertex)._1.flatMap(flattenTree)

    def xdfWith[C](vs: Seq[N], d: Context[N, A, B] => Seq[N], f: Context[N, A, B] => C): (Vector[Tree[C]], Graph[N, A, B]) =
      if (vs.isEmpty || graph.isEmpty) (Vector(), graph)
      else graph.decomp(vs.head) match {
        case Decomp(None, g) => g.xdfWith(vs.tail, d, f)
        case Decomp(Some(c), g) =>
          val (xs, _) = g.xdfWith(d(c), d, f)
          val (ys, g3) = g.xdfWith(vs.tail, d, f)
          (Node(f(c), xs.toStream) +: ys, g3)
      }
  }
}