import gradio as gr
import time
import base64
from io import BytesIO
import requests
import json
import io
import time
from io import BytesIO
from PIL import Image
from stability_sdk import client
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
import re
import random
import datetime
import pytz
from proxy_randomizer import RegisteredProviders
from cryptography.fernet import Fernet
import os

sd_decode = os.environ.get('sd_decode')

def freeimagehost(prompt, img, seed):
    try:
        rp = RegisteredProviders()
        rp.parse_providers()
        # get one proxy
        proxy = rp.get_random_proxy()
        proxy = proxy.get_proxy()
        # convert proxy to dictionary
        proxy = proxy.split(':')
        # convert list to dictionary
        proxy = {'http': 'http://' + proxy[0] + ':' + proxy[1]}
    except:
        print("Error getting proxy")
        proxy = {}

    try:
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue())
        key = "6d207e02198a847aa98d0a2a901485a5"
        payload = {'key': key, 'image': img_str, 'format': 'json', 'expiration': '3600'}
        r = requests.post("https://freeimage.host/api/1/upload", data=payload, proxies=proxy)
        r = json.loads(r.text)
        try:
            link = loglink = r['image']['url']
            link = '<a style="color:blue;" href="https://arinteli.com/app/?' + link + '+' + prompt + '+' + seed + '" target=”_blank”>'+ link +'</a>'
        except:
            return imagebb(prompt, img, seed)
    except: 
        return imagebb(prompt, img, seed)
    return link, loglink

bbtries = 0
def imagebb(prompt, img, seed):
    global bbtries
    try:
        rp = RegisteredProviders()
        rp.parse_providers()
        # get one proxy
        proxy = rp.get_random_proxy()
        proxy = proxy.get_proxy()
        # convert proxy to dictionary
        proxy = proxy.split(':')
        # convert list to dictionary
        proxy = {'http': 'http://' + proxy[0] + ':' + proxy[1]}
    except:
        print("Error getting proxy")
        proxy = {}

    try:
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue())
        with open('bb_encode.txt') as f:
            bbkeys = f.readlines()
            bbkey = random.choice(bbkeys)
            bbkeyid = bbkeys.index(bbkey) + 1
            print(bbkeyid)
            bbkey = Fernet(sd_decode).decrypt(bbkey.encode())
            # convert the sdkey from bytes to a string
            bbkey = bbkey.decode()
            bbkey = bbkey.strip()
        payload = {'key': bbkey, 'image': img_str, 'format': 'json', 'expiration': '3600'}
        r = requests.post("https://api.imgbb.com/1/upload", data=payload, proxies=proxy)
        r = json.loads(r.text)
        try:
            link = loglink = r['data']['url']
            link = '<a style="color:blue;" href="https://arinteli.com/app/?' + link + '+' + prompt + '+' + seed + '" target=”_blank”>'+ link +'</a>'
        except:
            print(r)
            bbtries += 1
            print(bbtries)
            if bbtries < 5:
                return imagebb(prompt, img, seed)
            else:
                bbtries = 0
                return imgur(prompt, img, seed)
    except Exception as e: 
        print(e)
        bbtries += 1
        if bbtries < 5:
            return imagebb(prompt, img, seed)
        else:
            bbtries = 0
            return imgur(prompt, img, seed)
    return link, loglink

def imgur(prompt, img, seed):
    try:
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue())
        url = "https://api.imgur.com/3/image"
        payload={'image': img_str}
        files=[
        ]
        headers = {
        'Authorization': 'Client-ID ab515931f7df961'
        }
        response = requests.request("POST", url, headers=headers, data=payload, files=files)
        response = json.loads(response.text)
        try:
            link = loglink = response['data']['link']
            link = '<a style="color:blue;" href="https://arinteli.com/app/?' + link + '+' + prompt + '+' + seed + '" target=”_blank”>'+ link +'</a>'
        except:
            link = loglink = "Error generating link (service rebooting)"
            
    except: 
        link = loglink = "Error (2) generating link"
    return link, loglink
    
