#!/usr/bin/env python3

import argparse
import pathlib
import subprocess
import os
import psycopg2
from pathlib import Path

from sqlalchemy.engine import URL
from sqlalchemy import create_engine
import paramiko
import datetime


def etablit_connexion_SFTP(p_hote, p_port, p_utilisateur, p_motdepasse):
    connexion_ssh = paramiko.SSHClient()
    connexion_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    connexion_ssh.connect(p_hote, p_port, p_utilisateur, p_motdepasse)
    return connexion_ssh.open_sftp()


def creer_chaine_connexion(p_utilisateur, p_pass, p_serveur, p_port,p_base):
    # return f"postgresql://{p_utilisateur}@{p_serveur}:{p_port}"
    url_geo = URL.create('postgresql+psycopg2', 
        username = p_utilisateur,
        password = p_pass,
        host = p_serveur,
        port = p_port,
        database = p_base,
    )
    return create_engine(url_geo)


def creer_base_de_donnee(p_base_a_creer, p_base_connexion, p_chaine_connexion):
    requete_SQL = f'CREATE DATABASE "{p_base_a_creer}" WITH OWNER logeproj'
    commande = [
        "psql",
        f"{p_chaine_connexion}/{p_base_connexion}",
        "-c",
        f"{requete_SQL}",
    ]
    commande_run = subprocess.run(commande)
    code_retour = commande_run.returncode
    if code_retour != 0:
        return False
    else:
        return True


def supprimer_base_de_donnee(p_base_a_supprimer, p_base_connexion, p_chaine_connexion):
    requete_SQL = f'DROP DATABASE "{p_base_a_supprimer}"'
    commande = [
        "psql",
        f"{p_chaine_connexion}/{p_base_connexion}",
        "-c",
        f"{requete_SQL}",
    ]
    commande_run = subprocess.run(commande)
    code_retour = commande_run.returncode
    if code_retour != 0:
        return False
    else:
        return True


def lister_bases_de_donnees(p_connexion):
    tableau_bases = []
    requete_SQL = "SELECT datname FROM pg_database WHERE datistemplate = false;"
    curseur = p_connexion
    curseur.execute(requete_SQL)
    for enregistrement in curseur:
        tableau_bases.append(enregistrement[0])
    return tableau_bases


def restaurer_sauvegarde(
    p_chemin_sauvegarde, p_base, p_base_connexion, p_chaine_de_connexion
):
    # connexion = psycopg2.connect(p_chaine_de_connexion)
    connexion = p_chaine_de_connexion
    if p_base in lister_bases_de_donnees(connexion):
        supprimer_base_de_donnee(p_base, p_base_connexion, p_chaine_de_connexion)
    creer_base_de_donnee(p_base, p_base_connexion, p_chaine_de_connexion)
    connexion.close()
    commande = [
        "pg_restore",
        "-d",
        f"{p_chaine_de_connexion}/{p_base}",
        p_chemin_sauvegarde,
    ]
    commande_run = subprocess.run(commande)
    code_retour = commande_run.returncode
    if code_retour != 0:
        return False
    else:
        return True


def restaurer_toutes_sauvegardes(
    p_chemin_sauvegarde, p_base_de_connexion, p_chaine_connexion
):
    repertoire_des_sauvegardes = pathlib.Path(p_chemin_sauvegarde)
    for sauvegarde in repertoire_des_sauvegardes.glob("*-logeproj-centrale.custom"):
        sauvegarde.resolve()
        chemin = str(sauvegarde)
        fichier = sauvegarde.name
        fichier_split = fichier.split("-")
        base = f"{fichier_split[0]}-logeproj-test"
        restaurer_sauvegarde(chemin, base, p_base_de_connexion, p_chaine_connexion)


