ScreenCastService.scala

package com.hanhuy.android.airshow

import java.net.{DatagramPacket, InetAddress, InetSocketAddress, DatagramSocket}
import java.nio.ByteBuffer
import java.nio.channels.DatagramChannel
import java.util.concurrent.TimeUnit

import android.content.res.Configuration
import android.graphics.Point
import android.hardware.display.VirtualDisplay
import android.media.{MediaCodecList, MediaCodecInfo, MediaFormat, MediaCodec}
import android.media.projection.{MediaProjection, MediaProjectionManager}
import android.view.{WindowManager, Surface}
import com.hanhuy.android.common._
import com.hanhuy.android.conversions._
import Futures._

import android.app.{PendingIntent, Activity, Notification, Service}
import android.content.{Context, BroadcastReceiver, Intent}
import ScreenCastService._

import scala.annotation.tailrec
import scala.concurrent.Future

object ScreenCastService {
  val log = Logcat("ScreenCastService")
  val EXTRA_RESULT_INTENT = "com.hanhuy.android.airshow.extra.RESULT_INTENT"
  val REQUEST_STOP = 5706
  val ACTION_STOP = "com.hanhuy.android.airshow.action.STOP"
}
class ScreenCastService extends Service {
  @volatile var encoderRunning = true
  val NOTIFICATION_ID = 0x31337
  var mediaProjection = Option.empty[MediaProjection]
  var mediaCodec = Option.empty[MediaCodec]
  var surface = Option.empty[Surface]
  var socket = Option.empty[DatagramChannel]
  var virtualDisplay = Option.empty[VirtualDisplay]
  override def onBind(intent: Intent) = ???

  @tailrec
  final def downScaleSize(w: Int, h: Int, dpi: Int, tw: Int, th: Int): (Int,Int,Int) = {
    if (w <= tw && h <= th) {
      (w, h, dpi)
    } else {
      val ratio = math.max(w.toDouble / tw, h.toDouble / th)
      downScaleSize((w / ratio).toInt, (h / ratio).toInt, (dpi / ratio).toInt, tw, th)
    }
  }
  override def onStartCommand(intent: Intent, flags: Int, startId: Int) = {
    val p = new Point
    this.systemService[WindowManager].getDefaultDisplay.getRealSize(p)
    getResources.getDisplayMetrics.xdpi
    encoderRunning = true
    registerReceiver(receiver, ACTION_STOP)
    val notification = new Notification.Builder(this)
      .setContentTitle("Sharing screen with ...")
      .setSmallIcon(R.drawable.ic_picture_in_picture_white_18dp)
      .addAction(android.R.drawable.ic_menu_close_clear_cancel, "Stop",
        PendingIntent.getBroadcast(this, REQUEST_STOP, new Intent(ACTION_STOP),
          PendingIntent.FLAG_UPDATE_CURRENT))
      .build

    val f = Future {
      val dc = DatagramChannel.open()
      socket = Option(dc)
    }
    f onSuccessMain { case ds =>
      val list = new MediaCodecList(MediaCodecList.REGULAR_CODECS)
      val infos = list.getCodecInfos
      val matched = infos.filter(_.isEncoder).filter(_.getSupportedTypes.contains(MediaFormat.MIMETYPE_VIDEO_AVC))
      matched.foreach { info =>
        val caps = info.getCapabilitiesForType(MediaFormat.MIMETYPE_VIDEO_AVC)
        log.v(s"color formats: ${caps.colorFormats mkString ":"}")
        log.v(s"profile levels: ${caps.profileLevels map (l => (l.level, l.profile)) mkString ":"}")
        log.v(s"default format: ${caps.getDefaultFormat}")
        val vidCaps = caps.getVideoCapabilities
        log.v(s"bitrate range: ${vidCaps.getBitrateRange} ha: ${vidCaps.getHeightAlignment}, widths: ${vidCaps.getSupportedWidths}, heights: ${vidCaps.getSupportedHeights}, fps: ${vidCaps.getSupportedFrameRates}")
        log.v(s"codec info: ${info.getName} ${info.getSupportedTypes.mkString(":")}")
      }

      val scaled = downScaleSize(p.x, p.y, getResources.getDisplayMetrics.xdpi.toInt, 1280, 720)
      log.v("scaling to: " + scaled)
      val mc = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
      val format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, scaled._1, scaled._2)
      format.setInteger(MediaFormat.KEY_COLOR_FORMAT,
        MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface)
      format.setInteger(MediaFormat.KEY_BIT_RATE, 10000000)
      format.setInteger(MediaFormat.KEY_FRAME_RATE, 30)
      format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 5)
      mc.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
      val s = mc.createInputSurface
      surface = Option(s)
      mediaCodec = Some(mc)
      mc.start()
      val mpm = this.systemService[MediaProjectionManager]
      val mp = mpm.getMediaProjection(
        Activity.RESULT_OK, intent.getParcelableExtra(EXTRA_RESULT_INTENT))
      virtualDisplay = Option(mp.createVirtualDisplay("airshow", scaled._1, scaled._2, scaled._3, 0, s, null, null))

