import time  
import requests
import json

URL = 'xxx'

def save_webpage(url, file_path):
    try:
        response = requests.get(url)
        if response.status_code == 200:
            with open(file_path, 'wb') as file:
                file.write(response.content)
            print(f"Webpage saved to {file_path}")
        else:
            print(f"Failed to retrieve the webpage.")
    except requests.exceptions.RequestException as e:
        print(f"An error occurred: {e}")

class MidJourneyBot(object):  
    def __init__(self):  
        pass
        
    def mj_create_img(self, query):  
        url = URL + "/submit/imagine"  
        if query.endswith("."):  
            query = query[:-1]
        if '--relax' not in query:
            query += ' --relax'
        bodyParams = {
            "prompt": query  
        }  
        response = requests.post(url, json=bodyParams)  
        print(response)  
        if (response.status_code == 200):  
            data = response.json()
            print(data)
            if (data['code'] == 1):  
                job_id = data['result']  
                print(job_id)
                return self.check_job_status(job_id)  
            return False
        return False

    def check_job_status(self, job_id):  
        url = f"{URL}/task/{job_id}/fetch"  
        try:  
            while True:  
                response = requests.get(url)
                data = response.json()
                progress = data['progress']
                status = data['status']
                if progress == '100%':
                    return data['imageUrl']
                if status == 'FAILURE':
                    return "False"
                if response.status_code == 200:  
                    data = response.json() 
                    print(data)                 
                else:  
                    raise Exception('Failed to check job status')  
                time.sleep(5)
        except Exception as e:  
            print(e)
            
file_in = "xxx"
with open(file_in, 'r') as file:
    data_list = json.load(file)

result_dict = []
MJ = MidJourneyBot()
for data in data_list:
    prompt = data['prompt']
    a = data['A']
    b = data['B']
    for i in range(5):
        sentence = prompt[i]
        url_to_save = img = MJ.mj_create_img(sentence)
        result = {"A": a, "B": b,'prompt':prompt[i],'img': img}
        result_dict.append(result)
    
        output_file_path = "images/f{prompt}.png"
        save_webpage(url_to_save, output_file_path)

json_str = json.dumps(result_dict)  
file_out = "xxx"
with open(file_out, "w") as file:
    file.write(json_str)