sdtries = 0
def sd(prompt, negprompt, version):
    if prompt == "":
        return Image.open('error.png'), "Error! Please enter a prompt"

    london = pytz.timezone('Europe/London')
    now = datetime.datetime.now(london)
    timestamp = (now.strftime("%H:%M:%S"))
    timestamp = str(timestamp)
    with open("keys_encode.txt", "r") as f:
        sdkeys = f.readlines()
        sdkey = random.choice(sdkeys)
        keyid = sdkeys.index(sdkey) + 1
        keyid = str(keyid)
        sdkey = sdkey.strip()
        # decode the sdkey with the Fernet key
        sdkey = Fernet(sd_decode).decrypt(sdkey.encode())
        # convert the sdkey from bytes to a string
        sdkey = sdkey.decode()

    if version == "2.1":
        engine = "stable-diffusion-512-v2-1"
    if version == "1.5":
        engine = "stable-diffusion-v1-5"

    # connect to dream API
    stability_api = client.StabilityInference(
        key = sdkey,
        verbose=True,
        engine=engine,
    )
    # convert string prompt to lowercase
    prompt = prompt.lower()
    negprompt = negprompt.lower()

    # set random seed
    seed = random.randint(0, 9999999)

    # create list bad_words_list from bad_words.txt and strip whitespace
    with open('bad_words.txt') as f:
        bad_words_list = f.readlines()
        bad_words_list = [x.strip() for x in bad_words_list]

        # check if bad_words_list is in prompt
        if any(re.findall(r'\b{}\b'.format(bad_word), prompt) for bad_word in bad_words_list):
            # check if any bad words are in prompt
            blocked_words = []
            for bad_word in bad_words_list:
                if re.findall(r'\b{}\b'.format(bad_word), prompt):
                    blocked_words.append(bad_word)
            
               
            blocked_words = ', '.join(blocked_words)
            print(keyid + ": " + timestamp + ": BLOCKED PROMPT \"" + prompt + "\" BANNED WORDS: " + blocked_words)  
            time.sleep(3)
            print("User has had their time wasted.\n")
            # javascript to log error to console
            return Image.open('unsafe.png'), "Your prompt contains blocked words! Please try a different prompt!"

        else:
            try:
                # send prompt to dream API
                # set random seed
                answers = stability_api.generate(
                    #prompt=prompt,
                    steps=15,
                    sampler=generation.SAMPLER_K_DPMPP_2S_ANCESTRAL,
                    guidance_preset=generation.GUIDANCE_PRESET_FAST_GREEN,
                    seed=seed,
                    prompt = [generation.Prompt(text=prompt,parameters=generation.PromptParameters(weight=5)), 
                    generation.Prompt(text=negprompt,parameters=generation.PromptParameters(weight=-5))],
                    )
        

                # convert seed to string
                seed = str(seed)

                # iterating over the generator produces the api response
                for resp in answers:
                    for artifact in resp.artifacts:
                        if artifact.finish_reason == generation.FILTER:
                            print("FILTERED PROMPT \"" + prompt + "\"")
                            try:
                                img = Image.open(io.BytesIO(artifact.binary))
                                print("API image filter\n")
                                return img, "<b>Sorry, Stability.AI detected your image as unsafe (likely a false positive). Please try regenerating your image again, or adjusting your prompt.</b>", seed
                            except:
                                img = Image.open('unsafe.png')
                                print("API word filter\n")
                                return img, "<b>Sorry, Stability.AI detected unsafe words in your prompt. Please remove them and try again.</b>", seed

                        if artifact.type == generation.ARTIFACT_IMAGE:
                            img = Image.open(io.BytesIO(artifact.binary))
                            link = freeimagehost(prompt, img, seed)
                            # time.sleep(2)
                            print(version + "," + keyid + ": " + timestamp + ": PROMPT \"" + prompt + ", Neg: " + negprompt + "\" WAS GENERATED SUCCESSFULLY AND IMAGE LINK: \n" + link[1] + "\n")
                            return img, 'Your download URL: ' + link[0] + '<br> Seed: ' + seed
            except Exception as e:
                print(keyid, '\n' ,e)
                global sdtries
                sdtries += 1
                if sdtries < 5:
                    print("Retrying...")
                    return sd(prompt, negprompt)
                else:
                    sdtries = 0
                    return Image.open('error.png'), "Error! Please try again"
    


examples = [["A hyperrealistic photograph of German architectural modern home in the suburbs of Hamburg, Germany, lens flares, cinematic, hdri, matte painting, concept art, celestial, soft render, highly detailed, cgsociety, octane render, trending on artstation, architectural HD, HQ, 4k, 8k", "tree, bush, woods, branches, leaves, green, forest", "2.1"],
["A dream of a distant galaxy, by Caspar David Friedrich, matte painting trending on artstation HQ", "", "1.5"],
["A Hyperrealistic photograph of ancient Malaysian architectural ruins in Borneo's East Malaysia, lens flares, cinematic, hdri, matte painting, concept art, celestial, soft render, highly detailed, cgsociety, octane render, trending on artstation, architectural HD, HQ, 4k, 8k", "", "2.1"],
["A dream of a distant galaxy, by Caspar David Friedrich, matte painting trending on artstation HQ", "", "1.5"],
["A small cabin on top of a snowy mountain in the style of Disney, artstation", "", "2.1"],
["A Hyperrealistic photograph of ancient Tokyo/London/Paris architectural ruins in a flooded apocalypse landscape of dead skyscrapers, lens flares, cinematic, hdri, matte painting, concept art, celestial, soft render, highly detailed, cgsociety, octane render, trending on artstation, architectural HD, HQ, 4k, 8k", "", "2.1"],
["a portrait of a beautiful blonde woman, fine - art photography, soft portrait shot 8 k, mid length, ultrarealistic uhd faces, unsplash, kodak ultra max 800, 85 mm, intricate, casual pose, centered symmetrical composition, stunning photos, masterpiece, grainy, centered composition", "blender, cropped, lowres, poorly drawn face, out of frame, poorly drawn hands, blurry, bad art, blurred, text, watermark, disfigured, deformed, closed eyes", "2.1"],
["A hyperrealistic painting of an astronaut inside of a massive futuristic metal mechawarehouse, cinematic, sci-fi, lens flares, rays of light, epic, matte painting, concept art, celestial, soft render, octane render, trending on artstation, 4k, 8k", "blender, cropped, lowres, out of frame, blurry, bad art, blurred, text, disfigured, deformed", "2.1"]]

