#!/usr/bin/python
# SPDX-License-Identifier: CC0
#
# Written by: Ferass El Hafidi <vitali64pmemail@protonmail.com>
#
from sys import exit
import json
import string
import socket
import requests
import miniirc
import hashlib
import yaml
import re
from time import sleep
from requests.adapters import HTTPAdapter, Retry

puppets = {}

def connect_new_puppet(nickname, user_id, gaccount):
    if not puppets.get(gaccount):
        puppets[gaccount] = {}

    try:
        print("* Puppet %s already connected, skipping" % puppets[gaccount][user_id])

        if puppets[gaccount][user_id].connected == True:
            return puppets[gaccount][user_id]
    except KeyError:
        pass

    # If there's too many puppets, disconnect the oldest one
    if len(puppets[gaccount]) == client_limit:
        puppets[gaccount][0].disconnect(msg="Connection closed for inactivity")
        puppets[gaccount].pop(0)

    # Sanitize nickname
    allowed_chars = string.digits + string.ascii_letters + "^|\\-_[]{}"
    sanitized_nickname = nickname
    for char in sanitized_nickname:
        if char not in allowed_chars:
            # Try to represent illegal characters with similar, allowed ones
            if char == "(" or char == "<":
                sanitized_nickname = sanitized_nickname.replace(char, "[")
            elif char == ")" or char == ">":
                sanitized_nickname = sanitized_nickname.replace(char, "]")
            elif char == "/":
                sanitized_nickname = sanitized_nickname.replace(char, "|")
            elif char == "." or char == " ":
                sanitized_nickname = sanitized_nickname.replace(char, "_")
            else:
                sanitized_nickname = sanitized_nickname.replace(char, "-")
    if sanitized_nickname[0] in string.digits + '-':
        sanitized_nickname = "_" + sanitized_nickname

    if (sanitized_nickname.endswith("Serv") or "." in nickname) and gaccount.startswith("irc."):
        print("* Not creating puppet of a network service/server (%s)", sanitized_nickname)
        return

    print("* Connecting new puppet %s (user_id %s) as %s" % (nickname,
            user_id, sanitized_nickname))

    puppets[gaccount][user_id] = miniirc.IRC(
        irc_host,
        irc_port,
        sanitized_nickname,
        realname="%s [%s]" % (user_id, gaccount),
        auto_connect=False,
        debug=False,
        ssl=irc_tls,
        persist=False,
        quit_message="quit"
    )
    puppets[gaccount][user_id].Handler('PRIVMSG', colon=False)(fail_on_pm)
    userid_hash = hashlib.sha3_224(bytes(user_id, "utf-8"))
    write_ident_file(str(userid_hash.hexdigest()))

    puppets[gaccount][user_id].connect()

    while puppets[gaccount][user_id].connected != True:
        pass

    return puppets[gaccount][user_id]

def write_ident_file(ident):
    if ident_enabled:
        print("* Writing to ident file for oidentd/ident2")
        with open(ident_file, 'w') as file:
            file.write(ident_fmt.replace("%user%", ident))

def get_puppet_client(gaccount, user_id):
    try:
        ret = puppets[gaccount][user_id]
    except KeyError:
        print("* Error finding puppet client %s (%s)" % (user_id, gaccount))
        return None
    return ret

def on_irc_msg(irc, hostmask, args):
    nick    = hostmask[0]
    channel = args[0]
    message = args[-1]

    for gw in puppets:
        for puppet in puppets[gw]:
            if puppets[gw][puppet].nick == nick:
                return # don't act on own messages

    # Get on which gateway it was sent
    for gw in gateways:
        if gateways[gw] == channel:
            gateway = gw
            break
    else:
        return

    # Craft API message
    api_message = {
        "text":     message,
        "username": nick,
        "userid":   hostmask[1] + "@" + hostmask[2],
        "gateway":  gateway
    }

    print("* Sending messages to gateway")

    req = requests.post(f"{api_host}/api/message", json=api_message)


