import toml
import os
import db_model as db
from peewee import SqliteDatabase, IntegrityError, PostgresqlDatabase
from sshtunnel import SSHTunnelForwarder


def config_laden():
    configfile = os.path.join(SKRIPTPFAD, "config.toml")
    with open(configfile) as file:
        return toml.loads(file.read())


SKRIPTPFAD = os.path.abspath(os.path.dirname(__file__))
CONFIG = config_laden()
QUELLEN = ["jhu", "who", "rki"]


def alle_country_vereinheitlichen():
    print("Betrete country vereinheitlichen")
    country_dict = {}
    for datensatz in db.CoronaCountry.select():
        neuer_name = datensatz.country_region.lower().capitalize()
        print(neuer_name)
        try:
            if neuer_name not in country_dict[datensatz.quelle]:
                country_dict[datensatz.quelle].append(neuer_name)
        except KeyError:
            country_dict[datensatz.quelle] = [neuer_name]
        with db.database.atomic() as transaction:
            try:
                db.CoronaDaten.update(country_region=neuer_name).where(
                    db.CoronaDaten.country_region == datensatz.country_region).execute()
                print(neuer_name, datensatz.country_region)
            except IntegrityError:
                print("Fehler")
                transaction.rollback()
    print("Schleife Beendet")
    db.CoronaCountry.delete().execute()
    with db.database.atomic():
        print("CountryCreate")
        for quelle, countrys in country_dict.items():
            print(quelle, countrys)
            for country in countrys:
                db.CoronaCountry.create(country_region=country, quelle=quelle)


def country_umbennnen():
    quelle = input("Quelle eingeben: ")
    name_altes_land = input("Name der alten Bezeichnung des Landes eingeben: ")
    name_neues_land = input("Name der neuen Bezeichnung des Landes eingeben: ")

    anzahl_daten = db.CoronaCountry.select().where(
        (db.CoronaCountry.country_region == name_neues_land) & (db.CoronaCountry.quelle == quelle)
    ).count()
    if anzahl_daten > 0:
        db.CoronaCountry.delete().where(
            db.CoronaCountry.country_region == name_altes_land
        ).execute()
    else:
        db.CoronaCountry.update(country_region=name_neues_land).where(
            db.CoronaCountry.country_region == name_altes_land
        ).execute()
    try:
        db.CoronaDaten.update(country_region=name_neues_land).where(
            db.CoronaDaten.country_region == name_altes_land
        ).execute()
    except IntegrityError:
        db.CoronaDaten.delete().where(
            db.CoronaDaten.country_region == name_altes_land
        ).execute()


def main():
    with SSHTunnelForwarder(
                (CONFIG["ssh"]["ip_server"], CONFIG["ssh"]["ssh_port"]), ssh_username=CONFIG["ssh"]["user"],
                ssh_password=CONFIG["ssh"]["pw"], remote_bind_address=('127.0.0.1', CONFIG["pg"]["pgport"])) as server:

        db.database.initialize(PostgresqlDatabase(CONFIG["pg"]["pgdb"],
                                                  user=CONFIG["pg"]["pguser"], password=CONFIG["pg"]["pgpw"],
                                                  host="127.0.0.1",
                                                  port=server.local_bind_port))
        country_umbennnen()


main()