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

STYLE = """
#container-col {
  
}

.custom-btn {
    border: none !important;
    background: none !important;
    box-shadow: none !important;
    display: block !important;
    text-align: left !important;
}
.custom-btn:hover {
    background: rgb(243 244 246) !important;
}
#prompt-txt > label > span {
    display: none !important;
}
#prompt-txt > label > textarea {
    border: transparent;
    box-shadow: none;
}
#chatbot {
    height: 800px; 
    overflow: auto;
    box-shadow: none !important;
    border: none !important;
}
#chatbot > .wrap {
    max-height: 780px;
}
#left-pane {
    background-color: #f9fafb;
    border-radius: 15px;
    padding: 10px;
}

#left-top {
    padding-left: 10px;
    padding-right: 10px;
    text-align: center;
    font-weight: bold;
    font-size: large;    
}

#chat-history-accordion {
    background: transparent;
    border: 0.8px !important;  
}
"""

get_local_storage = """
function() {
  globalThis.setStorage = (key, value)=>{
    localStorage.setItem(key, JSON.stringify(value));
  }
  globalThis.getStorage = (key, value)=>{
    return JSON.parse(localStorage.getItem(key));
  }
  var local_data = getStorage('local_data');
  const history = [];
  if(local_data) {
    console.log(local_data);
    local_data[0].pingpongs.forEach(element => 
      history.push([element.ping, element.pong])
    );
  }
  else {
    local_data = [
      {'ctx': '', 'pingpongs':[]},
      {'ctx': '', 'pingpongs':[]},
      {'ctx': '', 'pingpongs':[]},
      {'ctx': '', 'pingpongs':[]},
      {'ctx': '', 'pingpongs':[]},
    ];
    local_data[0].pingpongs.forEach(element => 
      history.push([element.ping, element.pong])
    );    
  }
  console.log(history);
  return [history, 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

with gr.Blocks(css=STYLE, elem_id='container-col') as block:
  idx = gr.State(0)
  local_data = gr.JSON({},visible=False)
  
  with gr.Row():
    with gr.Column(scale=1, min_width=180):
      gr.Markdown("GradioChat", elem_id="left-top")
        
      with gr.Column(elem_id="left-pane"):
        with gr.Accordion("Histories", elem_id="chat-history-accordion"):
          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=8):
      chatbot = gr.Chatbot(elem_id='chatbot')
      instruction_txtbox = gr.Textbox(
          placeholder="Ask anything", label="",
          elem_id="prompt-txt"
      )

  btns = [first, second, third, fourth, fifth]
  for btn in btns:
    btn.click(
      set_chatbot,
      [btn, local_data],
      [chatbot, idx]
    )
  
  instruction_txtbox.submit(
    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=[chatbot, local_data],
      _js=get_local_storage,
  )

block.launch(debug=True)