      mp.registerCallback(new MediaProjection.Callback {
        override def onStop() = {
          shutdown()
          super.onStop()
        }
      }, UiBus.handler)
      mediaProjection = Option(mp)
      log.v("Starting encode thread")
      encode()
    }
    f onFailure { case ex =>
        log.v("failed to setup connection", ex)
    }
    startForeground(NOTIFICATION_ID, notification)
    Service.START_NOT_STICKY
  }

  def encode(): Unit = {
    new Thread(() => {
      val bufferInfo = new MediaCodec.BufferInfo
      mediaCodec foreach { mc =>
        @tailrec
        def encodeBlock(): Unit = {
          val status = mc.dequeueOutputBuffer(bufferInfo, TimeUnit.MILLISECONDS.toMicros(16))
          if (status == MediaCodec.INFO_TRY_AGAIN_LATER) {
            if (!encoderRunning) {
              log.v("Exiting encoder loop because encoder exit requested")
              shutdownEncoder()
            } else
              encodeBlock()
          } else if (status >= 0) {
            // encoded sample
            val data = mc.getOutputBuffer(status)
            if (data != null) {
              val endOfStream = bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM
              // pass to whoever listens to
              if (endOfStream == 0) onEncodedSample(bufferInfo, data)
              // releasing buffer is important
              mc.releaseOutputBuffer(status, false)
              if (endOfStream != MediaCodec.BUFFER_FLAG_END_OF_STREAM)
                encodeBlock()
              else {
                log.v("End of stream detected")
                shutdownEncoder()
              }
            } else
              encodeBlock()
          } else if (status == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
            encodeBlock()
          } else {
            log.v("Exiting because status: " + status)
            shutdownEncoder()
          }
        }
        encodeBlock()
      }
    }, "Encoding Loop").start()
  }

  private[this] val sockaddr = new InetSocketAddress("192.168.9.2", 12346)
  private[this] val packet = ByteBuffer.allocate(1500)
  @tailrec
  private[this] def rtpize(sock: DatagramChannel, seq: Short, data: ByteBuffer, offset: Int, limit: Int, clock: Int): Short = {
    val header1: Byte = 0x70 // version 2, no padding, extension nor csrc
    val newseq = (seq + 1).toShort
    val packetSize = 1488 // 1500 mtu - 12 bytes for header
    packet.position(0)
    packet.limit(packet.capacity)
    packet.put(header1)
    data.position(offset)

    if (limit - offset <= packetSize) {
      val header2: Byte = (96 | 0x80).toByte // marker bit + payload type 96
      data.limit(limit)
      packet.put(header2)
      packet.putShort(newseq)
      packet.putInt(clock)
      packet.putInt(0) // ssrc 0
      packet.put(data)
      packet.flip()
      sock.send(packet, sockaddr)
      newseq
    } else {
      val header2: Byte = 96 // payload type 96, no marker
      val end = math.min(limit, offset + packetSize)
      data.limit(end)
      packet.put(header2)
      packet.putShort(newseq)
      packet.putInt(clock)
      packet.putInt(0) // ssrc 0
      packet.put(data)
      packet.flip()
      sock.send(packet, sockaddr)
      rtpize(sock, newseq, data, end, limit, clock)
    }
  }

  private[this] var datarates = Map.empty[Long,Int]
  private[this] var seqN = (util.Random.nextInt() & 0xffff).toShort
  private[this] var pps = Array.empty[Byte]
  private[this] var sps = Array.empty[Byte]
  def onEncodedSample(info: MediaCodec.BufferInfo, data: ByteBuffer): Unit = {
    val seconds = System.currentTimeMillis() / 1000
    datarates = datarates.updated(seconds, info.size + datarates.getOrElse(seconds, 0))
    if (datarates.size > 1) {
      val rates = datarates.toList.sortBy(_._1)
      log.v(s"Datarate: ${rates.head._2 / 1000 * 8}kbps")
      datarates = rates.lastOption.toMap
    }
    socket foreach { sock =>
      val ts = System.nanoTime()
      val clock = ((ts * 11111) & 0xffffffff).toInt
      data.position(info.offset + 4)
      @tailrec
      def send(limit: Int): Unit = {
        data.limit(limit)
        if (limit - data.position > 1500) {
          data.limit(data.position + 1500)
        }
        sock.send(data, sockaddr)
        if (limit - data.position > 0) {
          send(limit)
        }
      }
      data.position(info.offset)
      val x = List(data.get, data.get, data.get, data.get, data.get, data.get, data.get, data.get)
      log.v(f"${info.size}%6d => " + x.mkString(", "))
      data.position(info.offset)
      data.limit(info.offset + info.size)
      send(info.offset + info.size)
      /*
      if ((data.get & 0x1f) == 7) {
        // sps & pps data
        data.position(info.offset)
        data.limit(info.offset + info.size)
        val seq = (0 until info.size) map (_ => data.get)
        log.v("sps/pps data: " + seq.mkString(", "))
        val ppsIdx = seq.drop(4).indexOfSlice(Seq(0,0,0,1): Seq[Byte])
        data.limit(ppsIdx + 4) // because start seq was dropped
        data.position(info.offset + 4)
        sps = Array.ofDim[Byte](data.remaining)
        data.get(sps)
        data.limit(info.offset + info.size)
        data.position(ppsIdx + 4 + 4) // because start seq was dropped
        pps = Array.ofDim[Byte](data.remaining)
        data.get(pps)
        log.v("sps: " + sps.mkString(", "))
        log.v("pps: " + pps.mkString(", "))
        // STAP-A NAL header + NALU 1 (SPS) size + NALU 2 (PPS) size = 5 bytes
        val stapa = Array.ofDim[Byte](sps.length + pps.length + 5)

        // STAP-A NAL header is 24
        stapa(0) = 24

        // Write NALU 1 size into the array (NALU 1 is the SPS).
        stapa(1) = (sps.length >> 8).toByte
        stapa(2) = (sps.length & 0xFF).toByte

        // Write NALU 2 size into the array (NALU 2 is the PPS).
        stapa(sps.length + 3) = (pps.length >> 8).toByte
        stapa(sps.length + 4) = (pps.length & 0xFF).toByte

        // Write NALU 1 into the array, then write NALU 2 into the array.
        System.arraycopy(sps, 0, stapa, 3, sps.length)
        System.arraycopy(pps, 0, stapa, 5 + sps.length, pps.length)
        seqN = rtpize(sock, seqN, ByteBuffer.wrap(stapa), 0, stapa.length, clock)
      } else {
        data.limit(info.offset + info.size)
        seqN = rtpize(sock, seqN, data, info.offset + 4, info.offset + info.size, clock)
      }
      */
    }
  }

  val receiver: BroadcastReceiver = (context: Context, intent: Intent) => {
    intent.getAction match {
      case ACTION_STOP => shutdown()
      case _ =>
    }
  }

  override def onConfigurationChanged(newConfig: Configuration) = {
    log.v("orientation: " + newConfig.orientation)
    super.onConfigurationChanged(newConfig)
  }

  private[this] def shutdownEncoder(): Unit = {
    mediaCodec.foreach { mc =>
      mc.stop()
      mc.reset()
      mc.release()
    }
    mediaCodec = None
  }
  private[this] def shutdown(): Unit = {
    encoderRunning = false
    mediaProjection foreach { mp =>
      unregisterReceiver(receiver)
      stopForeground(true)
      stopSelf()
      mp.stop()
    }
    virtualDisplay.foreach { vd => vd.release() }
    socket foreach (_.close())
    surface foreach (_.release())
    mediaProjection = None
    surface = None
    virtualDisplay = None
  }
}