import random
import io
import zipfile
import requests

from PIL import Image


jwt_token = ''
url = "https://api.novelai.net/ai/generate-image"
headers = {}


def set_token(token):
    global jwt_token, headers
    if jwt_token == token:
        return
    jwt_token = token
    headers = {
            "Authorization": f"Bearer {jwt_token}",
            "Content-Type": "application/json",
            "Origin": "https://novelai.net",
            "Referer": "https://novelai.net/"
        }

def generate_novelai_image(
    input_text="", 
    negative_prompt="", 
    seed=-1, 
    scale=5.0, 
    width=1024, 
    height=1024, 
    steps=28, 
    sampler="k_euler",
    schedule='native',
    smea=False,
    dyn=False,
    dyn_threshold=False,
    cfg_rescale=0,
):
    # Assign a random seed if seed is -1
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)

    # Define the payload
    payload = {
        "action": "generate",
        "input": input_text,
        "model": "nai-diffusion-3",
        "parameters": {
            "width": width,
            "height": height,
            "scale": scale,
            "sampler": sampler,
            "steps": steps,
            "n_samples": 1,
            "ucPreset": 0,
            "add_original_image": False,
            "cfg_rescale": cfg_rescale,
            "controlnet_strength": 1,
            "dynamic_thresholding": dyn_threshold,
            "legacy": False,
            "negative_prompt": negative_prompt,
            "noise_schedule": schedule,
            "qualityToggle": True,
            "seed": seed,
            "sm": smea,
            "sm_dyn": dyn,
            "uncond_scale": 1,
        }
    }

    # Send the POST request
    response = requests.post(url, json=payload, headers=headers)

    # Process the response
    if response.headers.get('Content-Type') == 'application/x-zip-compressed':
        zipfile_in_memory = io.BytesIO(response.content)
        with zipfile.ZipFile(zipfile_in_memory, 'r') as zip_ref:
            file_names = zip_ref.namelist()
            if file_names:
                with zip_ref.open(file_names[0]) as file:
                    return file.read(), payload
            else:
                return "NAI doesn't return any images", response
    else:
        return "Generation failed", response




def image_from_bytes(data):
    img_file = io.BytesIO(data)
    img_file.seek(0)
    return Image.open(img_file)