# importation de ray
import ray
from email import policy
from email.parser import BytesParser
from ray.data import read_binary_files, Dataset
#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 un bloc d'e-mails
def parse_email_batch(batch):
    
    #instanciation locale de l'outil de calcul de la polarité
    sentiment = SentimentIntensityAnalyzer()
    
    #structure de résultats, dicttionnaire de listes
    results = {
        "message_id": [],
        "date": [],
        "from": [],
        "body": [],
        "score": []
    }
    
    #pour chaque message
    for content in batch["bytes"]:
    
        # 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"]
        
        results["message_id"].append(msg.get("Message-ID", ""))
        results["date"].append(msg.get("Date", ""))
        results["from"].append(msg.get("From", ""))
        results["body"].append(body)
        results["score"].append(score)
        
    return(results)


            
# programme principal
if __name__ == "__main__":
    
    # paramètre en ligne de commande
    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int,default=4)
    parser.add_argument("--size", type=int,default=1000)
    
    # 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}")  
    
    #taille du batch
    taille_batch = args.size
    print(f"Batch size = {taille_batch}")
        
    # Initialiser Ray
    ray.init(ignore_reinit_error=True,
             object_store_memory=6 * 1024**3, #6 GB
             num_cpus=num_cores)
    
    print("Debut du traitement...")
    
    #debut chrono
    debut = time.time()  
    
    # Lire tous les fichiers en parallèle avec Ray
    ds: Dataset = read_binary_files(EMAIL_FOLDER, include_paths=True)

    #Appliquer la fonction de parsing
    parsed_ds = ds.map_batches(
        parse_email_batch,
        batch_size=taille_batch
        )
   
    # 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()
