import os

import solara
from solara_enterprise import auth
from reacton import ipyvuetify as rv
from functools import partial

os.environ['SOLARA_OAUTH_API_BASE_URL']="dev-r2bw2n3hwf72bni6.us.auth0.com"  # replace with your domain
os.environ['SOLARA_OAUTH_CLIENT_ID']="lQl2yjR7dr7XVbK480OJ2sfgoWf6U7s2"  # replace with your client ID
os.environ['SOLARA_OAUTH_CLIENT_SECRET']="BsGtc_GLEZgsv4DsZnSya4YhB38nm7ohJp5kobfOVb2X8CeT9XTHMOxnXR_djQMd"  # not shown here, replace with your client secret

os.environ['REPLICATE_API_TOKEN']='r8_YSD53gSw5IQbT8Fmnekc8ypEIuy28wN4IG45C'

import replicate

@solara.component
def Message(message):
    icon = "mdi-robot" if message['user']=='Assistant' else 'mdi-account'
    icon = rv.Icon(children=[icon])
    user = rv.CardTitle(icon=icon,children=[solara.Columns([0.6,2,10],children=[icon,message['user'],rv.Spacer()])])
    text = rv.CardText(children=[message['text']])
    messageDiv = rv.Card(_class="mx-auto",width="100%",flat=True,children=[user,text])
    return messageDiv

def Redo():
    message = messages.value[-2]['text']
    messages.set(messages.value[:-2])
    sendMessage(message)

@solara.component
def InputText():
    InputText = solara.InputText("Enter some text",value=input_text)
    rv.use_event(InputText,"keypress",partial(handleEnter))
    InputText

def sendMessage(message):
    name = str(auth.user.value['userinfo']['name']).split(" ")[0]
    messages.set(messages.value+[{'user':name,'text':message}])
    prompt.value += f"[INST] {message} [/INST] \n"
    messages.set(messages.value+[{'user':'Assistant','text':''}])
    response = ''
    for token in replicate.stream(
        "meta/llama-2-7b-chat",
        input={
            "debug": False,
            "top_p": 1,
            "prompt": prompt.value,
            "temperature": 0.75,
            "system_prompt": f"You are a helpful, respectful and honest assistant. You are tasked with helping {name} with figuring out the correct European Binding Tariff Information for their shipment. Please address them by their name and their needs.",
            "max_new_tokens": 800,
            "min_new_tokens": -1,
            "prompt_template": "[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]",
            "repetition_penalty": 1
        },
    ):
        response += str(token)
        messages.set(messages.value[:-1]+[{'user':'Assistant', 'text':response}])
    prompt.value += f"{messages.value[-1]['text']}"
    print('Boom')
        

def handleEnter(*args):
    if str(args[2]['key'])=='Enter':
        message = args[0].v_model
        print(message)
        sendMessage(message)


prompt = solara.reactive('')
messages = solara.reactive([])
input_text = solara.reactive('')

@solara.component
def ColumnLayout(children):
    solara.Columns([2,6,2],children=[solara.Column(),solara.Column(children),solara.Column()])
    

@solara.component
def Page():
    with solara.AppLayout(title="EBTI GPT"):
        solara.AppBarTitle(children=[])
        with rv.Card(flat=True):
            with ColumnLayout():
                if not auth.user.value:
                    solara.Markdown(f"### Please Login ")
                    solara.Button("Login", icon_name="mdi-login", href=auth.get_login_url())
                else:
                    userinfo = auth.user.value['userinfo']
                    if 'name' in userinfo:
                        solara.Markdown(f"### Welcome {userinfo['name']}")
            with ColumnLayout():
                rv.Card(flat=True,children=[Message(message) for message in messages.value])
            with ColumnLayout():
                InputText()
            with ColumnLayout():
                solara.Button("Logout", icon_name="mdi-logout", href=auth.get_logout_url())
Page()