import os
import gradio as gr
from groq import Groq

# Initialize the Groq client
# On Hugging Face, GROQ_API_KEY should be set up as a Repository Secret
api_key = os.environ.get("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")

# Define the banking agent's system prompt to enforce boundaries and persona
BANKING_SYSTEM_PROMPT = """
You are a helpful, secure, and professional AI Banking Customer Support Agent. 
Your goal is to assist customers efficiently with the following queries:
- Account balance inquiry (remind them to never share pins/passwords)
- Branch timings (Standard hours: Mon-Fri 9:00 AM - 4:00 PM, Sat 9:00 AM - 1:00 PM)
- ATM locations and branch information
- Credit card bill due date queries
- Loan product information (Home, Personal, Auto loans)
- Complaint registration (Collect name, contact info, and details professionally)

Important Security Notice: Never ask for, accept, or display sensitive personal information such as full account numbers, passwords, PINs, or CVVs. If a user shares this, politely remind them to keep it confidential. Stay professional and polite at all times.
"""

def respond(message, chat_history, model_choice):
    if not api_key:
        return chat_history + [["", "⚠️ Groq API Key is missing. Please configure 'GROQ_API_KEY' in your Space's Repository Secrets."]]
    
    try:
        client = Groq(api_key=api_key)
        
        # Format history for the Groq chat completion endpoint
        messages = [{"role": "system", "content": BANKING_SYSTEM_PROMPT}]
        
        for user_msg, assistant_msg in chat_history:
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})
                
        # Append the latest user query
        messages.append({"role": "user", "content": message})
        
        # Call the Groq API
        completion = client.chat.completions.create(
            model=model_choice,
            messages=messages,
            temperature=0.5,
            max_tokens=1024,
        )
        
        bot_response = completion.choices[0].message.content
        chat_history.append((message, bot_response))
        return chat_history, ""
        
    except Exception as e:
        chat_history.append((message, f"❌ Error communicating with Groq API: {str(e)}"))
        return chat_history, ""

# Build Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🏦 AI Banking Customer Support Agent")
    gr.Markdown("Welcome! This AI agent is here to help you find answers regarding account inquiries, branch hours, ATM setups, loan products, or registering complaints.")
    
    with gr.Row():
        model_choice = gr.Dropdown(
            choices=["llama-3.1-8b-instant", "llama-3.3-70b-versatile"],
            value="llama-3.3-70b-versatile",
            label="Select LLM Architecture Models",
            interactive=True
        )
        
    chatbot = gr.Chatbot(label="Banking Support Chat")
    msg = gr.Textbox(label="Type your banking query here...", placeholder="e.g., What are the standard branch timings?")
    
    with gr.Row():
        submit_btn = gr.Button("Send", variant="primary")
        clear_btn = gr.Button("Clear Chat")

    # Wire up actions
    msg.submit(respond, inputs=[msg, chatbot, model_choice], outputs=[chatbot, msg])
    submit_btn.click(respond, inputs=[msg, chatbot, model_choice], outputs=[chatbot, msg])
    clear_btn.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    # Combined configuration for themes and sharing options in Gradio 6+
    demo.launch(share=True)