import os
from dotenv import load_dotenv
load_dotenv()

import chainlit as cl
from phi.assistant import Assistant
from phi.tools.duckduckgo import DuckDuckGo
from phi.llm.openai import OpenAIChat
from phi.tools.yfinance import YFinanceTools
from src.tools.coin_data_toolkit import CryptoDataTools


# assistant = Assistant(tools=[DuckDuckGo()], show_tool_calls=True)
# assistant.print_response("Whats happening in France?", markdown=True)

# f_assistant = Assistant(
#     llm=OpenAIChat(model="gpt-4o"),
#     tools=[YFinanceTools(stock_price=True, analyst_recommendations=True, company_info=True, company_news=True)],
#     show_tool_calls=True,
#     markdown=True,
# )
# f_assistant.print_response("What is the stock price of NVDA")
# f_assistant.print_response("Write a comparison between NVDA and AMD, use all tools available.")

@cl.on_chat_start
def start():
    is_dev_mode = True if os.getenv("DEV_MODE") else False

    # Initialize the assistant
    cxbt_assistant = Assistant(
        llm=OpenAIChat(model="gpt-4o"),
        tools=[CryptoDataTools(), DuckDuckGo(), YFinanceTools(stock_price=True)],
        show_tool_calls= is_dev_mode,
        markdown=True,
    )

    # Set the assistant in the user session
    cl.user_session.set("agent", cxbt_assistant)

@cl.on_message
async def main(message: cl.Message):
    # Retrieve the assistant from the user session
    agent = cl.user_session.get("agent")

    # Process the user message using the assistant
    # response = agent.run(message.content, stream=True)
    response = ""
    for delta in agent.run(message.content, stream=True):
        response += delta
    
    # Send the response back to the user
    await cl.Message(content=response).send()

# Run the Chainlit application
if __name__ == "__main__":
    cl.run()