# importation de ray
import ray
from email import policy
from email.parser import BytesParser
# importation du module RAY.DATA /!\
import ray.data
#import nltk
from nltk.sentiment.vader import SentimentIntensityAnalyzer
import os
import time
import argparse


# nécessaire une seule fois dans l'environnement
#nltk.download('vader_lexicon')

# Répertoire contenant les fichiers e-mails
EMAIL_FOLDER = "emails"

# pour éviter le warning
os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0"

# Fonction pour parser le contenu d'un email
def parse_email_content(content, sentiment):
    try:
        # parsing du message (date, expéditeur)
        msg = BytesParser(policy=policy.default).parsebytes(content)
        
        # récuperation du corps du message
        body = msg.get_payload(decode=True).decode(errors="ignore")
        
        # calcul du score de positivité
        score = sentiment.polarity_scores(body)["pos"]
        
        return {
            "message_id": msg.get("Message-ID", ""),
            "date": msg.get("Date", ""),
            "from": msg.get("From", ""),
            "body": body,
            "score": score
        }
    except Exception as e:
        return {
            "message_id": None,
            "date": None,
            "from": None,
            "body": None,
            "error": str(e)
        }


# programme principal
if __name__ == "__main__":
    
    # paramètre en ligne de commande
    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int,default=4)
    
    # récupération des paramètres
    args = parser.parse_args()

    #nombre de cores
    num_cores = args.n
    print(f"Nombre de coeurs autorisé à utiliser = {num_cores}")    
    
    # Initialiser Ray
    ray.init(ignore_reinit_error=True,
             object_store_memory=6 * 1024**3, #6 GB
             num_cpus=num_cores)
    
    #instanciation de l'outil de calcul de la polarité
    sentiment = SentimentIntensityAnalyzer()    
    
    print("Debut du traitement...")
    
    #debut chrono
    debut = time.time()  

    # Lire tous les fichiers en parallèle avec Ray
    ds: ray.data.Dataset = ray.data.read_binary_files(EMAIL_FOLDER, include_paths=True)

    #Appliquer la fonction de parsing
    parsed_ds = ds.map(fn = lambda row: parse_email_content(row["bytes"], sentiment))
   
    # transformer en pandas les résultats
    df = parsed_ds.to_pandas()
    
    # fin chrono
    duree = time.time() - debut
    print(f"\nFin du traitement, duree = {duree} sec.\n")    
    
    # afficher les premières lignes
    print(df.sort_values(by="message_id")[['date','score']].head(10))

    #fin
    ray.shutdown()
