import json
import gradio as gr
from pingpong import PingPong
from pingpong.gradio import GradioAlpacaChatPPManager

get_local_storage = """
    function() {
      globalThis.setStorage = (key, value)=>{
        localStorage.setItem(key, JSON.stringify(value))
      }
       globalThis.getStorage = (key, value)=>{
        return JSON.parse(localStorage.getItem(key))
      }
       const local_data =  getStorage('local_data')
       return [local_data];
      }
    """

def add_pingpong(idx, ld, ping):
  res = [
      GradioAlpacaChatPPManager.from_json(json.dumps(ppm))
      for ppm in ld
  ]

  ppm = res[idx]
  ppm.add_pingpong(PingPong(ping, "dang!!!!!!!"))
  return ppm.build_uis(), str(res)

def set_chatbot(btn, ld):
  choice = 0
  if btn == "1st":
    choice = 0
  elif btn == "2nd":
    choice = 1
  elif btn == "3rd":
    choice = 2
  elif btn == "4th":
    choice = 3
  elif btn == "5th":
    choice = 4

  res = [
      GradioAlpacaChatPPManager.from_json(json.dumps(ppm_str))
      for ppm_str in ld
  ]
  return res[choice].build_uis(), choice

def initialize(local_data):
  if local_data is None:
    list_conv = [
        GradioAlpacaChatPPManager()
        for _ in range(5)
    ]
    local_data = str(list_conv)

  return local_data

def initialize_t(ld):
  print(ld)

  if ld == {} or id is None:
    list_conv = [
        GradioAlpacaChatPPManager()
        for _ in range(5)
    ]
    local_data = str(list_conv)
  else:
    list_conv = [
        GradioAlpacaChatPPManager.from_json(json.dumps(ppm))
        for ppm in ld
    ]
    local_data = str(list_conv)

  return list_conv[0].build_uis(), local_data

with gr.Blocks() as block:
  idx = gr.State(0)
  ttt = gr.Textbox(visible=False)
  local_data = gr.JSON(
      {},
      label="Local Storage",
      visible=False
  )
  
  with gr.Row():
    with gr.Column(scale=1, min_width=100):
      first = gr.Button("1st")
      second = gr.Button("2nd")
      third = gr.Button("3rd")
      fourth = gr.Button("4th")
      fifth = gr.Button("5th")
    with gr.Column(scale=4):
      chatbot = gr.Chatbot(elem_id='chatbot')
      instruction_txtbox = gr.Textbox(placeholder="What do you want to say to AI?", label="Instruction")

  first.click(
    initialize,
    local_data,
    local_data
  ).then(
    set_chatbot,
    [first, local_data],
    [chatbot, idx]
  )

  second.click(
    initialize,
    local_data,
    local_data
  ).then(
    set_chatbot,
    [second, local_data],
    [chatbot, idx]
  )

  third.click(
    initialize,
    local_data,
    local_data
  ).then(
    set_chatbot,
    [third, local_data],
    [chatbot, idx]
  )
  
  fourth.click(
    initialize,
    local_data,
    local_data
  ).then(
    set_chatbot,
    [fourth, local_data],
    [chatbot, idx]
  )   

  fifth.click(
    initialize,
    local_data,
    local_data
  ).then(
    set_chatbot,
    [fifth, local_data],
    [chatbot, idx]
  )     

  instruction_txtbox.submit(
    initialize,
    local_data,
    local_data
  ).then(        
    add_pingpong,
    [idx, local_data, instruction_txtbox],
    [chatbot, local_data]
  ).then(
    None, local_data, None, 
    _js="(v)=>{ setStorage('local_data',v) }"
  )

  block.load(
      None,
      inputs=None,
      outputs=[local_data],
      _js=get_local_storage,
  )

block.launch(debug=True)