#!/usr/bin/python

'''
VideoRelay - Copy/transcode live stream to http stream.

(c) 2012-2014 Jan ONDREJ (SAL) <ondrejj(at)salstar.sk>

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

Examples:
  http://localhost:9090/raw/mark
  http://localhost:9090/x264/ch=joj/b=3000
'''

import sys, socket, os, time, re, threading, urllib, html
from httpserver import AuthHTTPRequestHandler, ThreadingHTTPServer
from StringIO import StringIO
from config import allowed_ips, auth

# Import gstreamer libraries
import gi
gi.require_version('Gst', '1.0')
from gi.repository import GObject, Gio, Gst
GObject.threads_init()
Gst.init(None)

__version__ = "0.1"

HOST = ""
PORT = 9090

GST_SOURCES = dict(
  file = """
    filesrc location=%(filename)s
    ! queue
  """,
  multicast = """
    udpsrc name=src multicast-group=%(ip)s port=%(port)s
    ! application/x-rtp, media=(string)video, clock-rate=90000, encoding-name=MP2T-ES, payload=33
    #! queue
  """
)
GST_DECODER = """
    ! decodebin name=dec
    ! avdeinterlace
    ! videoscale
    ! video/x-raw, width=%(w)s, height=%(h)s
    ! queue
"""
GST_VAAPI_DECODER = """
    ! vaapidecode name=dec
    ! avdeinterlace
    ! videoscale
    ! queue
"""
GST_SINK = """
    #! progressreport
    ! multisocketsink name=sink
    #! multifdsink name=sink
"""
GST_ENCODERS = dict(
  ogv = GST_DECODER + """
    ! theoraenc bitrate=%(kb)s
    ! mux. dec.
    ! audioconvert
    ! vorbisenc
    ! oggmux name=mux
  """ + GST_SINK,# + "max-lateness=1000000",
  vp8 = GST_DECODER + """
    ! vp8enc target-bitrate=%(b)s cpu-used=4 threads=8
    ! mux. dec.
    ! audioconvert
    ! vorbisenc
    ! webmmux name=mux
  """ + GST_SINK,# + "max-lateness=1000000",
  vp9 = GST_DECODER + """
    ! vp9enc target-bitrate=%(b)s cpu-used=4 threads=8
    ! mux. dec.
    ! audioconvert
    ! vorbisenc
    ! webmmux name=mux
  """ + GST_SINK,# + "max-lateness=1000000",
  x264 = GST_DECODER + """
    ! x264enc bitrate=%(kb)s
    ! video/x-h264, profile=main, width=%(w)s, height=%(h)s
    ! mux. dec.
    ! audioconvert
    ! lamemp3enc
    ! mpegtsmux name=mux
    #! avmux_mp4 name=mux
  """ + GST_SINK,
  x264va = GST_DECODER + """
    ! vaapiencode_h264 bitrate=%(kb)s
    ! video/x-h264, profile=main, width=%(w)s, height=%(h)s
    ! mux. dec.
    ! audioconvert
    ! lamemp3enc
    ! mpegtsmux name=mux
  """ + GST_SINK,
  mpeg4 = GST_DECODER + """
    ! avenc_mpeg4 bitrate=%(b)s
    ! video/mpeg, profile=main, width=%(w)s, height=%(h)s
    ! mux. dec.
    ! audioconvert
    ! twolamemp2enc
    ! mpegtsmux name=mux
  """ + GST_SINK,
  mpeg2va = GST_VAAPI_DECODER + """
    ! vaapiencode_mpeg2 bitrate=%(kb)s
    ! video/mpeg, profile=main, width=%(w)s, height=%(h)s
    ! mux. dec.
    ! audioconvert
    ! twolamemp2enc
    ! mpegtsmux name=mux
  """ + GST_SINK,
  mp3 = """
    ! decodebin name=dec
    ! audioconvert
    ! lamemp3enc bitrate=%(kb)s name=mux
  """ + GST_SINK,
  oga = """
    ! decodebin name=dec
    ! audioconvert
    ! vorbisenc bitrate=%(b)s
    ! oggmux name=mux
  """ + GST_SINK,
  raw = """
    ! rtpmp2tdepay
    #! queue
  """ + GST_SINK,
  rawaudio = """
    ! rtpmp2tdepay
    #! queue
  """ + GST_SINK,
  rawvideo = """
    ! rtpmp2tdepay
    #! queue
  """ + GST_SINK
)
GST_ENCODERS["auto"] = GST_ENCODERS["ogv"]

