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",
]
def get_webpage_text(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: {company_html}"""
    prompt_in=[
        {'role':'system','content':system_prompt},
        {'role':'user','content':prompt},
    ]
    generate_kwargs = dict(
        temperature=0.99,
        max_new_tokens=512, #total tokens - input tokens
        top_p=0.99,
        repetition_penalty=1.0,
        do_sample=True,
        seed=random.randint(1,100000000000000),
    )
    stream = client.text_generation(prompt_in, **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=""):
    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("black-forest-labs/FLUX.1-dev")
    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)")
            generate_button = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    generate_button.click(generate_image, inputs=[comp_name, comp_desc, comp_url], outputs=output_image)

# Launch the interface
demo.launch()