import os, json
from typing import Optional, Dict
from dotenv import load_dotenv

load_dotenv()

import chainlit as cl
from phi.assistant import Assistant
from phi.llm.openai import OpenAIChat
from phi.llm.ollama import Ollama
from phi.llm.groq import Groq
from phi.tools.duckduckgo import DuckDuckGo
from phi.tools.yfinance import YFinanceTools
from src.databases.postgres import sqlalchemy_engine
from src.knowledge_bases.combined import knowledge_base
from src.tools.crypto_swap_toolkit import CryptoSwapTools
from src.tools.crypto_bridge_toolkit import CrossChainSwapTools
from src.tools.crypto_data_toolkit import CryptoDataTools
from portkey_ai import PORTKEY_GATEWAY_URL, createHeaders
from src.config.portkey_config import generate_portkey_config
from src.tools.user_profile_toolkit import UserProfileToolkit
from phi.storage.assistant.postgres import PgAssistantStorage
from src.tools.crypto_evm_wallet_toolkit import CryptoEVMWalletTools
from src.tools.user_confirmation_pin_toolkit import UserConfirmationPinToolkit


storage = PgAssistantStorage(
    # stores runs in the ai.assistant_runs table
    table_name="assistant_runs",
    db_engine=sqlalchemy_engine
)

@cl.oauth_callback
def oauth_callback(
  provider_id: str,
  token: str,
  raw_user_data: Dict[str, str],
  default_user: cl.User,
) -> Optional[cl.User]:
    # step-todo: will handle auth from v2 api
    return cl.User(
        identifier=raw_user_data.get('email'),
        email=raw_user_data.get('email'),
        name=raw_user_data.get('name'),
        metadata={**raw_user_data, **({'image':default_user.metadata.get('image')})}
    )


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    # Fetch the user matching username from your database
    # and compare the hashed password with the value stored in the database
    if (username, password) == ("admin", "admin"):
        user = {
            "identifier": "admin", "email": username, "metadata": {"role": "admin", "provider": "credentials"}
        }
        cl.user_session.set("user_profile", user)
        return cl.User(
            identifier="admin", email=username, metadata={"role": "admin", "provider": "credentials"}
        )
    else:
        return None


@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Create a crypto wallet for me",
            message="create a crypto wallet for me",
            icon="/public/wallet-svgrepo-com.svg",
        ),

        cl.Starter(
            label="Latest News on defi, crypto and solana",
            message="What news are currentlly trending on defi.",
            icon="/public/news-svgrepo-com.svg",
            ),
        cl.Starter(
            label="Get price of BTC, ETH and PEPE",
            message="Get me the Price of BTC, ETH and PEPE",
            icon="/public/coins-electronics-svgrepo-com.svg",
            ),
        cl.Starter(
            label="Get trending stocks",
            message="Get latest stock",
            icon="/public/stockchart-svgrepo-com.svg",
            )
        ]

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

    portkey_local_gateway = True if os.getenv("PORTKEY_LOCAL_GATEWAY_URL") else False
    portkey_config = generate_portkey_config(local=portkey_local_gateway)

    # Initialize the assistant
    cxbt_assistant = Assistant(
        # llm=OpenAIChat(),
        # llm=Groq(model="llama3-70b-8192"),
        llm=Ollama(model="llama3"),
        tools=[
            DuckDuckGo(),
            CryptoDataTools(),
            CryptoSwapTools(),
            CrossChainSwapTools(),
            UserProfileToolkit(),
            CryptoEVMWalletTools(),
            UserConfirmationPinToolkit(),
            YFinanceTools(stock_price=True)
        ],
        show_tool_calls=is_dev_mode,
        markdown=True,
        knowledge_base=knowledge_base,
        storage=storage,
        search_knowledge=True,
        read_chat_history=True,
        add_references_to_prompt=True,
        add_chat_history_to_prompt=True,
        prevent_hallucinations=True,
        prevent_prompt_injection=True
    )
    cxbt_assistant.knowledge_base.load(recreate=False)

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

@cl.on_message
async def main(message: cl.Message):

    msg = cl.Message(content="")
    await msg.send()

    # Retrieve the assistant from the user session
    agent = cl.user_session.get("agent")

    # Process the user message using the assistant
    for delta in agent.run(message.content, stream=True):
        for part in delta:
            if token := part or "":
                # Send the response back to the user
                await msg.stream_token(token)

    await msg.update()

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