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

STYLE = """
.custom-btn {
    border: none;
    background: none;
    box-shadow: none;
}
"""

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');
  const history = [];

  console.log(1);

  if(local_data) {
    console.log(2);
    console.log(local_data);
    local_data[0].pingpongs.forEach(element => 
      history.push([element.ping, element.pong])
    );
  }

  console.log(history);
  return [local_data, history];
}
"""

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", elem_classes=["custom-btn"])
      second = gr.Button("2nd", elem_classes=["custom-btn"])
      third = gr.Button("3rd", elem_classes=["custom-btn"])
      fourth = gr.Button("4th", elem_classes=["custom-btn"])
      fifth = gr.Button("5th", elem_classes=["custom-btn"])
    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, chatbot],
      _js=get_local_storage,
  )

block.queue().launch(debug=True)