def kmg(bytes, unit="B"):
    sis = [[3, 'G'], [2, 'M'], [1, 'k']]
    for size, si in sis:
      if bytes>1024**size:
        return "%4.2f%s%s" % (float(bytes)/(1024**size), si, unit)
    return str(bytes)+unit

def dhm(seconds):
    if seconds>3600*24:
      return "%4.2f days" % (seconds/3600/24)
    elif seconds>3600:
      return "%4.2f hours" % (seconds/3600)
    return "%4.2f minutes" % (seconds/60)

def hms(seconds):
    return "%d:%02d:%02d" % (
      seconds/3600,
      seconds/60%60,
      seconds % 60
    )

class player_class:
  reg_strip_comments = re.compile("^ *(#.*)\n", re.M)
  reg_strip_spaces = re.compile("[ \t\n]+")
  bufsize = 65536
  def __init__(self, t='raw', ip=None, port=1234, file=None,
                     w=720, h=576, b=1000):
      if ip is None:
        raise ValueError, "No channel selected!"
      # start stream
      self.args = dict(
        t=t, ip=ip, port=port, filename=file,
        w=w, h=h, b=int(b)*1024, kb=int(b)
      )
      print self.args
      self.player_type = self.args['t']
      self.source_type = "multicast"
      self.channel_name = None
      if file:
        from config import file_browser_path
        if os.path.abspath(file).startswith(file_browser_path):
          self.source_type = "file"
      self.pipeline = Gst.Pipeline()
      self.cmd = (
        GST_SOURCES[self.source_type]
        +
        GST_ENCODERS[self.player_type]
      ) % self.args
      self.cmd = self.reg_strip_comments.sub("", self.cmd)
      print self.cmd.strip()
      self.playbin = Gst.parse_launch(self.cmd)
      self.src = self.playbin.get_by_name("src")
      self.sink = self.playbin.get_by_name("sink")
      self.sink_factory = self.sink.get_factory().get_name()
      self.clients = {}
      self.pipeline.add(self.playbin)
      self.pipeline.set_state(Gst.State.PAUSED)
      self.connect()
  def connect(self):
      self.sink.connect("client-added", self.add_client)
      if self.sink_factory=="multisocketsink":
        self.sink.connect("client-socket-removed", self.remove_client)
      else:
        self.sink.connect("client-fd-removed", self.remove_client)
  def handles(self):
      return self.sink.get_property("num-handles")
  def add_client(self, sink, gsock):
      print "Adding client:", gsock
      self.pipeline.set_state(Gst.State.PLAYING)
  def remove_client(self, sink, gsock):
      print "Removing socket:", gsock
      sink.emit("remove-flush", gsock)
      if gsock in self.clients:
        del self.clients[gsock]
      if self.handles()==0:
        print "Last client, stopping stream ..."
        sink.emit("clear")
        if self.pipeline:
          self.pipeline.set_state(Gst.State.NULL)
        self.src.set_state(Gst.State.NULL)
        self.pipeline = None
        #print self.src.get_state(10L**9) # 0.1 ns
  def remove_all_clients(self):
      for client in self.clients.keys():
        self.remove_client(self.sink, client)
  def handle(self, wfile, args={}):
      if self.sink_factory=="multisocketsink":
        client_id = Gio.Socket().new_from_fd(wfile.fileno())
      else: # multifdsink
        pipe = os.pipe()
        client_id = pipe[1]
        args["pipe"] = pipe
      self.sink.emit("add", client_id)
      self.clients[client_id] = args
      self.clients[client_id]["start_time"] = time.time()
      if self.channel_name is None and "channel_name" in args:
        self.channel_name = urllib.unquote_plus(args["channel_name"])
      if self.sink_factory=="multifdsink":
        pipe = self.clients[client_id]["pipe"]
        while client_id in self.clients:
          data = ''
          while len(data)<self.bufsize:
            data += os.read(pipe[0], self.bufsize)
          #print "Data:", len(data)
          #self.wfile.write("%x\r\n" % len(data)+data+"\r\n")
          wfile.write(data)
          wfile.flush()
        os.close(pipe[0])
        os.close(pipe[1])
        self.sink.emit("remove-flush", client_id)
      else: # multisocketsink
        while client_id in self.clients:
          #print "Waiting..."
          time.sleep(1)
  def get_cmd(self):
      return self.reg_strip_spaces.sub(" ", self.cmd)
  def client_info(self):
      for client, data in self.clients.items():
        stats = self.sink.emit("get-stats", client)
        ns = 1000000000.0 # nanosecond
        yield dict(
          channel = urllib.unquote_plus(data["channel_name"]),
          time = hms((time.time()-data["start_time"])),
          client = client,
          addr = data["addr"],
          port = data["port"],
          user_agent = data["user_agent"],
          bytes_sent = stats["bytes-sent"],
          bytes_sent_s = kmg(stats["bytes-sent"]),
          buffers_dropped = stats["buffers-dropped"],
          last_activity = time.time()-stats["last-activitity-time"]/ns,
          avg_speed = stats["bytes-sent"]/(stats["connect-duration"]/ns)*8,
          avg_speed_s = kmg(
            stats["bytes-sent"]/(stats["connect-duration"]/ns)*8, "b/s"
          )
        )
  def guess_content_type(self):
      mux = self.playbin.get_by_name("mux")
      if not mux:
        # possible a raw stream
        if self.args['t']=="rawaudio":
          return "audio/mp3"
        elif self.args['t']=="rawvideo":
          return "video/mpeg"
        return "application/octet-stream"
      muxname = mux.get_path_string().lower()
      if 'webmmux' in muxname:
        return "video/webm"
      elif 'mpegtsmux' in muxname:
        return "video/mpeg"
      elif 'mp4mux' in muxname:
        return "video/mp4"
      elif 'oggmux' in muxname:
        if self.args['t']=="ogv":
          return "video/ogg"
        else:
          return "audio/ogg"
      elif 'mp3' in muxname:
        return "audio/mp3"
      print "UNKNOWN MUX:", muxname
      return "application/octet-stream"