def cherche_chemin_rep_sauvegardes_du_jour(p_connexion_sftp, p_rep_racine, p_date, p_format=None):
    liste_rep_backups_du_jour = []
    date_backup = "1900-01-01"
    rep_split = ["1900", "01", "01"]

    p_connexion_sftp.chdir("/"+p_rep_racine)

    for rep in p_connexion_sftp.listdir():
        rep_split = rep.split("-")
        
        if len(rep_split) < 3 : continue
        if p_format is not None and rep.find(p_format) == -1: continue

        date_backup = f"{rep_split[0]}-{rep_split[1]}-{rep_split[2]}"

        if date_backup == p_date.strftime("%Y-%m-%d"):
            liste_rep_backups_du_jour.append(rep)

    nombre_de_rep = len(liste_rep_backups_du_jour)


    if nombre_de_rep > 1:
        print(
            f"Le serveur de sauvegarde dispose de plusieurs répertoires de sauvegardes datés du {p_date.strftime('%d/%m/%Y')} :"
        )
        for indice in range(0, nombre_de_rep):
            numero_rep = indice + 1
            print(f"- {numero_rep} : {liste_rep_backups_du_jour[indice]}")
        choix = int(
            input(
                "Entrer le numéro du répertoire dont vous souhaitez télécharger les sauvegardes : "
            )
        )
        rep_sauvegarde_du_jour = liste_rep_backups_du_jour[choix - 1]
    else:
        rep_sauvegarde_du_jour = liste_rep_backups_du_jour[0]

    chemin_sauvegardes_du_jour = f"/{p_rep_racine}/{rep_sauvegarde_du_jour}"
    return chemin_sauvegardes_du_jour


def telecharge_un_fichier(
    p_chemin_fichier_distant, p_chemin_fichier_local, p_connexionSFTP
):
    print(
        f"- Téléchargement de {p_chemin_fichier_distant} dans {p_chemin_fichier_local}"
    )
    p_connexionSFTP.get(p_chemin_fichier_distant, p_chemin_fichier_local)


def clos_connexion_SFTP(p_connexion_SFTP):
    p_connexion_SFTP.close()


def main():
    aujourdhui = datetime.date.today()
    chemin_rep_local_sauvegarde = "./"
    sftp_hote = "ftpmdv.myriadev.fr"
    sftp_port = "22060"
    sftp_utilisateur = "cen-isere"
    sftp_mot_de_passe = "gaic4Ao'Do1O"
    sftp_rep_racine_sauvegardes = 'cen-isere'

    repertoire_sauvegardes = "/cen-isere"
    serveur_postgresql = "91.134.194.221"
    port_postgresql = "5432"
    utilisateur_postgresql = "cen_admin"
    mot_de_passe_postgresql = "#CEN38\@venir"
    base_de_connexion = "logeproj"

    # Si le mot de passe de postgre est stocké dans une variable d'environnement, plus
    # besoin de le fournir par la suite
    os.environ["PGPASSWORD"] = mot_de_passe_postgresql

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-b", "--base",
        required=False,
        help="nom de la base à restaurer (si non fourni, le script restaure toutes les bases",
    )

    
    arguments = parser.parse_args()
    arguments.base = 'logeproj'

    connexion_sftp = etablit_connexion_SFTP(
        sftp_hote, sftp_port, sftp_utilisateur, sftp_mot_de_passe
    )
    chemin_rep_sauvegardes_du_jour = cherche_chemin_rep_sauvegardes_du_jour(
        connexion_sftp, sftp_rep_racine_sauvegardes, aujourdhui,'gz'
    )


    # Télécharger la base fournie en argument au format custom
    chemin_sauvegarde_distante = chemin_rep_sauvegardes_du_jour
    chemin_sauvegarde_locale = (
        f"{chemin_rep_local_sauvegarde}"+chemin_rep_sauvegardes_du_jour.rsplit('/',1)[-1]
    )
    telecharge_un_fichier(
        chemin_sauvegarde_distante, chemin_sauvegarde_locale, connexion_sftp
    )

    clos_connexion_SFTP(connexion_sftp)

    chaine_de_connexion = creer_chaine_connexion(
        utilisateur_postgresql,mot_de_passe_postgresql, serveur_postgresql, port_postgresql,base_de_connexion
    )

    if arguments.base is None:
        # Restaurer toutes les sauvegardes
        print(f'Restauration de toutes les sauvegardes')
        restaurer_toutes_sauvegardes(
            repertoire_sauvegardes, base_de_connexion, chaine_de_connexion
        )
    else:
        # Restaurer la base fournie en argument
        base=f'{arguments.base}'
        print(f'Restauration de {base}')
        fichier_sauvegarde = chemin_sauvegarde_locale
        path_sauvegarde = pathlib.Path(fichier_sauvegarde)
        restaurer_sauvegarde(
            p_chemin_sauvegarde = path_sauvegarde,
            p_base = base,
            p_base_connexion = base_de_connexion,
            p_chaine_de_connexion = chaine_de_connexion,
        )


if __name__ == "__main__":
    main()