#!/usr/bin/env python
# coding: utf-8

# # Program
# Transfer real-time stream to media server from onvif device

# # History
# * 2019/12/18 first release.
# * 2019/12/19 insert payload argument, 'host' change to 'host_dst' and 'host_src' insert to payload.
# * 2019/12/19 insert argument to ffmpeg command, -loglevel
# * 2019/12/19 insert argument to ffmpeg command, -report
# * 2019/12/19 insert environmental variables and insert argument to ffmpeg command, -r

# # FlowChart
# 1. Check required arguments
# 2. Test connectivity of source and destination device.
# 3. Create subprocess push real-time stream to media server by ffmpeg
# 4. Get stream list by send request to media server

# # Dockerfile
#     FROM ubuntu:18.04
#
#     RUN apt update
#
#     RUN apt install -y python3-dev python3-pip
#
#     RUN apt install -y ffmpeg
#
#     RUN pip3 install pytz
#
#     RUN pip3 install requests
#
#     RUN pip3 install paho-mqtt
#
#     WORKDIR /root
#
#     ENV MQTT_HOST "" MQTT_PORT "1883" MQTT_TOPIC ""
#
#     ENV RTSP_HOST "" RTSP_PORT="554" RTSP_KEY "" RTSP_USER "" RTSP_PSWD ""
#
#     ENV EADA_HOST "" EADA_PORT "554" EADA_API "10008"
#
#     ENV HEARTBEAT "" TIMEOUT "60"
#
#     COPY run.py run.py
#
#     CMD python3 -u run.py

# In[1]:


from os import environ as env


# In[3]:


from datetime import datetime


# In[4]:


from paho.mqtt.client import Client as mqttclient


# In[5]:


import time, json, pytz, socket, subprocess, requests


# In[6]:


mqtt = mqttclient()


# In[7]:


def publish(tag, message, host=None, port=None, topic=None, payload=None):
    now = datetime.now(tz=pytz.timezone('Asia/Taipei')).strftime('%Y/%m/%d %H:%M:%S.%L')

    if topic and payload:
        if not mqtt.is_connected: mqtt.reconnect()

        payload['time'] = now

        payload['message'] = "{0} {1}".format(tag, message)

        mqtt.publish(topic=topic, payload=json.dumps(payload), qos=0)

    print("[{0}][{1}][{2}:{3}]: {4}".format(now, tag, host, port, message) if host and port else "[{0}][{1}]: {4}".format(now, tag, message))


# In[8]:


def ping(tag, host, port, topic=None, payload=None):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    try:
        s.connect((host, port))

        s.shutdown(socket.SHUT_RDWR)

        publish(tag, "connection test completed", host, port, topic, payload)

        return True
    except ConnectionRefusedError:
        publish(tag, "connection test refused", host, port, topic, payload)
    except:
        publish(tag, "connection test error with unknown problem", host, port, topic, payload)
    return False


# In[9]:


def api_request(host, port, topic):
    req = requests.get("http://{0}:{1}/api/v1/pushers".format(host, port), params={ 'q': topic })

    json = req.json()

    return json['rows'] if req.status_code == requests.codes.ok and int(json['total']) > 0 else []


# In[10]:


payload = { "time": "", 'status': 0, 'host_src': "", 'host_dst': "", 'port': "", 'message': "" }

check_mqtt = 'MQTT_HOST' in env and 'MQTT_PORT' in env and 'MQTT_TOPIC' in env and len(env['MQTT_TOPIC']) > 0

check_rtsp = 'RTSP_HOST' in env and 'RTSP_PORT' in env and 'RTSP_KEY' in env and len(env['RTSP_KEY']) > 0

check_eada = 'EADA_HOST' in env and 'EADA_PORT' in env and 'EADA_API' in env

heartbeat = int(env['HEARTBEAT']) if 'HEARTBEAT' in env else 30

