import json
import socket # Just so we can properly handle hostname exceptions
import obspython as obs
import paho.mqtt.client as mqtt
import ssl
import pathlib


# Meta
__version__ = '1.0.2'
__version_info__ = (1, 0, 1)
__license__ = "AGPLv3"
__license_info__ = {
    "AGPLv3": {
        "product": "update_mqtt_status_homeassistant",
        "users": 0, # 0 being unlimited
        "customer": "Unsupported",
        "version": __version__,
        "license_format": "1.0",
    }
}
__author__ = 'HeedfulCrayon'

__doc__ = """\
Publishes real-time OBS status info to the given MQTT server/port/channel \
at the configured interval.
"""

# Default values for the configurable options:
INTERVAL = 5 # Update interval (in seconds)
MQTT_HOST = "localhost" # Hostname of your MQTT server
MQTT_USER = ""
MQTT_PW = ""
MQTT_PORT = 1883 # Default MQTT port is 1883
MQTT_BASE_CHANNEL = ""
MQTT_SENSOR_NAME = "obs"
CONTROL = False

# This is how we keep track of the current status:
STATE = "Initializing"
PREV_STATE = STATE
STATUS = {
    "recording": False,
    "streaming": False,
    "paused": False,
    "replay_buffer": False, # If it's active or not
    "fps": 0,
    "frame_time_ns": 0,
    "frames": 0,
    "lagged_frames": 0,
}
PREV_STATUS = STATUS.copy()

# MQTT Event Functions
def on_mqtt_connect(client, userdata, flags, rc):
    """
    Called when the MQTT client is connected from the server.  Just prints a
    message indicating we connected successfully.
    """
    print("MQTT connection successful")
    
    set_homeassistant_config()

def on_mqtt_disconnect(client, userdata, rc):
    """
    Called when the MQTT client gets disconnected.  Just logs a message about it
    (we'll auto-reconnect inside of update_status()).
    """
    print("MQTT disconnected.  Reason: {}".format(str(rc)))

def on_mqtt_message(client, userdata, message):
    """
    Handles MQTT messages that have been subscribed to
    """
    topic = pathlib.PurePosixPath(message.topic)
    base_topic = str(topic.parent.parent)
    message_type = topic.parent.stem # Stream or Record
    payload = str(message.payload.decode("utf-8"))
    print(message.topic + ":" + payload)
    start_action(base_topic, message_type, payload)

# OBS Script Function Exports
def script_description():
    return __doc__ # We wrote a nice docstring...  Might as well use it!

def script_load(settings):
    """
    Just prints a message indicating that the script was loaded successfully.
    """
    global STATE
    global PREV_STATE
    print("MQTT script loaded.")
    STATE = "Initializing"
    PREV_STATE = "Off"

def script_unload():
    """
    Publishes a final status message indicating OBS is off
    (so your MQTT sensor doesn't get stuck thinking you're
    recording/streaming forever) and calls `CLIENT.disconnect()`.
    """
    global STATE
    STATE = "Off"
    if CLIENT.is_connected():
        publish_state_and_attributes()
        CLIENT.disconnect()
    CLIENT.loop_stop()

def script_defaults(settings):
    """
    Sets up our default settings in the OBS Scripts interface.
    """
    obs.obs_data_set_default_string(settings, "mqtt_host", MQTT_HOST)
    obs.obs_data_set_default_string(settings, "mqtt_user", MQTT_USER)
    obs.obs_data_set_default_string(settings, "mqtt_pw", MQTT_PW)
    obs.obs_data_set_default_string(settings, "mqtt_base_channel", MQTT_BASE_CHANNEL)
    obs.obs_data_set_default_string(settings, "mqtt_sensor_name", MQTT_SENSOR_NAME)
    obs.obs_data_set_default_int(settings, "mqtt_port", MQTT_PORT)
    obs.obs_data_set_default_int(settings, "interval", INTERVAL)
    obs.obs_data_set_default_bool(settings, "controllable", CONTROL)

