import gradio as gr
import cohere 
import os

co = cohere.Client(
  api_key=os.getenv('api_key'), # This is your trial API key
) 

def cohere_response(msg):
  stream = co.chat_stream( 
  model='command-r-plus',
  message=msg,
  temperature=0.3,
  chat_history=[],
  prompt_truncation='AUTO',
) 
  output = ''
  for event in stream:
     if event.event_type == "text-generation":
        output += event.text
        print(event.text, end='')
  return output

def generate_mockup(problem):
  prompt_sum = f'''You are a UX designer who is an expert in generating front-end code for a given requirement from a product manager: {problem}. You need to carefully understand the requirement
  and generate the html code for the mockup. You can use bootstrap and popularly available visual libraries to generate the code. Do not generate placeholders and always make assumptions that makes the generated mockup useful to solve real-world problems. It should be visually appealing and should
  satisfy the requirement as a mockup. You are supposed to just generate the html code and nothing else as output.'''
  response = cohere_response(prompt_sum)
  return response

# Define a function to handle form submissions (this is just a placeholder)
def handle_form(requirements):
    print(requirements)
    emp_content = generate_mockup(requirements)
    #final_html_content = extract_html(emp_content)
    return emp_content, emp_content

# Create a Gradio interface to showcase the HTML content
gr.Interface(
    fn=handle_form,
    inputs=gr.Textbox(),
    outputs=[gr.HTML(),gr.Textbox()],
    title="Mockup Generator",
    description="Generate a mockup for a given requirement",
    examples=["Design an interface for a food delivery app payment checkout screen, a menu with categories like 'Asian Cuisine' and 'Italian Pasta', and a user profile section.","Generate a newsletter template for a tech company with sections for latest news, featured products, and upcoming events."],
).launch()