import chainlit as cl
# from ollama import chat
from source.aisearch_v1 import aiSearch
from source.convoFlow import classify_query,reformulate_query
import json



@cl.on_chat_start
def on_chat_start():
    cl.user_session.set("counter", 0)
    

@cl.step(type="tool",name="Deep Search", show_input=False)
async def deepsearch(msg):

    print(f"Message from user: {msg}")
    
    inst = aiSearch()
    stream = await inst.run_search(msg)
    return stream


@cl.on_message  # this function will be called every time a user inputs a message in the UI
async def main(message: cl.Message):
    """
    This function is called every time a user inputs a message in the UI.
    It sends back an intermediate response from the tool, followed by the final answer.

    Args:
        message: The user's message.

    Returns:
        None.
    """
    counter = cl.user_session.get("counter")
    msg = cl.Message(content="")
    if counter == 0:
        query = message.content
        cl.user_session.set("curnt_query", query)
        resp = classify_query(query)
        resp = json.loads(resp)
        calrity = resp['clarity']
        if calrity:
            streams = await deepsearch(query)
            async for update in streams:
                if update.choices:
                    await msg.stream_token(update.choices[0].delta.content or "")
            cl.user_session.set("counter", 0)
            await msg.update()
        else:
            quetn = "\n".join(resp['follow_up_questions'])
            response = "Please answer the following:\n\n" + quetn
            cl.user_session.set("clrf_queries", quetn)
            await cl.Message(content=response).send()
            counter += 1
            cl.user_session.set("counter", counter)
    elif counter == 1:
        or_query = cl.user_session.get("curnt_query")
        clr_queries = cl.user_session.get("clrf_queries")
        cur_resp = message.content
        resp = reformulate_query(or_query,cur_resp,clr_queries)
        streams = await deepsearch(or_query+""+resp)
        async for update in streams:
            if update.choices:
                await msg.stream_token(update.choices[0].delta.content or "")
        cl.user_session.set("counter", 0)
        await msg.update()