import random
import io
import zipfile
import requests
import json
import base64

from PIL import Image


jwt_token = ''
url = "https://image.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,
    ref_image=None,
    info_extract=1,
    ref_str=0.6,
    i2i_image=None,
    i2i_str=0.7,
    i2i_noise=0,
    overlay=True,
    inp_img=None,
    selection='i2i'
):
    # 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": True,
            "cfg_rescale": cfg_rescale,
            "controlnet_strength": 1,
            "dynamic_thresholding": dyn_threshold,
            "params_version": 1,
            "legacy": False,
            "legacy_v3_extend": False,
            "negative_prompt": negative_prompt,
            "noise": i2i_noise,
            "noise_schedule": schedule,
            "qualityToggle": True,
            "reference_information_extracted": info_extract,
            "reference_strength": ref_str,
            "seed": seed,
            "sm": smea,
            "sm_dyn": dyn,
            "uncond_scale": 1,
            "overlay": True
        }
    }
    if ref_image is not None:
        payload['parameters']['reference_image'] = image2base64(ref_image)
    if selection == 'inp' and inp_img is not None:
        payload['action'] = "infill"
        payload['model'] = 'nai-diffusion-3-inpainting'
        payload['parameters']['mask'] = image2base64(inp_img['layers'][0])
        payload['parameters']['image'] = image2base64(inp_img['background'])
        payload['parameters']['extra_noise_seed'] = seed
    if i2i_image is not None and selection == 'i2i':
        payload['action'] = "img2img"
        payload['parameters']['image'] = image2base64(i2i_image)
        payload['parameters']['strength'] = i2i_str
        payload['parameters']['extra_noise_seed'] = seed
    # Send the POST request
    try:
        response = requests.post(url, json=payload, headers=headers, timeout=20)
    except:
        return None, {'message': 'NAI response timeout'}

    # Process the response
    if response.headers.get('Content-Type') == 'binary/octet-stream':
        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", json.loads(response.content)
    else:
        return "Generation failed", json.loads(response.content)




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

def image2base64(img):
    output_buffer = io.BytesIO()
    img.save(output_buffer, format='PNG' if img.mode=='RGBA' else 'JPEG')
    byte_data = output_buffer.getvalue()
    base64_str = base64.b64encode(byte_data).decode()
    return base64_str