#!/usr/bin/env python3
"""RKI reader

Tool to download and print Covid-19 data from https://npgeo-corona-npgeo-de.hub.arcgis.com/

Requirements are:
 * requests
 * click
 * tabulate
 * colorama
"""
import json
import requests
import click
import tabulate
from colorama import Fore

BASE_URL = "https://services7.arcgis.com/mOBPykOjAyBO2ZKk/arcgis/rest/services/RKI_Landkreisdaten/FeatureServer/0/query"
BASE_PARAM = {
    "f": "json",
    "returnGeometry": False,
    "where": "1=1",
}


@click.group()
def cli():
    pass


@click.command()
@click.option(
    "--format", "_format", default="json", type=click.Choice(["text", "json"])
)
def list_countries(_format):
    params = BASE_PARAM.copy()
    params.update({"outFields": "BL,county"})
    r = requests.get(BASE_URL, params)
    result = list({x["attributes"]["BL"] for x in r.json()["features"]})
    result = sorted(result)
    if _format == "json":
        click.echo(json.dumps(result))
    elif _format == "text":
        for c in result:
            click.echo(c)


cli.add_command(list_countries)


@click.command()
@click.option("--country")
@click.option(
    "--format", "_format", default="json", type=click.Choice(["text", "json"])
)
def list_counties(country, _format):
    params = BASE_PARAM.copy()
    params.update({"outFields": "BL,county", "where": "BL='{}'".format(country)})
    r = requests.get(BASE_URL, params)
    result = list({x["attributes"]["county"] for x in r.json()["features"]})
    result = sorted(result)
    if _format == "json":
        click.echo(json.dumps(result))
    elif _format == "text":
        for c in result:
            click.echo(c)


cli.add_command(list_counties)


@click.command()
@click.option("--selector", "selectors", multiple=True, help="select a county in a state (e.g. 'Bayern/LK Bamberg'). Can be repeated to print data for multiple areas")
@click.option(
    "--format", "_format", default="json", type=click.Choice(["text", "json"]), help="select output format, text is a table in text mode, json is a json dump",
)
@click.option(
    "--sort", "sorts", multiple=True, type=click.Choice(["cases", "state", "county", "name"]), help=""
)
@click.option(
    "--color-limits", "color_limits", help="2 comma separated ints (e.g. '30,60') to define limits to colorize in green, yellow, red."
)
def get_selected(selectors, _format, sorts, color_limits):
    selectors = [s.split("/") for s in selectors]
    params = BASE_PARAM.copy()
    wheres = " or ".join(
        ["(BL='{}' and county='{}')".format(s[0], s[1]) for s in selectors]
    )
    params.update(
        {"outFields": "BL,county,cases7_per_100k,last_update", "where": wheres}
    )
    r = requests.get(BASE_URL, params)
    result = r.json()["features"]
    result = list([x["attributes"] for x in r.json()["features"]])
    if _format == "json":
        click.echo(json.dumps(result))
    elif _format == "text":
        if sorts:
            if sorts[0] == "state":
                result = sorted(result, key=lambda x: x["BL"])
            if sorts[0] == "cases":
                result = sorted(result, key=lambda x: x["cases7_per_100k"], reverse=True)
            if sorts[0] == "name":
                result = sorted(result, key=lambda x: x["county"][3:])
        for line in result:
            cases = line["cases7_per_100k"]
            color = ""
            if color_limits:
                l_1, l_2 = color_limits.split(",", 1)
                if cases < int(l_1):
                    color = Fore.GREEN
                elif cases < int(l_2):
                    color = Fore.YELLOW
                else:
                    color = Fore.RED
            line["cases7_per_100k"] = "{}{:.2f}{}".format(color, cases, Fore.RESET)
        table = [
            [c["BL"], c["county"], c["cases7_per_100k"], c["last_update"]]
            for c in result
        ]
        click.echo(
            tabulate.tabulate(
                table,
                #                headers=["Bundesland", "Kreis", "7 Tagesinzidenz", "aktualisiert"],
                headers=[
                    "State",
                    "County",
                    "new cases (last 7 days, per 100000 inhabitants) ",
                    "updated at",
                ],
                tablefmt="pretty",
            )
        )


cli.add_command(get_selected)


if __name__ == "__main__":
    cli()