import paramiko
import numpy as np
import io, os
from PIL import Image
import requests
import json
import random
from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_MSCOCO

def create_ssh_client(server, port, user, password):
    ssh = paramiko.SSHClient()
    ssh.load_system_host_keys()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(server, port, user, password)
    return ssh

def get_image_from_url(image_url):
    response = requests.get(image_url)
    response.raise_for_status()  # success
    return Image.open(io.BytesIO(response.content))

# def upload_image_via_sftp(ssh, image, remote_image_path):
#     if isinstance(image, str):
#         print("get url image")
#         image = get_image_from_url(image)
#     with ssh.open_sftp() as sftp:
#         with io.BytesIO() as image_byte_stream:
#             image.save(image_byte_stream, format='JPEG')
#             image_byte_stream.seek(0)  
#             sftp.putfo(image_byte_stream, remote_image_path)
#         print(f"Successfully uploaded image to {remote_image_path}")
        
# def upload_json_via_sftp(ssh, data, remote_json_path):
#     json_data = json.dumps(data, indent=4)
#     with ssh.open_sftp() as sftp:
#         with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
#             sftp.putfo(json_byte_stream, remote_json_path)
#         print(f"Successfully uploaded JSON data to {remote_json_path}")

def upload_image(states, output_dir):
    pass
    # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    # for i in range(len(states)):
    #     output_file = os.path.join(output_dir, f"{i}.jpg")
    #     upload_image_via_sftp(ssh, states[i].output, output_file)

def upload_informance(data, data_path):
    pass
    # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    # upload_json_via_sftp(ssh, data, data_path)
def get_random_mscoco_prompt():
    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    sftp = ssh.open_sftp()
    files = sftp.listdir(SSH_MSCOCO)
    txt_files = [file for file in files if file.endswith('.txt')]

    selected_files = random.sample(txt_files, 1) # get one prompt

    for file in selected_files:
        remote_file_path = os.path.join(SSH_MSCOCO, file)
        with sftp.file(remote_file_path, 'r') as f:
            content = f.read().decode('utf-8')
            print(f"Content of {file}:")
            print("\n")
    sftp.close()
    ssh.close()

    return content


def create_remote_directory(remote_directory):
    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)

    stdin, stdout, stderr = ssh.exec_command(f'mkdir -p {SSH_LOG}/{remote_directory}')
    error = stderr.read().decode('utf-8')
    if error:
        print(f"Error: {error}")
    else:
        print(f"Directory {remote_directory} created successfully.")
    # ssh.close()
    return f'{SSH_LOG}/{remote_directory}'
    
def upload_ssh_all(states, output_dir, data, data_path):
    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    output_file_list = []
    image_list = []
    for i in range(len(states)):
        output_file = os.path.join(output_dir, f"{i}.jpg")
        output_file_list.append(output_file)
        image_list.append(states[i].output)

    with ssh.open_sftp() as sftp:
        for i in range(len(output_file_list)):
            if isinstance(image_list[i], str):
                print("get url image")
                image_list[i] = get_image_from_url(image_list[i])
            with io.BytesIO() as image_byte_stream:
                image_list[i].save(image_byte_stream, format='JPEG')
                image_byte_stream.seek(0)  
                sftp.putfo(image_byte_stream, output_file_list[i])
            print(f"Successfully uploaded image to {output_file_list[i]}")
        json_data = json.dumps(data, indent=4)
        with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
            sftp.putfo(json_byte_stream, data_path)
        print(f"Successfully uploaded JSON data to {data_path}")