def api_loop():
    away = False
    sess = requests.Session()
    sess_retries = Retry(total=10000,
                    backoff_factor=0.1,
                    status_forcelist=[ 500, 502, 503, 504 ])

    sess.mount('http://', HTTPAdapter(max_retries=sess_retries))
    sess.mount('https://', HTTPAdapter(max_retries=sess_retries))

    while True:
        req = sess.get(f"{api_host}/api/stream", stream=True)

        try:
            for line in req.iter_lines():
                if line:
                    # Unset away status for all puppets
                    if away:
                        for gw in puppets:
                            for puppet in puppets[gw]:
                                puppets[gw][puppet].send("AWAY")
                        away = False
                    message = json.loads(line.decode('utf-8'))

                    if message.get("message") == "Not Found":
                        print("Error: Channel not found: %s" % channel)
                        return -1

                    if message["event"] == "api_connected":
                        print("Connected!")
                    elif message["event"] == "" or message["event"] == "user_action": # XXX: Probably a message
                        channel = gateways[message["gateway"]]
                        print("Bridging message event")
                        # Connect puppet if it does not exist already
                        puppet = connect_new_puppet(message["username"], message["userid"], message["account"])

                        if puppet:
                            # Have the puppet join the channel if it isn't joined already
                            puppet.send("JOIN", channel)

                            # Parse message for long code blocks
                            if paste_enabled:
                                code_blocks = re.findall("^\\`\\`\\`(.*?)^\\`\\`\\`", message["text"], flags=re.S + re.M)
                                for code_block in code_blocks:
                                    if code_block.count("\n") > paste_maxlines:
                                        cb_checksum = hashlib.sha3_224(bytes(code_block, "utf-8"))
                                        file = cb_checksum.hexdigest() + ".txt"
                                        message["text"] = message["text"].replace("```" + code_block + "```", "%s%s" % (paste_domain, file))
                                        with open(paste_dir + "/" + file, mode='w') as file_paste:
                                            print(code_block, file=file_paste)
                                    else:
                                        # Remove trailing ```'s
                                        message["text"] = message["text"].replace("```" + code_block + "```", code_block)

                            # Send the message
                            for line in message["text"].split("\n"):
                                if len(line) >= irc_msglen:
                                    # Hacky...
                                    buf = ""
                                    i = 0
                                    for char in line:
                                        i += 1
                                        if len(buf) < irc_msglen and i != (len(line) - 1):
                                            buf += char
                                        elif len(buf) == irc_msglen or i == (len(line) - 1):
                                            puppet.msg(channel, buf)
                                            buf = ""
                                else:
                                    if line != "":
                                        if message["event"] == "user_action":
                                            puppet.me(channel, line)
                                        else:
                                            puppet.msg(channel, line)
                                sleep(0.5)
                    elif message["event"] == "join":
                        channel = gateways[message["gateway"]]
                        print("Bridging join event")
                        # Connect puppet if it does not exist already
                        puppet = connect_new_puppet(message["username"], message["userid"], message["account"])

                        if puppet:
                            # Have the puppet join the channel
                            puppet.send("JOIN", channel)
                    elif message["event"] == "leave":
                        channel = gateways[message["gateway"]]
                        # Check if puppet exists
                        puppet = get_puppet_client(message["account"], message["userid"])

                        if puppet == None:
                            # We have nothing to do
                            continue
                        puppet.send("PART", channel, "Leaving")
        except (requests.exceptions.ConnectionError, 
                requests.exceptions.ChunkedEncodingError):
            print("Disconnected from Matterbridge")

            # Set away status for all puppets
            if away == False:
                away = True
                for gw in puppets:
                    for puppet in puppets[gw]:
                        puppets[gw][puppet].send("AWAY", "Matterbridge disconnected")
            continue

def fail_on_pm(irc, hostmask, args):
    nick    = hostmask[0]
    target  = args[0]

    if "\x01" in args[1]:
        return

    if target == irc.nick:
        irc.msg(nick, "[Automated message] This is a bridged puppet, private messages are currently unsupported - message not sent.")

def on_irc_kick(irc, hostmask, args):
    print("* Got kick")
    if irc.nick == hostmask[0]:
        print("* Trying to rejoin channel")
        irc.join(args[0])

def clear_ident_file(irc, hostmask, args):
    if ident_enabled:
        print("* Clearing ident file")
        # Clear ident file
        with open(ident_file, 'w') as file:
            file.write("\n")

def main():
    # Set CTCP VERSION reply
    miniirc.version = "Matterpuppeter · https://git.vitali64.duckdns.org/utils/matterpuppeter.git"

    # Set handlers
    miniirc.Handler('001')(clear_ident_file) # Global handler
    # Write ident for the bot
    write_ident_file(ident_bot)

    # Connect to IRC
    irc = miniirc.IRC(
        irc_host,
        irc_port,
        irc_nick,
        realname=irc_gecos,
        debug=False,
        ssl=irc_tls,
        ns_identity=sasl_auth if sasl_enabled else None,
        quit_message="shutting down",
        auto_connect=True,
        persist=True
    )
    irc.Handler('PRIVMSG', colon=False)(on_irc_msg)
    irc.Handler('NOTICE', colon=False)(on_irc_msg)
    irc.Handler('KICK', colon=False)(on_irc_kick)


    for gateway in gateways:
        irc.send("JOIN", gateways[gateway])

    api_loop()

if __name__ == '__main__':
    with open("config.yaml") as file:
        config = yaml.load(file, Loader=yaml.FullLoader)
        try:
            mb_account     = config["api"]["account"]
            irc_host       = config["irc"]["host"]
            irc_port       = config["irc"]["port"]
            irc_tls        = config["irc"]["tls"]
            irc_nick       = config["irc"]["nick"]
            irc_gecos      = config["irc"]["gecos"]
            irc_msglen     = config["irc"]["message_limit"]
            api_host       = config["api"]["host"]
            ident_enabled  = config["ident"]["enable"]
            if ident_enabled:
                ident_file     = config["ident"]["file"]
                ident_fmt      = config["ident"]["format"]
                ident_bot      = config["ident"]["username"]
            client_limit   = config["irc"]["client_limit"]
            paste_enabled  = config["paste"]["enable"]
            if paste_enabled:
                paste_dir      = config["paste"]["dir"]
                paste_domain   = config["paste"]["domain"]
                paste_maxlines = config["paste"]["maxlines"]
            sasl_enabled   = config["irc"]["sasl"]["enable"]
            if sasl_enabled:
                sasl_auth  = (config["irc"]["sasl"]["username"],
                              config["irc"]["sasl"]["password"])
            gateways       = config["gateway"]
        except KeyError:
            print("Error: Configuration file is missing options")
            exit(1)
    main()