import os
import socket
import time
import json
import IA_Generate

host = '192.168.10.56'
port = 3001
# Attendre 1 seconde
time.sleep(1)

# Créer un objet socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# Se débarrasser du message d'erreur "ConnectionRefusedError: [Errno 111] Connection refused"
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

# Se débarrasser du message d'erreur "Address already in use"
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

# Se connecter au serveur Node.js
s.connect((host, port))

socket.SO_VM_SOCKETS_BUFFER_MAX_SIZE = (1024 * 1024 * 1024)

s.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 2 ^ 20)  # Buffer size = 1MB

# Envoyer des données au serveur
message = 'Bonjour, serveur!'
print('Envoi de:', message)
s.send(message.encode())  # Envoyer le message encodé en bytes

pipe=IA_Generate.init_gpu()
print('init gpu !')

while True:
    # Recevoir les données du serveur
    web_data_receive = s.recv(1024)  # Recevoir jusqu'à 1024 bytes

    if web_data_receive:
        #print('Données reçues:', web_data_receive.decode())  # Décoder les bytes en string

        web_data = json.loads(web_data_receive.decode())

        web_data_info = web_data['ID']
        web_data_Prompt = web_data['Prompt']
        web_data_NB_img = web_data['NB_img']
        web_data_Width = web_data['Width']
        web_data_Height = web_data['Height']
        web_data_Seed = web_data['Seed']
        web_data_Neg_Prompt = web_data['Neg_Prompt']
        web_data_Guidance = web_data['Guidance']
        web_data_Step = web_data['Step']
        web_data_Utilisateur = web_data['Utilisateur']

        print('Données reçues:', web_data)  # Décoder les bytes en string

        # Envoyer des données au serveur
        message = 'Accusé de réception'
        # s.send(message.encode())  # Envoyer le message encodé en bytes

        # Envoyer les images comtenues dans le dossier "images-IA-py" au serveur node.js
        if web_data_info == 'images':
            seed=IA_Generate.generate_image(web_data_Prompt, web_data_Neg_Prompt, web_data_Width, web_data_Height, web_data_Guidance, web_data_Seed, web_data_Step, web_data_NB_img, web_data_Utilisateur, pipe)
            print('Envoi des images au serveur...')
            # Lire les images
            images = os.listdir(
                './images-IA-py/')
            for image in images:
                # Envoyer les images
                with open(
                        './images-IA-py/' +
                        image, 'rb') as f:
                    image_data = f.read()
                    # Envoi de la taille du fichier
                    size = 'size:' + str(len(image_data))
                    s.sendall(size.encode())
                    f.close()
                    time.sleep(1)
                    # Envoi du contenu du fichier
                    s.sendall(image_data)
                print('Image envoyée:', image, flush=True)
                f.close()
            # attendre que le serveur node.js ait fini de lire les images et qu'il nous ai dit 'Python Socket Server received data'
            while True:
                web_data_receive = s.recv(1024)
                if web_data_receive:
                    web_data = web_data_receive.decode()
                    if web_data == 'Python Socket Server received data':
                        break
            print('Envoi des images terminé.')
            # envoyer le nom de l'utilisateur (int)
            time.sleep(0.1)
            utilisateur='utilisateur:'+str(web_data_Utilisateur)
            s.send(utilisateur.encode())
            f.close()
            time.sleep(0.1)
            seed='seed:'+str(seed)
            s.send(seed.encode())
            f.close()
            time.sleep(0.1)
            print('Envoi des images terminé.')
            gpu_status = 'GPU disponible'
            s.send(gpu_status.encode())  # Envoyer le message encodé en bytes
            f.close()
            # effacer les images
            images = os.listdir('./images-IA-py/')
            for image in images:
                os.remove('./images-IA-py/' + image)
            print('Images effacées.', flush=True)

    if message == 'Fin':
        print('Fermeture de la connexion.')
        break

# Fermer la connexion
s.close()
 