def script_properties():
    """
    Makes this script's settings configurable via OBS's Scripts GUI.
    """
    props = obs.obs_properties_create()
    obs.obs_properties_add_text(props, "mqtt_host", "MQTT server hostname", obs.OBS_TEXT_DEFAULT)
    obs.obs_properties_add_text(props, "mqtt_user", "MQTT username", obs.OBS_TEXT_DEFAULT)
    obs.obs_properties_add_text(props, "mqtt_pw", "MQTT password", obs.OBS_TEXT_PASSWORD)
    obs.obs_properties_add_text(props, "mqtt_base_channel", "MQTT Base channel",obs.OBS_TEXT_DEFAULT)
    obs.obs_properties_add_text(props, "mqtt_sensor_name", "MQTT Sensor Name",obs.OBS_TEXT_DEFAULT)
    obs.obs_properties_add_int(props, "mqtt_port", "MQTT TCP/IP port", MQTT_PORT, 65535, 1)
    obs.obs_properties_add_int(props, "interval", "Update Interval (seconds)", 1, 3600, 1)
    obs.obs_properties_add_bool(props, "controllable", "Control Streaming/Recording via MQTT")
    return props

def script_update(settings):
    """
    Applies any changes made to the MQTT settings in the OBS Scripts GUI then
    reconnects the MQTT client.
    """
    # Apply the new settings
    global MQTT_HOST
    global MQTT_USER
    global MQTT_PW
    global MQTT_PORT
    global MQTT_BASE_CHANNEL
    global MQTT_SENSOR_NAME
    global INTERVAL
    global CONTROL
    mqtt_host = obs.obs_data_get_string(settings, "mqtt_host")
    if mqtt_host != MQTT_HOST:
        MQTT_HOST = mqtt_host
    mqtt_user = obs.obs_data_get_string(settings, "mqtt_user")
    if mqtt_user != MQTT_USER:
        MQTT_USER = mqtt_user
    mqtt_pw = obs.obs_data_get_string(settings, "mqtt_pw")
    if mqtt_pw != MQTT_PW:
        MQTT_PW = mqtt_pw
    mqtt_base_channel = obs.obs_data_get_string(settings, "mqtt_base_channel")
    if mqtt_base_channel != MQTT_BASE_CHANNEL:
        MQTT_BASE_CHANNEL = mqtt_base_channel
    mqtt_sensor_name = obs.obs_data_get_string(settings, "mqtt_sensor_name")
    if mqtt_sensor_name != MQTT_SENSOR_NAME:
        MQTT_SENSOR_NAME = mqtt_sensor_name
    mqtt_port = obs.obs_data_get_int(settings, "mqtt_port")
    if mqtt_port != MQTT_PORT:
        MQTT_PORT = mqtt_port
    INTERVAL = obs.obs_data_get_int(settings, "interval")
    CONTROL = obs.obs_data_get_bool(settings, "controllable")

    # Disconnect (if connected) and reconnect the MQTT client
    CLIENT.disconnect()
    try:
        if MQTT_PW != "" and MQTT_USER != "":
            CLIENT.username_pw_set(MQTT_USER, password=MQTT_PW)
        CLIENT.connect_async(MQTT_HOST, MQTT_PORT, 60)

        # Publish our initial state
        publish_state_and_attributes()
    except (socket.gaierror, ConnectionRefusedError) as e:
        print("NOTE: Got a socket issue: %s" % e)
        pass # Ignore it for now

    # Remove and replace the timer that publishes our status information
    obs.timer_remove(update_status)
    obs.timer_add(update_status, INTERVAL * 1000)
    CLIENT.loop_start()

# Event Helper Functions
def set_homeassistant_config():
    """
    Sends initial configuration state and attributes topic
    for autodiscovery in Home Assistant
    """
    initMqttChannel = MQTT_BASE_CHANNEL + "/sensor/" + MQTT_SENSOR_NAME
    initConfig = {
        "name": "OBS",
        "unique_id": MQTT_SENSOR_NAME,
        "state_topic": initMqttChannel + "/state",
        "json_attributes_topic": initMqttChannel + "/attributes"
    }

    CLIENT.publish(initMqttChannel + "_state/config", json.dumps(initConfig))

    if CONTROL:
        # Set up switches for autodiscovery
        mqttChannel = MQTT_BASE_CHANNEL + "/switch/" + MQTT_SENSOR_NAME
        controlConfig = {
            "_stream": {
                "name": MQTT_SENSOR_NAME + " Stream",
                "unique_id": MQTT_SENSOR_NAME + "_stream",
                "state_topic": mqttChannel + "/stream/state",
                "command_topic": mqttChannel + "/stream/set",
                "payload_on": "ON",
                "payload_off": "OFF",
                "icon": "mdi:broadcast"
            },
            "_record": {
                "name": MQTT_SENSOR_NAME + " Record",
                "unique_id": MQTT_SENSOR_NAME + "_record",
                "state_topic": mqttChannel + "/record/state",
                "command_topic": mqttChannel + "/record/set",
                "payload_on": "ON",
                "payload_off": "OFF",
                "icon": "mdi:record"
            }
        }
        for config in controlConfig:
            CLIENT.publish(mqttChannel + config + "/config", json.dumps(controlConfig[config]))
        
        CLIENT.subscribe(mqttChannel + "/stream/set")
        CLIENT.subscribe(mqttChannel + "/record/set")
        print("Subscribed to record and stream")
        CLIENT.publish(mqttChannel + "/stream/set", "OFF")
        CLIENT.publish(mqttChannel + "/record/set", "OFF")
        print("Setting Record and Stream to OFF")

