import os
import re
import spaces
import random
import string
import torch
import requests
import gradio as gr
import numpy as np
from lxml.html import fromstring
#from transformers import pipeline
from torch import multiprocessing as mp
from torch.multiprocessing import Pool as Pool
#from pathos.multiprocessing import ProcessPool as Pool
#from pathos.threading import ThreadPool as Pool
#from diffusers.pipelines.flux import FluxPipeline
#from diffusers.utils import export_to_gif
#from huggingface_hub import hf_hub_download
#from safetensors.torch import load_file
from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
from diffusers.utils import load_image
#import jax
#import jax.numpy as jnp

def port_inc():
    env = os.getenv("CUSTOM_PORT")
    if env == None:
        os.environ["CUSTOM_PORT"]="7860"
    else:
        os.environ["CUSTOM_PORT"]=str(int(env)+1)

def pipe_t2i():
    PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
    return PIPE
    
def pipe_i2i():
    PIPE = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
    PIPE.unet = torch.compile(PIPE.unet, mode="reduce-overhead", fullgraph=True)
    return PIPE

def translate(text,lang):
    if text == None or lang == None:
        return ""       
    text = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', text)).lower().strip()
    lang = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', lang)).lower().strip()    
    if text == "" or lang == "":
        return ""
    if len(text) > 38:
        raise Exception("Translation Error: Too long text!")
    user_agents = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 13_1) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15'
    ]
    padded_chars = re.sub("[(^\-)(\-$)]","",text.replace("","-").replace("- -"," ")).strip()
    query_text = f'Please translate {padded_chars}, into {lang}'
    url = f'https://www.google.com/search?q={query_text}'
    resp = requests.get(
        url = url,
        headers = {
            'User-Agent': random.choice(user_agents)
        }
    )
    content = resp.content
    html = fromstring(content)
    translated = text
    try:
        src_lang = html.xpath('//*[@class="source-language"]')[0].text_content().lower().strip()
        trgt_lang = html.xpath('//*[@class="target-language"]')[0].text_content().lower().strip()
        src_text = html.xpath('//*[@id="tw-source-text"]/*')[0].text_content().lower().strip()
        trgt_text = html.xpath('//*[@id="tw-target-text"]/*')[0].text_content().lower().strip()
        if trgt_lang == lang:
            translated = trgt_text
    except:
        print(f'Translation Warning: Failed To Translate!')
    ret = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', translated)).lower().strip()
    print(ret)
    return ret

def generate_random_string(length):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))

@spaces.GPU(duration=75,enable_queue=True)
def Piper(name,name2,posi,neg):
    print("starting piper")
    
    ret1 = pp1(
        posi,
        height=512,
        width=512,
        num_inference_steps=4,
        max_sequence_length=256,
        guidance_scale=0
    )
    ret1.images[0].save(name)
    
    img = load_image(name).convert("RGB")
    
    ret2 = pp2(
        prompt=posi,
        negative_prompt=neg,
        image=img
    )
    ret2.images[0].save(name2)

def tok(pipe,txt):
    toks = pipe.tokenizer(txt)['input_ids']
    print(toks)
    return toks

css="""
input, input::placeholder {
    text-align: center !important;
}
*, *::placeholder {
    direction: ltr !important;
    font-family: Suez One !important;
}
h1,h2,h3,h4,h5,h6,span,p,pre {
    width: 100% !important;
    text-align: center !important;
    display: block !important;
}
footer {
    display: none !important;
}
#col-container {
    margin: 0 auto !important;
    max-width: 15cm !important;
 }
.image-container {
    aspect-ratio: 512 / 512 !important;
}
.dropdown-arrow {
    display: none !important;
}
*:has(.btn), .btn {
    width: 100% !important;
    margin: 0 auto !important;
}
"""

js="""
function custom(){
    document.querySelector("div#prompt input").setAttribute("maxlength","38")
    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
}
"""

def infer(p):
        print("infer: started")
    
        p1 = p["a"]
        p2 = p["b"]
        name = generate_random_string(12)+".png"
        name2 = "_"+name

        _do = ['beautiful', 'playful', 'photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable coloring', 'rough texture', 'best quality', 'focused']
        if p1 != "":
            _do.append(f'{p1}')
        posi = " ".join(_do)

        neg = " "
        if p2 != "":
            neg = f'{p2} where in {p1}'
        
        Piper(name,name2,posi,neg)
    
        return name2

def run(p1,p2,*result):
    
        p1_en = translate(p1,"english")
        p2_en = translate(p2,"english")
        p = {"a":p1_en,"b":p2_en}
        ln = len(result)
        print("images: "+str(ln))
        rng = list(range(ln))
    
        arr = [p for _ in rng]
        pool = Pool(ln)
        out = list(pool.imap(infer,arr))
        pool.close()
        pool.join()
        pool.clear()

        return out

def main():

    global result
    global pp1
    global pp2

    result=[]
    pp1=pipe_t2i()
    pp2=pipe_i2i()

    mp.set_start_method("spawn", force=True)
    
    with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
        with gr.Column(elem_id="col-container"):
            gr.Markdown(f"""
                # MULTI-LANGUAGE IMAGE GENERATOR
            """)
            with gr.Row():
                prompt = gr.Textbox(
                    elem_id="prompt",
                    placeholder="INCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                prompt2 = gr.Textbox(
                    elem_id="prompt2",
                    placeholder="EXCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                run_button = gr.Button("START",elem_classes="btn",scale=0)
            with gr.Row():
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))

        run_button.click(fn=run,inputs=[prompt,prompt2,*result],outputs=result)
        demo.queue().launch()

if __name__ == "__main__":
    main()