import os
import gradio as gr
from http import HTTPStatus
import openai
from typing import Generator, List, Optional, Tuple, Dict
from urllib.error import HTTPError
default_system = 'You are a helpful assistant.'

API_URL = os.getenv('API_URL')
API_KEY = os.getenv('API_KEY')
oai_client = openai.OpenAI(api_key=API_KEY, base_url=API_URL)

History = List[Tuple[str, str]]
Messages = List[Dict[str, str]]

def clear_session() -> History:
    return '', []

def modify_system_session(system: str) -> str:
    if system is None or len(system) == 0:
        system = default_system
    return system, system, []

def history_to_messages(history: History, system: str) -> Messages:
    messages = []
    for h in history:
        messages.append({'role': 'user', 'content': h[0]})
        messages.append({'role': 'assistant', 'content': h[1]})
    return messages

def messages_to_history(messages: Messages) -> Tuple[str, History]:
    system = messages[0]['content']
    history = []
    for q, r in zip(messages[1::2], messages[2::2]):
        history.append([q['content'], r['content']])
    return system, history

def model_chat(query: Optional[str], history: Optional[History]) -> Generator[Tuple[str, str, History], None, None]:
    if query is None:
        query = ''
    if history is None:
        history = []
    messages = history_to_messages(history, system)
    messages.append({'role': 'user', 'content': query})
    gen = oai_client.chat.completions.create(
        model='dicta-il/dictalm2.0-instruct',
        messages=messages,
        temperature=0.7,
        max_tokens=512,
        top_p=0.9,
        stream=True
    )
    messages.append({'role': 'assistant', 'content': ''})
    for completion in gen:
        text = completion.choices[0].text
        messages[-1]['content'] += text
        system, history = messages_to_history(messages)
        yield '', history, system
    
with gr.Blocks() as demo:
    gr.Markdown("""<center><font size=8>DictaLM2.0-Instruct Chat Demo</center>""")
    
    chatbot = gr.Chatbot(label='dicta-il/dictalm2.0-instruct')
    textbox = gr.Textbox(lines=2, label='Input')

    with gr.Row():
        clear_history = gr.Button("🧹 Clear history")
        sumbit = gr.Button("🚀 Send")

    sumbit.click(model_chat,
                 inputs=[textbox, chatbot],
                 outputs=[textbox, chatbot])
    clear_history.click(fn=clear_session,
                        inputs=[],
                        outputs=[textbox, chatbot])

demo.queue(api_open=False).launch(max_threads=10,height=800, share=False)