Main.scala

import scalaz.nio._
import scalaz.nio.channels.{AsynchronousServerSocketChannel, AsynchronousSocketChannel}
import scalaz.zio.clock.Clock
import scalaz.zio.console._
import scalaz.zio.{App, ZIO, _}


object Main extends App {
  override def run(args: List[String]): ZIO[Environment, Nothing, Int] = {
    theSocket.foldM(
      err => putStrLn(s"Execution Failed with: $err") *> ZIO.succeed(1),
      _   => ZIO.succeed(0)
    )
  }
  val theSocket =
    for {
      address     <- SocketAddress.inetSocketAddress("127.0.0.1", 1337)
      socket      <- AsynchronousServerSocketChannel()
      _           <- socket.bind(address)
      ref         <- Ref.make[Int](0)
      connections <- awaitConnection(ref, socket, doWork)
    } yield ()


  /*
     * Accept a connection from the server, fork the worker on it, and
     * loop to wait for next connection
     */
  def awaitConnection(ref: Ref[Int], socket: AsynchronousServerSocketChannel, worker: (Int, Int) => AsynchronousSocketChannel => ZIO[Console with Clock, Throwable, Unit]): ZIO[Console with Clock, Throwable, Nothing] = {
    for {
      i    <- ref.update(_+1)
      _    <- putStrLn("accept")
      _    <- socket.accept.flatMap(s => doWork(i)(s).ensuring(s.close.orDie).fork)
      _    <- putStrLn("next")
      loop <- awaitConnection(ref, socket, worker)
    } yield loop
  }
  /*
     * From a connected AsynchronousSocketChannel, read forever, until connection termination
     */
  def doWork(i:Int, rep: Int = 0)(work: AsynchronousSocketChannel): ZIO[Console with Clock, Throwable, Unit] = {
      for {
        chunk <- work.read(16)
        str = chunk.toArray.map(_.toChar).map( c => if (c.isControl) s" 0x${c.toHexString.toUpperCase} " else s"$c" ).mkString
        _ <- putStrLn(s"content for con $i: " + str)
        isOpen <- work.isOpen
        _ <- putStrLn(s"connection status: $isOpen")
        loop <- if (isOpen && rep < 5) doWork(i, rep + 1)(work) else ZIO.unit
      } yield loop
  }
}