import gradio as gr
from huggingface_hub import InferenceClient
from gradio_client import Client
from PIL import Image
import requests
from io import BytesIO
import re

# Function to fetch the API endpoint from the URL using regular expressions
def fetch_api_endpoint(url):
    try:
        headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'}
        response = requests.get(url, headers=headers)
        response.raise_for_status()  # Raise exception for HTTP errors
        
        # Use regular expressions to find the endpoint URL in the script content
        endpoint_match = re.search(r'root":"(https://[^"]+)', response.text)
        if endpoint_match:
            endpoint_url = endpoint_match.group(1)
            return endpoint_url
        else:
            return None
    except requests.RequestException as e:
        print("Error fetching URL:", e)
        return None
    except Exception as e:
        print("An unexpected error occurred:", e)
        return None

# Define the URL to fetch the API endpoint from
url = "https://playgroundai-playground-v2-5.hf.space/"

# Fetch the API endpoint
endpoint = fetch_api_endpoint(url)

if endpoint is None:
    print("Failed to fetch the API endpoint.")
    exit(1)

# Initialize the HuggingFace Inference Client with the specified model
client_mistral = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
# Initialize the Playground AI client with the fetched endpoint
client_playground = Client(endpoint)

def format_prompt(logo_request):
    system_prompt = """
    You are an advanced language model designed to create detailed and creative image prompts for logo generation. Based on the user's input, generate an elaborate and descriptive image prompt that can be used to create a high-quality icon logo. Ensure that the prompt is clear, imaginative, and provides specific details that will guide the logo creation process effectively. The logo should be a simple, stylized, and abstract icon without any text, focusing solely on graphical elements suitable for a logo.
    """
    prompt = f"<s>[SYS] {system_prompt} [/SYS][INST] {logo_request} [/INST]</s>"
    return prompt

def generate_improved_prompt(logo_request, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = {
        "temperature": temperature,
        "max_new_tokens": max_new_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "seed": 42,
    }

    formatted_prompt = format_prompt(logo_request)
    stream = client_mistral.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
    return output

def generate_image(prompt, width=1024, height=1024, guidance_scale=3):
    result = client_playground.predict(
        prompt,
        "",  # negative prompt
        False,  # use negative prompt
        0,  # seed
        width,  # width
        height,  # height
        guidance_scale,  # guidance scale
        True,  # randomize seed
        api_name="/run"
    )

    # Extract the image URL from the result
    image_path = result[0][0]["image"]
    image_url = endpoint + "/file=" + image_path
    
    # Fetch and display the result image
    response = requests.get(image_url)
    
    if response.headers['Content-Type'].startswith('image'):
        img = Image.open(BytesIO(response.content))
        return img, image_url
    else:
        return None, None

css = """
#mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
}
"""

def process_request(logo_request, width, height, guidance_scale):
    improved_prompt = generate_improved_prompt(logo_request)
    image, image_url = generate_image(improved_prompt, width, height, guidance_scale)
    return image

with gr.Blocks(css=css) as app:
    with gr.Row():
        with gr.Column(scale=2):
            gr.HTML("<h1>Settings</h1>")
            logo_input = gr.Textbox(label="Input your logo request", placeholder="Describe the logo you want...")
            width = gr.Slider(label="Width", minimum=256, maximum=1536, value=1024)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, value=1024)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20, value=3)
            examples = gr.Examples(examples=["coffeeshop website logo", "bio food website logo", "car Selling website logo"], inputs=logo_input)
        
        with gr.Column(scale=3):
            gr.HTML("<h1><center>Magical AI Logo Generator</h1><center>")
            generate_button = gr.Button("Generate")
            image_output = gr.Image(label="Generated Logo")
            generate_button.click(
                fn=process_request,
                inputs=[logo_input, width, height, guidance_scale],
                outputs=[image_output]
            )

    with gr.Row():
        gr.HTML("<h2>Generated Examples</h2>")
        gr.Markdown("""
        | **Preview Image** | **Used Prompt**                | **Width**  | **Height**  | **Guidance Scale** | **Seed** |
        |-------------------|--------------------------------|------------|-------------|--------------------|----------|
        | ![Preview](https://i.imgur.com/CBvJ3sy.png) | Coffeeshop website logo  | 1024 | 1024 | 3 | 0 |
        | ![Preview](https://i.imgur.com/i2XpinX.png) | Bio food website logo   | 1024 | 1024 | 3  | 0 |
        | ![Preview](https://i.imgur.com/HeEjh4T.png) | Car Selling website logo | 1024 | 1024 | 3  | 0 |
        """)

    gr.Markdown("""
    ---
    ### Meta Information
    **Project Title**: Magical AI Logo Generator
    
    **Github**: [https://github.com/pacnimo/](https://github.com/pacnimo/)
    
    **Description**: Magical AI Logo Generator is Free and Easy to Use. Create an AI Website Logo for any Niche. 1 Click Logo Generator.
    
    **Footer**: © 2024 by [pacnimo](https://github.com/pacnimo/). All rights reserved.
    """)

app.launch(debug=True)