#!/opt/maint/venv/bin/python3
"""Post CloudLinux user stats to grafana"""
from argparse import ArgumentParser
import shlex
import subprocess
import json
import sys
import traceback
from socket import gethostname
from datetime import datetime
from pathlib import Path
from time import time
import requests

maint_dir = Path("/opt/maint/etc")
GLOBAL_IDS = {}

CRITICAL = 2
WARNING = 1


def run_script(cmd, check_exit=True):
    return subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        check=check_exit,
        encoding='utf-8',
    ).stdout


class StatsAPI:
    def __init__(self, conf):
        self.url = conf["host"]
        self.key = conf["key"]

    def post_stats(self, stats, period):
        server = gethostname()
        data = {"server": gethostname(), "period": period, "stats": stats}

        response = requests.post(
            f"{self.url}/v1/userstats/?server={server.split('.')[0]}",
            json=data,
            headers={"Authorization": f"{self.key}"},
            timeout=10,
        )
        return response


def fatal_error(msg):
    print(f"Error: {msg}")
    msg = f"{datetime.now()}: {msg}"

    err_file = maint_dir / "userstats.err"
    footer = (
        "Remove this file if you've reviewed the error message"
        "or fixed the issue."
    )
    err_file.write_text(f"{msg}\n\n{footer}")
    sys.exit(1)


def load_conf():
    conf = maint_dir / "userstats.cfg"
    if not conf.exists():
        fatal_error(f"Configuration file {conf} not found.")
    return json.loads(conf.read_text())


def parse_args():
    parser = ArgumentParser(description="Export user usage statistics")

    parser.add_argument(
        "--check",
        action="store_true",
    )

    return parser.parse_args()


def check():
    err_file = maint_dir / "userstats.err"

    status = run_script(["systemctl", "status", "lvestats"], check_exit=False)
    if "Active: active" not in status:
        print("CRITICAL: lvestats is not running.")
        sys.exit(CRITICAL)

    if err_file.exists():
        now = datetime.now().astimezone().timestamp()
        delta = now - err_file.stat().st_mtime
        if delta < 60 * 15:  # 5 minutes
            print(
                f"CRITICAL: Userstats failfile tripped. Error in: {err_file} "
                "When the issue is fixed - the command runs normally - "
                "remove the file"
            )
            sys.exit(CRITICAL)
        else:
            print(
                f"WARNING: {err_file} exists and is "
                f"{int(delta/60)} minutes old."
            )
            sys.exit(WARNING)

    if "ERROR" in status:
        print("WARNING: lvestats service encountered an error.")
        sys.exit(WARNING)

    print("OK.")
    sys.exit(0)


def detect_period(average_period: int, touch_threshold=False):
    """
    Use touch file to detect period

    touch_threshold determines if we should touch the file regardless of
    average_period. True means we only touch it if delta is > average_period.
    """
    touch_file = Path(f"/run/clstats-touch-{average_period}")
    if not touch_file.exists():
        touch_file.touch()
        return average_period
    mtime = int(touch_file.stat().st_mtime)
    now = int(time())
    time_delta = now - mtime
    minutes = int(round(time_delta / 60))
    if not touch_threshold or minutes >= average_period:
        touch_file.touch()
        return max(average_period, minutes)
    return minutes


def get_stats(period):
    stats = [
        "/sbin/cloudlinux-statistics",
        "--period",
        period,
        "--json",
    ]
    try:
        usage_stats = run_script(stats)
    except subprocess.CalledProcessError as exc:
        sys.exit(
            f"Command returned exit code {exc.returncode}: "
            f"{shlex.join(exc.cmd)}"
        )
    usage_data = json.loads(usage_stats)
    rows = []

    userdata = usage_data["users"]

    for user_details in userdata:
        usage = user_details["usage"]
        limits = user_details["limits"]
        faults = user_details["faults"]
        username = user_details["username"]
        domain = user_details["domain"]
        reseller = user_details["reseller"]
        row = {}

        def convert_int(stat):
            return min(int(round(stat["lve"])), 2147483647)

        usage_convert_table = {
            "cpu": lambda x: min(int(round(x["lve"])), 65535),
            "ep": lambda x: min(int(x["lve"]), 65535),
            "io": lambda x: min(int(x["lve"] / 1024 / 1024), 65535),  # mb
            "iops": lambda x: min(int(x["lve"]), 65535),
            "pmem": lambda x: min(int(x["lve"] / 1024 / 1024), 65535),  # mb
            "vmem": lambda x: min(int(x["lve"] / 1024 / 1024), 65535),  # mb
            "nproc": lambda x: min(int(x["lve"]), 65535),
        }
        row["usage"] = {}
        for key, value in usage.items():
            converted = usage_convert_table[key](value)
            key_name = key
            if key in ["io", "pmem", "vmem"]:
                key_name = f"{key}_m"

            row["usage"][key_name] = converted

        row["limits"] = {}
        for key, value in limits.items():
            row["limits"][key] = convert_int(value)

        row["faults"] = {}
        for key, value in faults.items():
            row["faults"][key] = convert_int(value)

        row["domain"] = domain
        row["reseller"] = reseller
        row["username"] = username
        rows.append(row)

    return rows


def main():
    args = parse_args()
    if args.check:
        check()

    if not Path("/sbin/cloudlinux-statistics").exists():
        fatal_error(
            "cloudLinux-statistics not found. "
            "This tool only works on CloudLinux systems."
        )
    conf = load_conf()
    try:
        api = StatsAPI(conf["api"])

        period = detect_period(1)
        rows = get_stats(f"{period}m")
        if rows:
            post1 = api.post_stats(rows, 1)
            if post1.status_code != 200:
                fatal_error(
                    "Failed to post user stats to Grafana: "
                    f"{post1.status_code}\n{post1.text}"
                )

    except Exception as e:
        stack_trace = traceback.format_exc()
        fatal_error(f"An error occurred: {str(e)}\n{stack_trace}")


if __name__ == "__main__":
    main()
