import random


import argilla as rg

from chat_interface_preference import ChatInterface


def random_response(message, history, request):
    response = random.choice(["Yes", "No"])
    for char in response:
        yield char


style = "<style>.user-message,.system-message{display:flex;margin:10px}.user-message .message-content{background-color:#c2e3f7;color:#000}.system-message .message-content{background-color:#f5f5f5;color:#000}.message-content{padding:10px;border-radius:10px;max-width:70%;word-wrap:break-word}.container{display:flex;justify-content:space-between}.column{width:48%}</style>"
client = rg.Argilla(api_url="https://davidberenstein1957-argilla-gradio.hf.space", api_key="owner.apikey")

required_settings = rg.Settings(
    fields=[rg.TextField(name="conversation")],
    questions=[
        rg.TextQuestion(name="chosen"),
        rg.TextQuestion(name="rejected"),
    ],
)
name = "test"
if client.datasets(name=name).exists():
    dataset: rg.Dataset = client.datasets(name=name)

else:
    dataset = rg.Dataset(name=name, settings=required_settings)
    dataset.create()

demo = ChatInterface(random_response, cache_examples=False, css=style, rg_dataset=dataset)

if __name__ == "__main__":
    demo.launch()