class players_class(dict):
  def cleanup(self):
      running_players = 0
      connected_clients = 0
      # remove unused players
      for player_path in self.keys():
        player = dict.__getitem__(self, player_path)
        if not player.pipeline:
          print "Cleaning player instance:", player
          del self[player_path]
        else:
          running_players += 1
        connected_clients += len(player.clients)
      print "Still running: %d player(s), %d client(s)" \
            % (running_players, connected_clients)
  def __getitem__(self, path):
      self.cleanup()
      # find an runnnig player
      if path in self.keys():
        p = dict.__getitem__(self, path)
        return p
      args = list([
        x
        for x in path.strip("&").split("&")
        if "=" not in x
      ])
      kwargs = dict([
        x.split("=", 1)
        for x in path.strip("&").split("&")
        if '=' in x and not x.startswith("tmp=")
      ])
      key = "_".join(args+kwargs.values())
      if key not in self.keys():
        self[key] = player_class(*args, **kwargs)
      return dict.__getitem__(self, key)
players = players_class()

class request_handler(AuthHTTPRequestHandler):
  realm = "VideoRelay"
  server_version = "%s/%s %s" \
    % (realm, __version__, AuthHTTPRequestHandler.server_version)
  from config import users
  def info(self):
      loads = open("/proc/loadavg").read().split(" ")
      yield "Threads: %d, Server load: %4.2f %4.2f %4.2f, Uptime: %s" % (
        len(threading.enumerate()),
        float(loads[0]),
        float(loads[1]),
        float(loads[2]),
        dhm(time.time()-server_start)
      )
      #players.cleanup()
      for player in players.keys():
        p = players[player]
        yield "Player: %s: '%s', %d client(s)" \
              % (player, p.channel_name, p.handles())
        yield "  "+p.get_cmd().strip()
        for client in p.client_info():
          yield "  client: %(time)s, %(addr)s:%(port)d [%(user_agent)s]" \
                % client
          yield "          bytes: %(bytes_sent_s)s, avg speed: %(avg_speed_s)s, " \
                "last activity: %(last_activity)3.1fs, dropped=%(buffers_dropped)d" \
                % client
  def debug(self):
      import IPython
      IPython.embed()
  def send_content_type(self, ctype):
      print "Content-type:", ctype
      self.send_header("Content-type", ctype)
  def do_GET(self):
      if self.client_address[0] not in allowed_ips:
        if "basic" in auth:
          if not self.check_basic_auth():
            return
        elif "digest" in auth:
          if not self.check_digest_auth():
            return
        allowed_ips.append(self.client_address[0])
      print "Client request:", self.remote_user, self.path, self.client_address
      if not re.compile(r"^[?/&A-Za-z0-9=.%+-]*$").search(self.path):
        self.send_error(403) # Forbidden
        return
      if self.path=="/favicon.ico":
        self.send_error(404) # Not found
        return
      if self.path[1:] in html.FILES:
        self.send_response(200)
        self.send_header("Content-type", html.FILES[self.path[1:]])
        self.end_headers()
        if self.path=="/":
          self.wfile.write(
            html.page().index("http://%s:%s" % (socket.gethostname(), PORT))
          )
        elif self.path=="/playlist.m3u":
          self.wfile.write(
            html.page().m3u("http://%s:%s" % (socket.gethostname(), PORT))
          )
        else:
          self.wfile.write(open(self.path[1:]).read())
        return
      if self.path=="/info":
        self.send_response(200)
        self.send_header("Content-type", "text/plain")
        self.end_headers()
        self.wfile.write("\n".join(self.info()))
        return
      elif self.path=="/debug":
        self.send_error(404) # Not found
        self.debug()
        return
      if '&' not in self.path:
        if '//' in self.path:
          self.path, channel_name = self.path.rsplit("//", 1)
        else:
          channel_name = None
        self.path = self.path.replace("/", '&')
      if self.path.startswith("/?"):
        self.path = self.path[2:]
      player = players[self.path]
      self.send_response(200)
      self.send_content_type(player.guess_content_type())
      # This is real stream. Don't cache!
      self.send_header("Cache-Control", "no-cache, no-store, must-revalidate")
      self.send_header("Pragma", "no-cache")
      self.send_header("Expires", "0")
      self.send_header("Vary", "*")
      self.send_header("Etag", str(time.time()).replace(".", ""))
      self.send_header("Connection", "close")
      #self.send_header("Content-length", str(1024**4)) # 1TB
      #self.send_header("Transfer-encoding", "chunked")
      self.end_headers()
      start = time.time()
      try:
        player.handle(self.wfile, dict(
          addr = self.client_address[0],
          port = self.client_address[1],
          user_agent = self.headers.getheader("User-agent"),
          path = self.path,
          channel_name = channel_name
        ))
      except socket.error, err:
        print "Socket error:", err
        # Socket is already closed.
        # Simulate an open file to allow flushing data.
        self.wfile = StringIO()
      print "Request end: %s - %4.2f minutes" % (
        self.client_address, (time.time()-start)/60)
      players.cleanup()
  def do_OPTIONS(self):
      self.wfile.write("200 OK\r\n\r\n")

if __name__ == "__main__":
  # change directory to script
  os.chdir(os.path.dirname(sys.argv[0]))
  try:
    from config import PORT
    from config import HOST
  except:
    pass # use default HOST, PORT
  if HOST=="":
    HOST = socket.gethostname()
  # run server
  httpd = ThreadingHTTPServer((HOST, PORT), request_handler)
  print "Running http server for http://%s:%d/" % (HOST, PORT)
  server_start = time.time()
  try:
    httpd.serve_forever()
  except KeyboardInterrupt:
    print "\rExiting ..."
    # remove clients
    for player in players.keys():
      players.get(player).remove_all_clients()
    # stop all threads
    for thread in threading.enumerate():
      try:
        thread._Thread__stop()
      except:
        print thread.getName(), 'could not be terminated'