import gradio as gr
from datetime import datetime, timedelta

from gemini_api import model_api, sentiment, category, ord_num, NO_ORDER, NO_ITEM, food_return, cloth_return, item_identy, item_match, generate_add
from openai_api import model_api as openai_model_api, sentiment as openai_sentiment, image_gen

cust_qry_resp = {"senti":"", "cat":"", "num":""}

#********* UI Code ***********#
with gr.Blocks(title="Customer Support Assistant",
               analytics_enabled=False) as app:
  # States for triggering events
  # Order Num, Sentiment, Category, User Input
  state_order_num = gr.State([NO_ORDER, None, None, None, None])
  # Category, User Input
  state_ret_pol_cat = gr.State([None, None])
  # Item, Category
  state_match_item = gr.State([None, None])
  # Item, theme, age group
  state_gen_add = gr.State([None, None, None])
  
  gr.Markdown("# Customer Support Assistant")
  llm_api = gr.Radio(["gemini-1.0-pro", "gpt-3.5-turbo"], label="Choose LLM", value="gemini-1.0-pro")
  # Inputs from user
  with gr.Row():
    cust_qry = gr.Textbox(lines=5, type="text", label="Customer Query")

  btn_cust_qry = gr.Button("Analyze Query")

  # Model Output
  @gr.render(inputs=[cust_qry, llm_api], triggers=[btn_cust_qry.click])
  # Function for prediction
  def invoke_model(user_input, llm_api):
    if len(user_input) == 0:
      gr.Markdown("## No Customer Query Provided")
    else:  
      if llm_api == "gemini-1.0-pro":
        senti = model_api(user_input, sentiment)
      else:
        senti = openai_model_api(user_input, openai_sentiment)
      cat = model_api(user_input, category)
      num = model_api(user_input, ord_num)
      item = model_api(user_input, item_identy)
      
      def btn_ord_det_clk(ord_num):
        return [ord_num, senti, cat, user_input, item]
      
      # Output response
      gr.Textbox(lines=1, type="text", label="Customer Sentiment", value=senti)
      gr.Textbox(lines=1, type="text", label="Order Category", value=cat)
      ord_num_txt = gr.Textbox(lines=1, type="text", label="Order Number", value=num, interactive=True)
      gr.Textbox(lines=1, type="text", label="Order Item", value=item)
      
      # Decision Rules
      # Order Details
      if senti in ["NEGATIVE", "MIXED"]:
        if num != NO_ORDER:
          btn_ord_det = gr.Button("Fetch Order Details")
          btn_ord_det.click( btn_ord_det_clk, ord_num_txt, state_order_num)
        else:
          with gr.Row():
            gr.Textbox(lines=1, type="text", label="Next Step", value="Ask Order Number")
      else:
        btn_item_match = gr.Button("Fetch Matching Items from Adds Store")
        btn_item_match.click(lambda x: [item, cat], state_match_item, state_match_item)  
          
  @gr.render(inputs=state_order_num)
  def fetch_order_det(ip_arr):
    print("Get order Details")
    ord_num = ip_arr[0]
    ord_senti = ip_arr[1]
    ord_cat = ip_arr[2]
    usr_ip = ip_arr[3]
    ord_item = ip_arr[4]
    if ord_num != NO_ORDER:
      pur_dt = datetime.now() + timedelta(days=-2)
      ord_det = f"Order Number: {ord_num}\nPurchase Date: {pur_dt}\nItem Ordered: {ord_item}"
      gr.Textbox(lines=1, type="text", label="Order Details", value=ord_det)
      if ord_senti in ["NEGATIVE", "MIXED"] and ord_num != NO_ORDER:
        with gr.Row():
          btn_ret_pol = gr.Button("Fetch Return / Replacement Response")
          btn_ret_pol.click(lambda x: [ord_cat, usr_ip], state_ret_pol_cat, state_ret_pol_cat)
  
  @gr.render(inputs=state_ret_pol_cat)
  def fetch_ret_response(ip_arr):
    print("Return Response")
    ord_cat = ip_arr[0]
    user_input = ip_arr[1]
    if ord_cat is not None:
      if ord_cat == "Food":
        resp, policy = food_return(user_input)
      if ord_cat == "Clothing":
        resp, policy = cloth_return(user_input)        
      with gr.Row():
        with gr.Column(scale=1):
          gr.Textbox(lines=1, type="text", label="Return Response", value=resp)
        with gr.Column(scale=1):
          gr.Textbox(lines=1, type="text", label="Return Policy Used", value=policy)
  
  @gr.render(inputs=state_match_item)
  def fetch_match_item(ip_arr):
    print("Return Matching Items")
    ord_item = ip_arr[0]
    ord_cat = ip_arr[1]
    if ord_item is not None:
      ret_item, ret_theme = item_match(ord_item)
      def btn_clk(rd_ip):
        return [ret_item, ret_theme, rd_ip]
      
      with gr.Row():
        gr.Textbox(lines=1, type="text", label="Nearest Matching Item from Adds Store", value=ret_item)
        gr.Textbox(lines=1, type="text", label="Add Theme", value=ret_theme, interactive=True)
      with gr.Row():
        age_cat = gr.Radio(["less than 13 years", "14 to 30 years", "30 to 50 years", "more than 50 years"],
                           label="Select target age group", value="30 to 50 years", interactive=True)
        btn_gen_add = gr.Button("Generate Add")
        btn_gen_add.click(btn_clk, age_cat, state_gen_add)
  
  @gr.render(inputs=state_gen_add)
  def gen_add(ip_arr):
    print("Generate Add")
    ord_item = ip_arr[0]
    theme = ip_arr[1]
    age_grp = ip_arr[2]
    if ord_item is not None:
      ret_add = generate_add(ord_item, theme, age_grp)
      image_url = image_gen(f"{ord_item} {theme} {age_grp}")
      with gr.Row():
        gr.Textbox(type="text", label="Genarated Add", value=ret_add, interactive=True)
      with gr.Row():
        gr.Image(height=1024, width=1024, value=image_url)
  
if __name__ == "__main__":
  app.launch(server_name="0.0.0.0")