import gradio as gr
from huggingface_hub import InferenceClient
import base64
from io import BytesIO
from PIL import Image
import requests
import random
import bs4
import lxml
# Define the list of models
models = ["Qwen/Qwen2.5-Coder-32B-Instruct"]
im_models = [
    "black-forest-labs/FLUX.1-dev",
    "ostris/Flex.1-alpha",
]
def format_prompt(message, mod, system):
    eos=f"<|im_end|>\n"
    bos=f"<|im_start|>\n"
    prompt=""
    prompt+=bos
    prompt+=system
    prompt+=eos
    prompt+=bos
    prompt += message
    prompt+=eos
    prompt+=bos
    return prompt
def get_webpage_text(url):
    print("url ", url)
    source = requests.get(url)
    #isV('status: ', source.status_code)
    if source.status_code ==200:
        soup = bs4.BeautifulSoup(source.content,'lxml')
        rawp=f'RAW TEXT RETURNED: {soup.text}'
        return rawp
    else:
        return "ERROR couldn't find, "+url

def generate_prompt(company_name, company_html, company_descp):
    client = InferenceClient(models[0])
    output=""
    system_prompt=f"""You are a Master Generative Image Prompt Writer, you know just the perfect prompt secrets for every situation
    Today you will be generating Company Logo's
    You will be given a Company Name, Description, and HTML artifacts from their website, use this to generate a sufficiently long and detailed image generation prompt to satisfy the users request, make sure that the company name is the focal point of the image
    """
    prompt=f"""Company Name: {company_name}
    Company Description: {company_descp}
    HTML from Company Website: {str(company_html)}"""

    formatted_prompt=format_prompt(prompt,0,system_prompt)
    generate_kwargs = dict(
        temperature=0.99,
        max_new_tokens=512,
        top_p=0.99,
        repetition_penalty=1.0,
        do_sample=True,
        seed=random.randint(1,100000000000000),
    )
    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
    for response in stream:
        output += response.token.text
    return output

        
def generate_image(comp_name, comp_desc="", comp_url="", im_model=""):
    print(comp_name +"|"+comp_desc+"|"+comp_url)
    if comp_url != "":
        web_out=get_webpage_text(comp_url)
    else:
        web_out="No Web Data returned"
    print(web_out)
    prompt_out=generate_prompt(comp_name,web_out, comp_desc)
    client = InferenceClient(im_model)
    response = client.text_to_image(prompt_out)
    return response

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Smart Logo Maker")
    with gr.Row():
        with gr.Column():
            comp_name = gr.Textbox(label="Enter Company Name")
            comp_desc = gr.Textbox(label="Enter Company Description (optional)")
            comp_url = gr.Textbox(label="Enter Company URL (optional)")
            im_model = gr.Dropdown(label="Image Model", choices=im_models, interactive=True)
            generate_button = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    generate_button.click(generate_image, [comp_name, comp_desc, comp_url], output_image)

# Launch the interface
demo.launch()