#!/usr/bin/python3

import subprocess
from enum import Enum
from typing import List, Dict


class ShellCommandRunner:
    class ErrorData:
        def __init__(self, name, message):
            self._name = name
            self._message = message

        @property
        def name(self):
            return self._name

        @property
        def message(self):
            return self._message

    @classmethod
    def _build_cmd(cls):
        pass

    @classmethod
    def exec(cls):
        output = None
        error = None
        cmd = cls._build_cmd() or []
        try:
            output = subprocess.run(
                cmd, check=True, capture_output=True, text=True
            )
        except subprocess.CalledProcessError as err:
            error_name = err.__class__.__name__
            message = str(err.stderr)
            error = cls.ErrorData(error_name, message)
        except Exception as err:
            error_name = err.__class__.__name__
            message = str(err)
            error = cls.ErrorData(error_name, message)
        return output, error

class SystemCtlCommandRunner(ShellCommandRunner):
    _systemctl_bin = ['/usr/bin/systemctl']
    _systemctl_cmd = ['show']
    _systemctl_service = []
    _systemctl_options = []

    @classmethod
    def _build_cmd(cls):
        cmd = []
        cmd.extend(cls._systemctl_bin)
        cmd.extend(cls._systemctl_cmd)
        cmd.extend(cls._systemctl_service)
        cmd.extend(cls._systemctl_options)
        return cmd


class SystemCtlShowState(SystemCtlCommandRunner):
    _systemctl_cmd = ['show']
    _systemctl_options = ['-p', 'ActiveState']

    class SystemCtlState(Enum):
        active = 'active'
        inactive = 'inactive'
        failed = 'failed'
        unknown = 'unknown'

    @classmethod
    def _parse(cls):
        output, error = cls.exec()
        if error:
            return 'ActiveState', cls.SystemCtlState.unknown.value
        return output.stdout.split('=')

    @classmethod
    def state(cls):
        key, value = cls._parse()
        return value.strip()

    @classmethod
    def is_active(cls):
        current_state = cls.state()
        return current_state == cls.SystemCtlState.active.value

    @classmethod
    def is_inactive(cls):
        current_state = cls.state()
        return current_state == cls.SystemCtlState.inactive.value

    @classmethod
    def is_failed(cls):
        current_state = cls.state()
        return current_state == cls.SystemCtlState.failed.value


class PSQLCommandRunner(ShellCommandRunner):
    _psql_bin = ['/usr/bin/psql']
    _psql_host = ['-h', '/run/postgresql/']
    _psql_port = ['-p', '5432']
    _psql_db = ['']
    _psql_cmd = ['-c', '\?']
    _psql_options = []

    @classmethod
    def _build_cmd(cls):
        cmd = []
        cmd.extend(cls._psql_bin)
        cmd.extend(cls._psql_host)
        cmd.extend(cls._psql_port)
        cmd.extend(cls._psql_db)
        cmd.extend(cls._psql_cmd)
        cmd.extend(cls._psql_options)
        return cmd


class PSQLFormattedOutput(PSQLCommandRunner):
    _prefix = 'postgres'
    _psql_options = ['-xA']

    @staticmethod
    def _is_number(number: str):
        try:
            float(number)
            return True
        except ValueError:
            return False

    @classmethod
    def _parse(cls):
        """
        Expects output like (thanks to -xA options in psql command):
        name|coredns
        host|
        port|5432
        database|coredns
        pool_size|10
        ...

        name|glance
        host|
        port|5432
        database|glance
        pool_size|10
        ...
        :return: List[Dict], where keys are table column names, values are table records
        """
        output, error = cls.exec()

        if error:
            return [{
                'name': error.name,
                'message': error.message,
                'metric_collector_error': 1
            }]

        d = dict()
        records = []
        for line in output.stdout.split('\n'):
            if line == '':
                records.append(d)
                d = dict()
                continue
            key, value = line.split('|')
            d[key] = value
        return records

    @classmethod
    def _label_set(cls, records: List[Dict]):
        """
        This method tries to guess what keys are labels or metrics guided by next logic:
        any metric value can be converted to float type, if not - it's a label value.
        You can override this method and specify your own list of labels.
        :param records: List of records
        :return: Set of label names
        """
        return {
            k for k, v in records[0].items() if not cls._is_number(v)
        }

    @classmethod
    def _metric_name(cls, metric: str):
        if cls._prefix:
            metric_name = '{0}_{1}'.format(cls._prefix, metric)
        else:
            metric_name = metric
        return metric_name

    @staticmethod
    def _metric_help_description(metric_name: str, metric_description: str = None):
        if not metric_description:
            metric_description = metric_name.lower().capitalize().replace('_', ' ')
        return '# HELP {0} {1}'.format(metric_name, metric_description)

    @staticmethod
    def _metric_help_type(metric_name: str, metric_type: str = None):
        if metric_type is None:
            metric_type = 'gauge'
        return '# TYPE {0} {1}'.format(metric_name, metric_type)

    @staticmethod
    def _metric_record(metric_name: str, metric: str, labels: List[str], record: List[Dict]):
        labels = ','.join(
            ['{0}="{1}"'.format(k, record[k]) for k in labels]
        )
        labels = '{{{0}}}'.format(labels)
        return '{0}{1} {2}'.format(metric_name, labels, record[metric])

    @classmethod
    def print_metrics(cls):
        """
        Print Prometheus-like metrics
        :return: None
        """
        records = cls._parse()
        labels = cls._label_set(records)
        metrics = records[0].keys() - labels
        for metric in metrics:
            metric_name = cls._metric_name(metric)
            print(cls._metric_help_description(metric_name))
            print(cls._metric_help_type(metric_name))
            for record in records:
                print(cls._metric_record(metric_name, metric, labels, record))


class PGBouncerFormattedOutput(PSQLFormattedOutput):
    _psql_host = ['-h', '/run/pgbouncer/']
    _psql_port = ['-p', '6432']
    _psql_db = ['pgbouncer']
    _psql_cmd = ['show help']
    _prefix = 'pgbouncer_help'


class VStorageUIBackendShowState(SystemCtlShowState):
    _systemctl_service = ['vstorage-ui-backend']


class PGBouncerPoolList(PGBouncerFormattedOutput):
    _psql_cmd = ['-c', 'show pools']
    _prefix = 'pgbouncer_pool'


class PGBouncerStatList(PGBouncerFormattedOutput):
    _psql_cmd = ['-c', 'show stats']
    _prefix = 'pgbouncer_stat'


class PGBouncerDBList(PGBouncerFormattedOutput):
    _psql_cmd = ['-c', 'show databases']
    _prefix = 'pgbouncer_database'


if VStorageUIBackendShowState.is_active():
    PGBouncerPoolList.print_metrics()
    PGBouncerStatList.print_metrics()
    PGBouncerDBList.print_metrics()
