import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

example_questions_en = [
    "What is a TFSA?",
    "What are good long-term stocks?",
    "How does compound interest work?",
    "What is the difference between ETFs and stocks?",
    "How do I start investing with a small budget?",
    "What are dividend stocks and how do they work?",
    "How do I diversify my portfolio?",
    "What are the risks of investing in stocks?"
]

example_questions_fr = [
    "Qu'est-ce qu'un CELI?",
    "Quelles sont de bonnes actions à long terme?",
    "Comment fonctionne l'intérêt composé?",
    "Quelle est la différence entre les FNB et les actions?",
    "Comment commencer à investir avec un petit budget?",
    "Qu'est-ce qu'une action à dividendes et comment fonctionne-t-elle?",
    "Comment diversifier mon portefeuille?",
    "Quels sont les risques d'investir en bourse?"
]

def respond(message, history, language):
    if history is None:
        history = []
    
    if language == "English":
        system_prompt = (
            "You are a financial assistant specializing in investment strategies. "
            "Provide concise and direct answers to each question. "
            "Keep answers short and simple, offering general information, not specific advice."
        )
    else:
        system_prompt = (
            "Vous êtes un assistant financier spécialisé dans les stratégies d'investissement. "
            "Fournissez des réponses concises et directes à chaque question. "
            "Gardez les réponses courtes et simples, en offrant des informations générales, pas des conseils spécifiques."
        )
    
    messages = [{"role": "system", "content": system_prompt}]
    
    for user_msg, bot_reply in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_reply:
            messages.append({"role": "assistant", "content": bot_reply})
    
    messages.append({"role": "user", "content": message})
    
    reply = client.chat_completion(
        messages=messages,
        max_tokens=256,
        temperature=0.7,
        top_p=0.95,
    )
    
    response = reply.choices[0].message.content
    return response

def submit_message(message, history, language):
    if not message.strip():
        return history, ""
    bot_response = respond(message, history, language)
    history.append((message, bot_response))
    return history, ""

def update_suggested_questions(language):
    questions = example_questions_en if language == "English" else example_questions_fr
    return gr.update(choices=questions)

def reset_chat():
    return [], "", []

with gr.Blocks() as demo:
    state = gr.State([])
    
    with gr.Row():
        chatbot = gr.Chatbot(height=450)
    
    language_toggle = gr.Radio(
        ["English", "Français"],
        label="Select Language / Sélectionner la langue",
        value="English"
    )
    
    suggested_questions = gr.Dropdown(choices=example_questions_en, label="Suggested Questions")
    user_input = gr.Textbox(label="Your Message", elem_id="user_input")
    send_button = gr.Button("Send", elem_id="send_button")
    new_chat_button = gr.Button("New Chat", elem_id="new_chat_button")

    send_button.click(
        fn=submit_message,
        inputs=[user_input, state, language_toggle],
        outputs=[chatbot, user_input],
        queue=False
    )

    language_toggle.change(
        fn=update_suggested_questions,
        inputs=language_toggle,
        outputs=suggested_questions
    )

    suggested_questions.change(
        fn=lambda q: ([], q),
        inputs=suggested_questions,
        outputs=[state, user_input]
    )

    new_chat_button.click(
        fn=reset_chat,
        inputs=[],
        outputs=[chatbot, user_input, state]
    )

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