#!/usr/bin/env python3

import re
import queue
import threading
import traceback
import inspect
import os
import pprint
import json
import uuid

from datetime import datetime
from time import sleep

import znc

class mssql(znc.Module):
    module_types = [znc.CModInfo.GlobalModule]
    description = "Log all channels to a Microsoft SQL Server database"
    has_args = True

    log_queue = None
    worker_threads = None
    internal_log = None
    hook_debugging = False
    debug_log_saves = False
    default_workers = 4
    num_workers = 4
    max_workers = 32
    batch_size = 50
    db_config = None
    next_worker_id = 0

    def put_log(
            self,
            code,
            channel,
            host,
            user,
            ident,
            user_mode,
            date,
            target_user,
            message,
            network):
        """
        Adds the log line to database write queue.
        """

        decolouriser = re.compile(
            r"\x1f|\x02|\x1d|\x16|\x11|\x1e|\x0f|\x03(?:\d{1,2}(?:,\d{1,2})?)?|\x04(?:\d{1,2}(?:,\d{1,2})?)?",
            re.UNICODE)

        if message is not None:
            message = message.encode('utf8', 'replace').decode('utf8')
            message = decolouriser.sub("", message)

        if self.log_queue is None:
            return

        queue_len = self.log_queue.qsize() if self.log_queue is not None else 0
        if queue_len > 10:
            with self.internal_log.error() as target:
                target.write(
                    'Queue length high: {} entries at {}\n'.format(
                        queue_len, datetime.utcnow()))

        try:
            self.log_queue.put({
                'znc_user': self.GetUser().GetUserName() if self.GetUser() is not None else None,
                'code': code,
                'channel': channel,
                'host': host,
                'user': user,
                'ident': ident,
                'user_mode': user_mode,
                'date': str(date),
                'target_user': target_user,
                'message': message,
                'network': network,
                '_id': uuid.uuid4()
            })
        except Exception as e:
            with self.internal_log.error() as target:
                target.write(
                    'Failed to enqueue log event: {0} {1}\n'.format(
                        type(e), str(e)))
                target.write('Stack trace: ' + traceback.format_exc())
                target.write('\n')

    def debug_hook(self):
        """
        Dumps parent calling method name and its arguments to debug logfile.
        """

        if (self.hook_debugging is not True
                or not self.internal_log.debug_enabled):
            return

        frameinfo = inspect.stack()[1]
        argvals = frameinfo.frame.f_locals

        with self.internal_log.debug() as target:
            target.write('Called method: ' + frameinfo.function + '()\n')
            for argname in argvals:
                if argname == 'self':
                    continue
                target.write(
                    '    ' +
                    argname +
                    ' -> ' +
                    pprint.pformat(
                        argvals[argname]) +
                    '\n')
            target.write('\n')

    def _start_workers(self, count):
        """Start `count` more worker threads on the shared queue."""
        for _ in range(count):
            idx = self.next_worker_id
            self.next_worker_id += 1
            db = MSSQLDatabase(self.db_config)
            worker_thread = threading.Thread(
                target=DatabaseThread.worker_safe,
                args=(
                    db,
                    self.log_queue,
                    self.internal_log,
                    self.debug_log_saves,
                    self.batch_size),
                daemon=True,
                name=f'DatabaseWorker-{idx}')
            worker_thread.start()
            self.worker_threads.append(worker_thread)

    def OnModCommand(self, command):
        """
        Commands: Status, AddWorkers <n>, Debug [on|off|status]
        """
        parts = str(command).split()
        if not parts:
            return
        cmd = parts[0].lower()

        if cmd == 'status':
            alive = len([t for t in self.worker_threads if t.is_alive()]
                        ) if self.worker_threads else 0
            qsize = self.log_queue.qsize() if self.log_queue is not None else 0
            self.PutModule(
                'Workers alive: {}, queue length: {}, batch size: {}, '
                'debug: {}'.format(
                    alive,
                    qsize,
                    self.batch_size,
                    'on' if self.internal_log.debug_enabled else 'off'))
        elif cmd == 'addworkers':
            try:
                n = int(parts[1])
            except (IndexError, ValueError):
                self.PutModule('Usage: AddWorkers <n>')
                return
            current = len(self.worker_threads)
            n = max(0, min(n, self.max_workers - current))
            if n == 0:
                self.PutModule(
                    'No workers added (limit is {}, currently {})'.format(
                        self.max_workers, current))
                return
            self._start_workers(n)
            self.PutModule(
                'Started {} workers, now {}'.format(
                    n, len(
                        self.worker_threads)))
        elif cmd == 'debug':
            setting = parts[1].lower() if len(parts) > 1 else 'status'
            if setting == 'status':
                state = 'on' if self.internal_log.debug_enabled else 'off'
                self.PutModule('Debug logging is {}'.format(state))
            elif setting in ('on', 'off'):
                enabled = setting == 'on'
                self.internal_log.set_debug_enabled(enabled)
                self.PutModule('Debug logging {}'.format(
                    'enabled' if enabled else 'disabled'))
            else:
                self.PutModule('Usage: Debug [on|off|status]')
        else:
            self.PutModule(
                'Commands: Status, AddWorkers <n>, Debug [on|off|status]')

    def OnLoad(self, args, message):
        """
        This module hook is called when a module is loaded.
        """
        # Initialize internal log first for debugging
        self.internal_log = InternalLog(self.GetSavePath())

        # Parse num_workers from args (default 4)
        self.num_workers = self.default_workers
        args_str = ""
        if args:
            try:
                args_str = str(args) if not hasattr(args, 's') else args.s
                if args_str:
                    self.num_workers = int(args_str.strip())
                    if self.num_workers < 1:
                        self.num_workers = 1
                    if self.num_workers > self.max_workers:
                        self.num_workers = self.max_workers
            except (ValueError, AttributeError, TypeError) as e:
                with self.internal_log.debug() as target:
                    target.write(
                        'Failed to parse args: {} (args_str: "{}", type: {})\n'.format(
                            e, args_str, type(args)))

        with self.internal_log.debug() as target:
            target.write(
                'mssql module loaded at: {}\n'.format(
                    datetime.utcnow()))
            target.write('save_path: {}\n'.format(self.GetSavePath()))
            target.write(
                'raw args: {} (type: {})\n'.format(
                    repr(args), type(args).__name__))
            target.write('parsed args_str: "{}"\n'.format(args_str))
            target.write('num_workers: {}\n'.format(self.num_workers))
            target.write('\n')

        with self.internal_log.error() as target:
            target.write(
                'mssql module loaded at: {}\n'.format(
                    datetime.utcnow()))
            target.write('hook_debugging: {}\n'.format(self.hook_debugging))
            target.write('num_workers: {}\n'.format(self.num_workers))
            target.write('\n')

        self.debug_hook()

        try:
            self.log_queue = queue.Queue()
            self.db_config = {
                'host': "dbx.spd.internal",
                'user': "_srv_znc",
                'password': "znc",
                'database': "znc",
                'port': '1433',
                'appname': 'ZNC Logger'
            }

            self.worker_threads = []
            self.next_worker_id = 0
            self._start_workers(self.num_workers)

            return True
        except Exception as e:
            message.s = str(e)

            with self.internal_log.error() as target:
                target.write(
                    'Could not initialize module caused by: {} {}\n'.format(
                        type(e), str(e)))
                target.write('Stack trace: ' + traceback.format_exc())
                target.write('\n')

            return False

    def OnUnload(self):
        if self.log_queue is None:
            return

        with self.internal_log.debug() as target:
            target.write(
                'Unloading module, draining queue length {}\n'.format(
                    self.log_queue.qsize()))

        # Send one sentinel value per worker thread
        if self.worker_threads:
            for _ in self.worker_threads:
                self.log_queue.put(None)

        self.log_queue.join()

        if self.worker_threads is not None:
            for worker_thread in self.worker_threads:
                worker_thread.join(timeout=30)

        with self.internal_log.debug() as target:
            target.write('Unload complete, queue drained\n')

    def __del__(self):
        # Terminate worker threads cleanly if possible.
        try:
            if self.log_queue is not None:
                # Send sentinel values for all workers
                if self.worker_threads:
                    for _ in self.worker_threads:
                        self.log_queue.put_nowait(None)
        except Exception:
            pass

    def OnIRCConnected(self):
        """
        This module hook is called after a successful login to IRC.
        :rtype: None
        """
        self.debug_hook()
        self.put_log("CONNECT", None, None, None, None, None,
                     datetime.now(), None, None, str(self.GetNetwork()))
        return znc.CONTINUE

    def OnIRCDisconnected(self):
        """
        This module hook is called when a user gets disconnected from IRC.
        :rtype: None
        """
        self.debug_hook()
        self.put_log("DISCONNECT", None, None, None, None, None,
                     datetime.now(), None, None, str(self.GetNetwork()))
        return znc.CONTINUE

    def OnOp(self, user, target_user, channel, noChange):
        self.debug_hook()
        self.put_log("OP",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     target_user.GetNick(),
                     None,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnChanNotice(self, nick, channel, message):
        """
        Called when we receive a channel NOTICE message from IRC.
        :type nick: CNick &
        :type channel: CChan &
        :type message: CString &
        :rtype: EModRet
        """
        self.debug_hook()
        self.put_log("NOTICE",
                     channel.GetName(),
                     nick.GetHost(),
                     nick.GetNick(),
                     nick.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnPrivNotice(self, user, message):
        self.debug_hook()
        self.put_log("NOTICE",
                     None,
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnUserNotice(self, target, message):
        """
        This module hook is called when a user sends a NOTICE message.
        :type target: CString &
        :type message: CString &
        :rtype: EModRet
        """
        self.debug_hook()
        network = self.GetNetwork()
        if network:
            self.put_log("NOTICE",
                         None,
                         target.GetHost(),
                         target.GetNick(),
                         target.GetIdent(),
                         None,
                         datetime.now(),
                         target.GetNick(),
                         message.s,
                         str(self.GetNetwork()))
        return znc.CONTINUE

    def OnDeop(self, user, target_user, channel, noChange):
        self.debug_hook()
        self.put_log("DEOP",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     target_user.GetNick(),
                     None,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnVoice(self, user, target_user, channel, noChange):
        self.debug_hook()
        self.put_log("VOICE",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     target_user.GetNick(),
                     None,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnDevoice(self, user, target_user, channel, noChange):
        self.debug_hook()
        self.put_log("DEVOICE",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     target_user.GetNick(),
                     None,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnMode(self, user, channel, mode, argument, added, noChange):
        self.debug_hook()
        mode = chr(mode)
        code = "MODE"
        if mode == "b":
            code = "BAN" if added else "UNBAN"
        else:
            sign = "+" if added else "-"
            argument = "{0}{1} {2}".format(sign, mode, argument)

        self.put_log(code,
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     argument,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnQuit(self, user, message, channels):
        self.debug_hook()
        for channel in channels:
            self.put_log("QUIT",
                         channel.GetName(),
                         user.GetHost(),
                         user.GetNick(),
                         user.GetIdent(),
                         None,
                         datetime.now(),
                         None,
                         message,
                         str(self.GetNetwork()))
        return znc.CONTINUE

    def OnNick(self, user, new_nick, channels):
        self.debug_hook()
        for channel in channels:
            self.put_log("NICK",
                         channel.GetName(),
                         user.GetHost(),
                         user.GetNick(),
                         user.GetIdent(),
                         None,
                         datetime.now(),
                         new_nick,
                         None,
                         str(self.GetNetwork()))
        return znc.CONTINUE

    def OnKick(self, user, target_nick, channel, message):
        self.debug_hook()
        self.put_log("KICK",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     target_nick,
                     message,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnJoin(self, user, channel):
        self.debug_hook()
        self.put_log("JOIN",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     None,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnPart(self, user, channel, message):
        self.debug_hook()
        self.put_log("PART",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnUserAction(self, target, message):
        self.debug_hook()
        (channel, user) = self.resolveTarget(target)
        if channel is None:
            local_nick = self.GetNetwork().GetIRCNick()
            self.put_log(
                "ME",
                None,
                local_nick.GetHost() if local_nick is not None else None,
                local_nick.GetNick() if local_nick is not None else None,
                local_nick.GetIdent() if local_nick is not None else None,
                None,
                datetime.now(),
                target.s,
                message.s,
                str(self.GetNetwork()),
            )
            return znc.CONTINUE
        return self.OnChanAction(user, channel, message)

    def OnPrivAction(self, nick, message):
        """
        Called when we receive a private CTCP ACTION ("/me" in query) from IRC.
        :type nick: CNick &
        :type message: CString &
        :rtype: EModRet
        """
        self.debug_hook()
        self.put_log("ME",
                     None,
                     nick.GetHost(),
                     nick.GetNick(),
                     nick.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnPrivMsg(self, nick, message):
        self.debug_hook()
        self.put_log("MSG",
                     None,
                     nick.GetHost(),
                     nick.GetNick(),
                     nick.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnPrivCTCP(self, nick, message):
        self.debug_hook()
        self.put_log("CTCP",
                     None,
                     nick.GetHost(),
                     nick.GetNick(),
                     nick.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnChanCTCP(self, user, channel, message):
        self.debug_hook()
        self.put_log("CTCP",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnUserCTCP(self, target, message):
        self.debug_hook()
        local_nick = self.GetNetwork().GetIRCNick()
        if local_nick is not None:
            self.put_log(
                "CTCP",
                None,
                local_nick.GetHost(),
                local_nick.GetNick(),
                local_nick.GetIdent(),
                None,
                datetime.now(),
                target.s,
                message.s,
                str(self.GetNetwork()),
            )
        else:
            self.put_log(
                "CTCP",
                None,
                None,
                None,
                None,
                None,
                datetime.now(),
                target.s,
                message.s,
                str(self.GetNetwork()),
            )
        return znc.CONTINUE

    def OnUserMsg(self, target, message):
        self.debug_hook()
        (channel, user) = self.resolveTarget(target)
        if channel is None:
            local_nick = self.GetNetwork().GetIRCNick()
            if local_nick is not None:
                self.put_log(
                    "MSG",
                    None,
                    local_nick.GetHost(),
                    local_nick.GetNick(),
                    local_nick.GetIdent(),
                    None,
                    datetime.now(),
                    target.s,
                    message.s,
                    str(self.GetNetwork()),
                )
            else:
                self.put_log(
                    "MSG",
                    None,
                    None,
                    None,
                    None,
                    None,
                    datetime.now(),
                    target.s,
                    message.s,
                    str(self.GetNetwork()),
                )
            return znc.CONTINUE
        return self.OnChanMsg(user, channel, message)

    def OnChanAction(self, user, channel, message):
        self.debug_hook()
        self.put_log("ME",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return True

    def OnChanMsg(self, user, channel, message):
        self.debug_hook()
        self.put_log("SAY",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     self.findMode(channel,
                                   user),
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def OnTopic(self, user, channel, message):
        self.debug_hook()
        self.put_log("TOPIC",
                     channel.GetName(),
                     user.GetHost(),
                     user.GetNick(),
                     user.GetIdent(),
                     None,
                     datetime.now(),
                     None,
                     message.s,
                     str(self.GetNetwork()))
        return znc.CONTINUE

    def findMode(self, channel, user):
        realUser = channel.FindNick(user.GetNick())
        if not realUser:
            return None
        return str(realUser.GetPermChar())

    def resolveTarget(self, target):
        target = target.s
        channel = self.GetNetwork().FindChan(target)
        if channel is None:
            return (None, None)
        user = self.GetNetwork().GetIRCNick()
        return (channel, user)

    def colourstrip(self, data):
        find = data.find('\x03')
        while find > -1:
            done = False
            data = data[0:find] + data[find + 1:]
            if len(data) <= find + 1:
                done = True
            try:
                assert not done
                assert int(data[find])
                while True:
                    assert int(data[find])
                    data = data[0:find] + data[find + 1:]
            except BaseException:
                if not done:
                    if data[find] != ',':
                        done = True

            if (not done) and (len(data) > find + 1) and (data[find] == ','):
                try:
                    assert not done
                    assert int(data[find + 1])
                    data = data[0:find] + data[find + 1:]
                    data = data[0:find] + data[find + 1:]
                except BaseException:
                    done = True
                try:
                    assert not done
                    while True:
                        assert int(data[find])
                        data = data[0:find] + data[find + 1:]
                except BaseException:
                    pass

            find = data.find('\x03')

        data = data.replace('\x02', '')
        data = data.replace('\x1d', '')
        data = data.replace('\x1f', '')
        data = data.replace('\x16', '')
        data = data.replace('\x0f', '')
        return data


class DatabaseThread:
    @staticmethod
    def worker_safe(
            db,
            log_queue,
            internal_log,
            debug_log_saves=False,
            batch_size=50) -> None:
        while True:
            try:
                DatabaseThread.worker(
                    db, log_queue, internal_log, debug_log_saves, batch_size)
                break
            except Exception as e:
                with internal_log.error() as target:
                    target.write(
                        'Worker thread crashed: {0} {1}\n'.format(
                            type(e), str(e)))
                    target.write('Stack trace: ' + traceback.format_exc())
                    target.write(
                        'Restarting worker thread in 5 seconds...\n\n')
                sleep(5)
                continue

    @staticmethod
    def worker(
            db,
            log_queue,
            internal_log,
            debug_log_saves=False,
            batch_size=50) -> None:
        stop = False
        while not stop:
            first = log_queue.get()
            if first is None:
                log_queue.task_done()   # count the sentinel too, or join() hangs
                return

            # Collect more rows without blocking, up to batch_size
            taken = 1
            items = [first]
            while len(items) < batch_size:
                try:
                    nxt = log_queue.get_nowait()
                except queue.Empty:
                    break
                taken += 1
                if nxt is None:
                    stop = True
                    break
                items.append(nxt)

            try:
                DatabaseThread.save_batch(
                    db, items, log_queue, internal_log, debug_log_saves)
            finally:
                # exactly one task_done() per get(), after the batch is
                # finished
                for _ in range(taken):
                    log_queue.task_done()

    @staticmethod
    def save_batch(
            db,
            items,
            log_queue,
            internal_log,
            debug_log_saves) -> None:
        while True:
            try:
                db.ensure_connected(internal_log, str(items[0]['_id']))
                db.insert_batch(items)
                if debug_log_saves:
                    with internal_log.debug() as target:
                        target.write(
                            'Saved {} log items queue_size={}\n'.format(
                                len(items), log_queue.qsize()))
                return
            except Exception as e:
                sleep_for = 10

                with internal_log.error() as target:
                    target.write(
                        'Could not save to database caused by: {0} {1}\n'.format(
                            type(e), str(e)))
                    target.write('Stack trace: ' + traceback.format_exc())
                    target.write(
                        'Batch of {} rows, first log: '.format(
                            len(items)))
                    item_copy = items[0].copy()
                    item_copy['_id'] = str(item_copy['_id'])
                    json.dump(item_copy, target)
                    target.write('\n\n')
                    target.write('Retry in {} s\n'.format(sleep_for))

                sleep(sleep_for)


class InternalLog:
    def __init__(self, save_path: str):
        self.save_path = save_path
        self.debug_enabled = False

    def debug(self):
        if not self.debug_enabled:
            from io import StringIO
            return StringIO()
        return self.open('debug')

    def error(self):
        return self.open('error')

    def set_debug_enabled(self, enabled: bool):
        self.debug_enabled = enabled

    def open(self, level: str):
        target = open(os.path.join(self.save_path, level + '.log'), 'a')
        return target


class Database:
    def __init__(self, dsn: dict):
        self.dsn = dsn
        self.conn = None


class MSSQLDatabase(Database):
    def connect(self) -> None:
        import pymssql
        self.conn = pymssql.connect(**self.dsn)

    def ensure_connected(self, internal_log, log_id):
        # No per-row SELECT 1. A dead connection is detected when an
        # insert fails; insert_batch() then drops it and the retry reconnects.
        if self.conn is None:
            self.connect()

    def insert_batch(self, rows):
        """Insert all rows, then commit once. All or nothing: on any error
        the transaction is rolled back and the connection is dropped."""
        try:
            for row in rows:
                params = (
                    row['znc_user'],
                    row['date'],
                    row['code'],
                    row['network'],
                    row['channel'],
                    row['host'],
                    row['user'],
                    row['ident'],
                    row['user_mode'],
                    row['target_user'],
                    row['message'])
                cur = self.conn.cursor()
                try:
                    cur.callproc("add_log", params)
                finally:
                    try:
                        cur.close()
                    except Exception:
                        pass
            self.conn.commit()
        except Exception:
            try:
                self.conn.rollback()
            except Exception:
                pass
            try:
                self.conn.close()
            except Exception:
                pass
            self.conn = None
            raise

    def insert_into(self, table, row):
        self.insert_batch([row])