with gr.Blocks(
) as demo:
    # create radio for to use either v1.5 or v2.0
    version = gr.Radio(
        label="Stable Diffusion Version",
        choices=["1.5", "2.1"],
        value="1.5",
        elem_id="version",
    )
    
    download = gr.HTML(elem_id="download")
    output = gr.Image(label="Image Generation", elem_id="output")
    name = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate. Longer and more detailed prompts work better.", elem_id="name") 
    negprompt = gr.Textbox(label="Negative Prompt", placeholder="Describe the image you want to avoid. Longer and more detailed prompts work better.", elem_id="negprompt")
    greet_btn = gr.Button("Generate Image", elem_id="greet_btn")
    bin = gr.HTML("", visible=False)
    ex = gr.Examples(examples=examples, fn=sd, inputs=[name, negprompt, version], outputs=[output, download], cache_examples=True)
    greet_btn.click(
            None,
            [],
            bin,
            _js="""
            () => {
                /*! js-cookie v3.0.1 | MIT */
                !function(e,t){"object"==typeof exports&&"undefined"!=typeof module?module.exports=t():"function"==typeof define&&define.amd?define(t):(e=e||self,function(){var n=e.Cookies,o=e.Cookies=t();o.noConflict=function(){return e.Cookies=n,o}}())}(this,(function(){"use strict";function e(e){for(var t=1;t<arguments.length;t++){var n=arguments[t];for(var o in n)e[o]=n[o]}return e}return function t(n,o){function r(t,r,i){if("undefined"!=typeof document){"number"==typeof(i=e({},o,i)).expires&&(i.expires=new Date(Date.now()+864e5*i.expires)),i.expires&&(i.expires=i.expires.toUTCString()),t=encodeURIComponent(t).replace(/%(2[346B]|5E|60|7C)/g,decodeURIComponent).replace(/[()]/g,escape);var c="";for(var u in i)i[u]&&(c+="; "+u,!0!==i[u]&&(c+="="+i[u].split(";")[0]));return document.cookie=t+"="+n.write(r,t)+c}}return Object.create({set:r,get:function(e){if("undefined"!=typeof document&&(!arguments.length||e)){for(var t=document.cookie?document.cookie.split("; "):[],o={},r=0;r<t.length;r++){var i=t[r].split("="),c=i.slice(1).join("=");try{var u=decodeURIComponent(i[0]);if(o[u]=n.read(c,u),e===u)break}catch(e){}}return e?o[e]:o}},remove:function(t,n){r(t,"",e({},n,{expires:-1}))},withAttributes:function(n){return t(this.converter,e({},this.attributes,n))},withConverter:function(n){return t(e({},this.converter,n),this.attributes)}},{attributes:{value:Object.freeze(o)},converter:{value:Object.freeze(n)}})}({read:function(e){return'"'===e[0]&&(e=e.slice(1,-1)),e.replace(/(%[\dA-F]{2})+/gi,decodeURIComponent)},write:function(e){return encodeURIComponent(e).replace(/%(2[346BF]|3[AC-F]|40|5[BDE]|60|7[BCD])/g,decodeURIComponent)}},{path:"/"})}));
                // create cookie
                if (Cookies.get('creations') == undefined) {
                Cookies.set('creations', '1', { expires: 250 / 24 / 60 / 60 });
                } else {
                Cookies.set('creations', parseInt(Cookies.get('creations')) + 1, { expires: 250 / 24 / 60 / 60 }); }
                console.log(Cookies.get('creations'));
            }""")
    greet_btn.click(fn=sd, inputs=[name, negprompt, version], outputs=[output, download])


demo.queue(concurrency_count=3, max_size=20).launch(max_threads=150)