import pickle
from pathlib import Path
import os
import requests
import tqdm

def generate_image(prompt, api_key: str, n=1, size="1024x1024"):
    """
    Generates an image using OpenAI's DALL-E API.

    :param prompt: The prompt to generate an image for.
    :param api_key: Your OpenAI API key.
    :param n: The number of images to generate.
    :param size: The size of the generated images.
    :return: The response from the API call.
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    payload = {
        "model": "dall-e-3",
        "prompt": prompt,
        "n": n,
        "size": size
    }
    
    response = requests.post("https://api.openai.com/v1/images/generations", 
                             headers=headers, 
                             json=payload)
    return response

def get_prompt(persona: dict):
    prompt = f"""A character named {persona['name']} is a {persona['age']} years old. The basic information of this character is as follows:
- Name: {persona['name']}
- Age: {persona['age']}
- Gender: {persona['gender']}
- Routine: {persona['routine']}
- Personality: {', '.join(persona['personality'])}
- Occupation: {persona['occupation']}
- Thoughts: {persona['thoughts']}
- Lifestyle: {persona['lifestyle']}
Please generate a full-face photo of this character.
"""
    return prompt

def get_persona_avatar_bytes(persona: dict, api_key: str) -> bytes:
    prompt = get_prompt(persona)
    response = generate_image(prompt, api_key)
    if response.status_code == 200:
        # 这里可以添加代码来处理响应体，例如保存图像或进一步的处理
        image_data = response.json()['data']
        assert len(image_data) == 1
        image_url = image_data[0]['url']
        # download image
        image_response = requests.get(image_url)
        image = image_response.content
        return image
    else:
        print(f"Failed to generate image: {response.status_code} - {response.text}")

def generate_for_cache(api_key: str):
    root_path = Path("static/exist_characters")
    character_dirs = os.listdir(root_path)
    character_dirs = [root_path / character_dir for character_dir in character_dirs]

    for character_dir in tqdm.tqdm(character_dirs):
        character_path = character_dir / f"{character_dir.name}.pkl"
        character_data = pickle.load(open(character_path, 'rb'))
        image_path = character_dir / f"avatar.jpg"
        image_bytes = get_persona_avatar_bytes(character_data['persona'], api_key)
        if image_bytes:
            with open(image_path, 'wb') as f:
                f.write(image_bytes)


if __name__ == "__main__":
    api_key = os.getenv("OPENAI_API_KEY")
    generate_for_cache(api_key)