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

# Free Google Translate API endpoint (without API key)
GOOGLE_TRANSLATE_API_URL = "https://translate.googleapis.com/translate_a/single?client=gtx&sl=auto&tl=en&dt=t&q="

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

# Function to translate text using Google Translate API
def translate_prompt(prompt):
    url = GOOGLE_TRANSLATE_API_URL + requests.utils.quote(prompt)
    response = requests.get(url)
    translated_text = json.loads(response.text)[0][0][0]
    return translated_text

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):
    translated_prompt = translate_prompt(prompt)  # Translate the prompt to English
    result = prodia_client.generate({
        "prompt": translated_prompt,
        "negative_prompt": 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):
    translated_prompt = translate_prompt(prompt)  # Translate the prompt to English
    result = prodia_client.transform({
        "imageData": image_to_base64(input_image),
        "denoising_strength": denoising,
        "prompt": translated_prompt,
        "negative_prompt": 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, label="Sampling Method", choices=prodia_client.list_samplers())
                                
                            with gr.Column(scale=1):
                                steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=25, value=20, step=1)
    
                        with gr.Row():
                            with gr.Column(scale=1):
     # Function for translating the prompt using the free Google Translate API
def translate_prompt(prompt, target_language="en"):
    try:
        # Set up the translation endpoint and parameters
        url = "https://translate.googleapis.com/translate_a/single"
        params = {
            "client": "gtx",
            "sl": "auto",  # Source language is automatically detected
            "tl": target_language,  # Target language
            "dt": "t",  # Request translation only
            "q": prompt,  # The text to translate
        }

        # Make the GET request to the translation API
        response = requests.get(url, params=params)

        # Extract the translated text from the response
        translated_text = response.json()[0][0][0]
        return translated_text

    except Exception as e:
        print(f"Error translating prompt: {e}")
        return prompt  # In case of error, return the original prompt


# Modify txt2img function to include prompt translation
def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
    translated_prompt = translate_prompt(prompt)
    translated_negative_prompt = translate_prompt(negative_prompt)
    
    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"]


# Modify img2img function to include prompt translation
def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
    translated_prompt = translate_prompt(prompt)
    translated_negative_prompt = translate_prompt(negative_prompt)

    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"]

# Add a translation option in the UI
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")
                    translate_option = gr.Checkbox(label="Translate Prompt to English", value=True)
                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, label="Sampling Method",
                                                      choices=prodia_client.list_samplers())

                            with gr.Column(scale=1):
                                steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=25, value=20, step=1)

                        with gr.Row():
                            with gr.Column(scale=1):
                                width = gr.Slider(label="Width", maximum=1024, value=512, step=8)
                                height = gr.Slider(label="Height", maximum=1024, value=512, step=8)

                            with gr.Column(scale=1):
                                batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
                                batch_count = gr.Slider(label="Batch Count", maximum=1, value=1)

                        cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=1)
                        seed = gr.Number(label="Seed", value=-1)

                with gr.Column(scale=2):
                    image_output = gr.Image(
                        value="https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png")

            text_button.click(fn=txt2img, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed],
                              outputs=image_output, concurrency_limit=64) 
            demo.queue(max_size=40, api_open=False).launch(max_threads=128, show_api=False)