def start_action(channel, message_type, payload):
    """
    Starts or stops action based on message type and payload.
    Also publishes state after action is run
    """
    if message_type == "stream":
        if payload == "ON":
            obs.obs_frontend_streaming_start()
        else:
            obs.obs_frontend_streaming_stop()
    elif message_type == "record":
        if payload == "ON":
            obs.obs_frontend_recording_start()
        else:
            obs.obs_frontend_recording_stop()

    CLIENT.publish(channel + "/" + message_type + "/state", payload)

# Helper Functions
def publish_state_and_attributes():
    """
    Publishes the STATE to the configured MQTT_BASE_CHANNEL and MQTT_SENSOR_NAME state and
    the STATUS to the configured MQTT_BASE_CHANNEL and MQTT_SENSOR_NAME attributes.
    """
    CLIENT.publish(MQTT_BASE_CHANNEL + "/sensor/" + MQTT_SENSOR_NAME + "/state", STATE)
    CLIENT.publish(MQTT_BASE_CHANNEL + "/sensor/" + MQTT_SENSOR_NAME + "/attributes", json.dumps(STATUS))

def set_state():
    """
    Updates the STATE global with the current state (Recording/Streaming/Stopped)
    """
    global STATE
    if STATUS["recording"] and STATUS["streaming"]:
        STATE = "Streaming and Recording"
    elif STATUS["streaming"]:
        STATE = "Streaming"
    elif STATUS["recording"]:
        STATE = "Recording"
    else:
        STATE = "Stopped"

def update_status():
    """
    Updates the STATE and the STATUS global with the stats of the current session.
    This info if published (JSON-encoded) to the configured MQTT_HOST/MQTT_PORT/MQTT_BASE_CHANNEL.
    Meant to be called at the configured INTERVAL.
    """
    global PREV_STATUS
    global STATUS
    global PREV_STATE
    global STATE
    STATUS["recording"] = obs.obs_frontend_recording_active()
    STATUS["streaming"] = obs.obs_frontend_streaming_active()
    STATUS["paused"] = obs.obs_frontend_recording_paused()
    STATUS["replay_buffer"] = obs.obs_frontend_replay_buffer_active()
    STATUS["fps"] = obs.obs_get_active_fps()
    STATUS["frame_time_ns"] = obs.obs_get_average_frame_time_ns()
    STATUS["frames"] = obs.obs_get_total_frames()
    STATUS["lagged_frames"] = obs.obs_get_lagged_frames()

    set_state()

    # Publish final stopped message
    if PREV_STATE != "Stopped" and STATE == "Stopped":
        print("Publishing Final Stopped Message")
        publish_state_and_attributes()
    # Publish final off message
    if PREV_STATE != "Off" and STATE == "Off":
        print("Publishing Final Off Message")
        publish_state_and_attributes()
    # Only start publishing regular messages if we're streaming or recording
    if STATUS["recording"] or STATUS["streaming"]:
        print("Publishing Status Update Message")
        publish_state_and_attributes()

    PREV_STATUS = STATUS.copy()
    PREV_STATE = STATE


# Using a global MQTT client variable to keep things simple:
CLIENT = mqtt.Client()
CLIENT.on_connect = on_mqtt_connect
CLIENT.on_disconnect = on_mqtt_disconnect
CLIENT.on_message = on_mqtt_message