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 a series of mockups that together form a clickable prototype. 
  You can use Bootstrap and popularly available visual libraries to generate the code.Each screen should be clearly delimited by the header - Screen 1, Screen 2, Screen 3, etc. Do not generate placeholders and always make assumptions that make the generated 
  mockups useful to solve real-world problems. Each mockup should link to the next, creating a seamless user experience. The prototype should be visually appealing and 
  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()