import gradio as gr
import requests
import time
import json
import base64
import os
from io import BytesIO
import html
import re

# Translation Function
async def translateToEnglish(text):
    apiUrl = f'https://translate.googleapis.com/translate_a/single?client=gtx&sl=my&tl=en&dt=t&q={encodeURIComponent(text)}'
    try:
        response = await fetch(apiUrl)
        data = await response.json()
        if data and data[0] and data[0][0] and data[0][0][0]:
            return data[0][0][0]
        else:
            raise Error('Translation to English failed')
    except Exception as error:
        print('Translation to English failed:', error)
        raise error

class Prodia:
    def __init__(self, api_key, base=None):
        self.base = base or "https://api.prodia.com/v1"
        self.headers = {
            "X-Prodia-Key": api_key
        }
    
    def generate(self, params):
        response = self._post(f"{self.base}/sd/generate", params)
        return response.json()
    
    def transform(self, params):
        response = self._post(f"{self.base}/sd/transform", params)
        return response.json()
    
    def controlnet(self, params):
        response = self._post(f"{self.base}/sd/controlnet", params)
        return response.json()
    
    def get_job(self, job_id):
        response = self._get(f"{self.base}/job/{job_id}")
        return response.json()

    def wait(self, job):
        job_result = job

        while job_result['status'] not in ['succeeded', 'failed']:
            time.sleep(0.25)
            job_result = self.get_job(job['job'])

        return job_result

    def list_models(self):
        response = self._get(f"{self.base}/sd/models")
        return response.json()

    def list_samplers(self):
        response = self._get(f"{self.base}/sd/samplers")
        return response.json()

    def _post(self, url, params):
        headers = {
            **self.headers,
            "Content-Type": "application/json"
        }
        response = requests.post(url, headers=headers, data=json.dumps(params))

        if response.status_code != 200:
            raise Exception(f"Bad Prodia Response: {response.status_code}")

        return response

    def _get(self, url):
        response = requests.get(url, headers=self.headers)

        if response.status_code != 200:
            raise Exception(f"Bad Prodia Response: {response.status_code}")

        return response

def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue())
    return img_str.decode('utf-8')

def remove_id_and_ext(text):
    text = re.sub(r'.*$', '', text)
    extension = text[-12:].strip()
    if extension == "safetensors":
        text = text[:-13]
    elif extension == "ckpt":
        text = text[:-4]
    return text

def get_data(text):
    results = {}
    patterns = {
        'prompt': r'(.*)',
        'negative_prompt': r'Negative prompt: (.*)',
        'steps': r'Steps: (\d+),',
        'seed': r'Seed: (\d+),',
        'sampler': r'Sampler:\s*([^\s,]+(?:\s+[^\s,]+)*)', 
        'model': r'Model:\s*([^\s,]+)',
        'cfg_scale': r'CFG scale:\s*([\d\.]+)',
        'size': r'Size:\s*([0-9]+x[0-9]+)'
        }
    for key in ['prompt', 'negative_prompt', 'steps', 'seed', 'sampler', 'model', 'cfg_scale', 'size']:
        match = re.search(patterns[key], text)
        if match:
            results[key] = match.group(1)
        else:
            results[key] = None
    if results['size'] is not None:
        w, h = results['size'].split("x")
        results['w'] = w
        results['h'] = h
    else:
        results['w'] = None
        results['h'] = None
    return results

def send_to_txt2img(image):
    result = {tabs: gr.update(selected="t2i")}
    try:
        text = image.info['parameters']
        data = get_data(text)
        result[prompt] = gr.update(value=data['prompt'])
        result[negative_prompt] = gr.update(value=data['negative_prompt']) if data['negative_prompt'] is not None else gr.update()
        result[steps] = gr.update(value=int(data['steps'])) if data['steps'] is not None else gr.update()
        result[seed] = gr.update(value=int(data['seed'])) if data['seed'] is not None else gr.update()
        result[cfg_scale] = gr.update(value=float(data['cfg_scale'])) if data['cfg_scale'] is not None else gr.update()
        result[width] = gr.update(value=int(data['w'])) if data['w'] is not None else gr.update()
        result[height] = gr.update(value=int(data['h'])) if data['h'] is not None else gr.update()
        result[sampler] = gr.update(value=data['sampler']) if data['sampler'] is not None else gr.update()
        if model in model_names:
            result[model] = gr.update(value=model_names[model])
        else:
            result[model] = gr.update()
        return result
    except Exception as e:
        print(e)
        return result

prodia_client = Prodia(api_key=os.getenv("PRODIA_API_KEY"))
model_list = prodia_client.list_models()
model_names = {}
for model_name in model_list:
    name_without_ext = remove_id_and_ext(model_name)
    model_names[name_without_ext] = model_name