timeout = int(env['TIMEOUT']) if 'TIMEOUT' in env else 60

# In[ ]:

if check_mqtt and ping("MQTT", env['MQTT_HOST'], int(env['MQTT_PORT'])):
    mqtt_host = env['MQTT_HOST']

    mqtt_port = int(env['MQTT_PORT'])

    mqtt_topic = env['MQTT_TOPIC']

    mqtt.connect(host=mqtt_host, port=mqtt_port)

    if check_eada and ping("EasyDarwin", env['EADA_HOST'], int(env['EADA_PORT']), mqtt_topic, payload):
        out_host = env['EADA_HOST']

        out_port = int(env['EADA_PORT'])

        # out_fps = int(env['EADA_FPS']) if 'EADA_FPS' in env else 30

        rtsp_out = 'rtsp://{0}:{1}/{2}'.format(out_host, out_port, mqtt_topic)

        payload['host_dst'] = out_host

        payload['port'] = out_port

        if check_eada and ping("EasyDarwin", env['EADA_HOST'], int(env['EADA_API']), mqtt_topic, payload):
            url_host = env['EADA_HOST']

            url_port = int(env['EADA_API'])

            if check_rtsp and ping("RTSP", env['RTSP_HOST'], int(env['RTSP_PORT']), mqtt_topic, payload):
                in_host = env['RTSP_HOST']

                in_port = int(env['RTSP_PORT'])

                in_key = env['RTSP_KEY']

                # in_fps = int(env['RTPS_FPS']) if 'RTPS_FPS' in env else 30

                in_user = env['RTSP_USER'] if 'RTSP_USER' in env else ""

                in_pswd = env['RTSP_PSWD'] if 'RTSP_PSWD' in env else ""

                rtsp_in = 'rtsp://{0}:{1}@'.format(in_user, in_pswd) if len(in_user) > 0 and len(in_pswd) > 0 else 'rtsp://'

                rtsp_in = '{0}{1}:{2}/{3}'.format(rtsp_in, in_host, in_port, in_key)

                ffmpeg = [
                    'ffmpeg',                 # Transfer real-time stream to media server
                    '-report',                # Ouput ffmpeg info
                    '-loglevel', 'quiet',     # Hide ffmpeg extra info
                    '-rtsp_transport', 'tcp', # The protocol of source stream
                    # '-r', str(in_fps),        # Input video frame count
                    '-i', rtsp_in,            # The URL of source stream
                    '-c:v', 'copy',           # Output video encoding
                    '-an',                    # Output audio encoding 指定輸出聲音編碼 (an: 無聲音)
                    '-rtsp_transport', 'tcp', # The protocol of destination stream
                    # '-r', str(out_fps),       # Output video frame count
                    '-f','rtsp', rtsp_out,    # The URL of destination stream
                ]

                p = subprocess.Popen(ffmpeg)

                payload['host_src'] = in_host

                payload['status'] = 0

                publish(mqtt_topic, "Start", out_host, out_port, mqtt_topic, payload)

                time.sleep(timeout)

                while (p.returncode == None):
                    rows = api_request(url_host, url_port, mqtt_topic)

                    if len(rows) > 0:
                        payload['status'] = 1

                        message = "In {0} / Out {1}".format(rows[0]['inBytes'], rows[0]['outBytes'])

                        publish(mqtt_topic, message, out_host, out_port, mqtt_topic, payload)
                    else:
                        break

                    time.sleep(heartbeat)

                payload['status'] = 0

                publish(mqtt_topic, "Stop", out_host, out_port, mqtt_topic, payload)

                p.kill()
            else:
                publish("RTSP", "parameter error", topic=mqtt_topic, payload=payload)
        else:
            publish("EasyDarwin API", "parameter error", topic=mqtt_topic, payload=payload)
    else:
        publish("EasyDarwin", "parameter error", topic=mqtt_topic, payload=payload)
else:
    publish("MQTT", "parameter error")

