#!/usr/bin/env python

import netrc
from datetime import datetime

import click
import requests
from tabulate import tabulate
from prompt_toolkit import print_formatted_text as print
from prompt_toolkit import HTML
from prompt_toolkit import prompt
from prompt_toolkit.completion import FuzzyWordCompleter
from prompt_toolkit.lexers import PygmentsLexer
from prompt_toolkit.styles import style_from_pygments_cls
from pygments.lexers.promql import PromQLLexer
from pygments.styles.solarized import SolarizedLightStyle


class promcli(object):
    def __init__(self, user, password, promhost):
        self.promhost = promhost
        self.auth = (user, password)
        self.df = "%Y-%m-%d %H:%M:%S"

    def get_series(self):
        resp = requests.get(
            f"https://{self.promhost}/api/v1/label/__name__/values",
            auth=self.auth
        )
        return resp.json().get("data", [])

    def get_labels(self, series):
        resp = requests.post(
            f"https://{self.promhost}/api/v1/labels",
            auth=self.auth,
            data={"match[]": f"{series}"},
        )
        return resp.json().get("data", [])

    def get_query(self, query):
        resp = requests.post(
            f"https://{self.promhost}/api/v1/query",
            data={"query": query},
            auth=self.auth
        )
        return resp.json().get("data", {}).get("result", [])

    def format_query(self, results):
        data = []
        for result in results:
            metric = result.get("metric")
            timestamp, value = result.get("value", [0, None])
            data.append(
                [
                    datetime.fromtimestamp(timestamp).strftime(self.df),
                    metric.get("__name__"),
                    metric.get("host", ""),
                    metric.get("job", ""),
                    value,
                ]
            )
        return tabulate(data, headers=["", "NAME", "HOST", "JOB", "VALUE"])


@click.group()
@click.option("-u", "--user")
@click.option("-p", "--password")
@click.option("-h", "--promhost")
@click.option("-n", "--netrc", "use_netrc", is_flag=True)
@click.pass_context
def cli(ctx, user, password, promhost, use_netrc):
    ctx.ensure_object(dict)
    if use_netrc:
        netrc_instance = netrc.netrc()
        user, _, password = netrc_instance.authenticators(promhost)
    ctx.obj['promcli'] = promcli(user, password, promhost)


@cli.command()
@click.option("-q", "--query")
@click.pass_context
def query(ctx, query):
    promclio = ctx.obj['promcli']
    if not query:
        query = prompt(
            "query: ",
            completer=FuzzyWordCompleter(promclio.get_series()),
            lexer=PygmentsLexer(PromQLLexer),
            style=style_from_pygments_cls(SolarizedLightStyle),
        )
        if '{' not in query:
            params = prompt(
                "{ ",
                completer=FuzzyWordCompleter(promclio.get_labels(query)),
                complete_in_thread=True,
                lexer=PygmentsLexer(PromQLLexer),
                style=style_from_pygments_cls(SolarizedLightStyle),
            )
            query = f"{query}{{{params}}}"
    print(HTML(f"Querying <lightblue>{query}</lightblue>"))
    print(promclio.format_query(promclio.get_query(f"{query}")))


@cli.command()
@click.pass_context
def alerts(ctx):
    promclio = ctx.obj['promcli']
    promalerts = promclio.get_query('ALERTS{alertstate="firing"}')
    data = []
    for alert in promalerts:
        metric = alert.get("metric", {})
        timestamp, value = alert.get("value", [0, None])
        data.append(
            [
                datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S"),
                metric.get("alertname", ""),
                metric.get("severity", ""),
                metric.get("host", ""),
                metric.get("job", ""),
                value
            ]
        )
        print(tabulate(data, headers=["", "ALERT", "SEVERITY", "HOST", "JOB", "VALUE"]))


if __name__ == "__main__":
    cli(obj={})