def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
    # Translate prompt from Myanmar to English
    translated_prompt = asyncio.run(translateToEnglish(prompt))
    translated_negative_prompt = asyncio.run(translateToEnglish(negative_prompt)) if negative_prompt else None

    result = prodia_client.generate({
        "prompt": translated_prompt,
        "negative_prompt": translated_negative_prompt,
        "model": model,
        "steps": steps,
        "sampler": sampler,
        "cfg_scale": cfg_scale,
        "width": width,
        "height": height,
        "seed": seed
    })

    job = prodia_client.wait(result)
    return job["imageUrl"]

def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
    # Translate prompt from Myanmar to English
    translated_prompt = asyncio.run(translateToEnglish(prompt))
    translated_negative_prompt = asyncio.run(translateToEnglish(negative_prompt)) if negative_prompt else None

    result = prodia_client.transform({
        "imageData": image_to_base64(input_image),
        "denoising_strength": denoising,
        "prompt": translated_prompt,
        "negative_prompt": translated_negative_prompt,
        "model": model,
        "steps": steps,
        "sampler": sampler,
        "cfg_scale": cfg_scale,
        "width": width,
        "height": height,
        "seed": seed
    })

    job = prodia_client.wait(result)
    return job["imageUrl"]

css = """
#generate {
    height: 100%;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column(scale=6):
            model = gr.Dropdown(interactive=True, value="absolutereality_v181.safetensors [3d9d4d2b]", show_label=True, label="Stable Diffusion Checkpoint", choices=prodia_client.list_models())
        with gr.Column(scale=1):
            gr.Markdown(elem_id="powered-by-prodia", value="AUTOMATIC1111 Stable Diffusion Web UI.<br>Powered by [Prodia](https://prodia.com).<br>For more features and faster generation times check out our [API Docs](https://docs.prodia.com/reference/getting-started-guide).")

    with gr.Tabs() as tabs:
        with gr.Tab("txt2img", id='t2i'):
            with gr.Row():
                with gr.Column(scale=6, min_width=600):
                    prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k", placeholder="Prompt", show_label=False, lines=3)
                    negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, value="3d, cartoon, anime, (deformed eyes, nose, ears, nose), bad anatomy, ugly")
                with gr.Column():
                    text_button = gr.Button("Generate", variant='primary', elem_id="generate")
                    
            with gr.Row():
                with gr.Column(scale=3):
                    with gr.Tab("Generation"):
                        with gr.Row():
                            with gr.Column(scale=1):
                                sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True
                                                   sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True, label="Sampler", choices=prodia_client.list_samplers())
                            with gr.Column(scale=1):
                                steps = gr.Slider(minimum=1, maximum=150, value=50, step=1, show_label=True, label="Steps")
                            with gr.Column(scale=1):
                                cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, value=7.5, step=0.1, show_label=True, label="CFG Scale")
                        with gr.Row():
                            with gr.Column(scale=1):
                                width = gr.Slider(minimum=256, maximum=2048, value=1024, step=64, show_label=True, label="Width")
                            with gr.Column(scale=1):
                                height = gr.Slider(minimum=256, maximum=2048, value=1024, step=64, show_label=True, label="Height")
                            with gr.Column(scale=1):
                                seed = gr.Number(value=-1, show_label=True, label="Seed", precision=0)
                with gr.Column():
                    output_img = gr.Image(label="Generated Image", interactive=False)
                    output_img.style(height=512)
                    output_gallery = gr.Gallery(show_label=False).style(grid=[2], height="auto")
                    
        text_button.click(fn=txt2img, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=[output_img])
        
        with gr.Tab("img2img"):
            with gr.Row():
                with gr.Column(scale=6):
                    input_image = gr.Image(source='upload', label="Input Image", type="pil", interactive=True)
                    prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k", placeholder="Prompt", show_label=False, lines=3)
                    negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, value="3d, cartoon, anime, (deformed eyes, nose, ears, nose), bad anatomy, ugly")
                with gr.Column():
                    denoising = gr.Slider(minimum=0.1, maximum=1.0, value=0.75, step=0.01, show_label=True, label="Denoising Strength")
                    img2img_button = gr.Button("Transform", variant='primary')

            with gr.Row():
                with gr.Column(scale=3):
                    sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True, label="Sampler", choices=prodia_client.list_samplers())
                    steps = gr.Slider(minimum=1, maximum=150, value=50, step=1, show_label=True, label="Steps")
                    cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, value=7.5, step=0.1, show_label=True, label="CFG Scale")
                    width = gr.Slider(minimum=256, maximum=2048, value=1024, step=64, show_label=True, label="Width")
                    height = gr.Slider(minimum=256, maximum=2048, value=1024, step=64, show_label=True, label="Height")
                    seed = gr.Number(value=-1, show_label=True, label="Seed", precision=0)
            with gr.Row():
                output_img_img2img = gr.Image(label="Transformed Image", interactive=False)
                
        img2img_button.click(fn=img2img, inputs=[input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=[output_img_img2img])

    with gr.Tab("Settings"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("# Model Settings")
                model = gr.Dropdown(interactive=True, value="absolutereality_v181.safetensors [3d9d4d2b]", show_label=True, label="Stable Diffusion Checkpoint", choices=prodia_client.list_models())

demo.launch()