import re
import gradio as gr
import vanna as vn
from groq import Groq
from vanna.remote import VannaDefault

# Set up Vanna.ai and Groq client as before
MY_VANNA_MODEL = "llama3-8b"

vn = VannaDefault(model=MY_VANNA_MODEL, api_key='30efac58cfee46d1967e28b8d6bdf5db')

vn.connect_to_mssql(odbc_conn_str=r'DRIVER={ODBC Driver 17 for SQL Server};SERVER=YISC1100715LT\SQLEXPRESS;DATABASE=master;Trusted_Connection=yes;') # Connect to your database

# Set up Groq client
groq_client = Groq(api_key="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")

def get_order_status(order_number):
    sql = f"SELECT know_history.status FROM know_history WHERE know_history.erpordernumber = {order_number};"
    result = vn.run_sql(sql)
    # if result and len(result) > 0:
        # return result['STATUS']
    return result['status']

def generate_response(user_input):
    # Check for order number in the input
    order_match = re.search(r'#?(\d{5})', user_input)
    
    if order_match:
        order_number = order_match.group(1)
        status = get_order_status(order_number)
        
        # if status:
            # Use Groq to generate a conversational response
        prompt = f"""Given an order status '{status}' for order number {order_number}, 
            generate a friendly, conversational response to the customer's query: "{user_input}".
            The response should be informative and reassuring."""

        chat_completion = groq_client.chat.completions.create(
                messages=[
                    {
                        "role": "system",
                        "content": "You are a helpful customer service chatbot for an e-commerce company."
                    },
                    {
                        "role": "user",
                        "content": prompt,
                    }
                ],
                model="llama3-8b-8192",
                max_tokens=150,
                temperature=0.7,
            )
            
        return chat_completion.choices[0].message.content.strip()
    #     else:
    #         return f"I'm sorry, but I couldn't find any information for order #{order_number}. Could you please check if the order number is correct?"
    else:
        # Handle general queries
        prompt = f"""As a customer service chatbot for an e-commerce company, 
        provide a helpful response to the following customer query: "{user_input}"."""

        chat_completion = groq_client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful customer service chatbot for an e-commerce company."
                },
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="llama3-8b-8192",
            max_tokens=150,
            temperature=0.7,
        )
        
        return chat_completion.choices[0].message.content.strip()

def chat_interface(message, history):
    response = generate_response(message)
    return response

iface = gr.ChatInterface(
    chat_interface,
    title="E-Commerce Customer Service Chatbot",
    description="Ask about your order status or any other questions!",
    examples=[
        "Where is my order #12345?",
        "What is the status of my order #67890?",
        "How can I track my order?",
        "Can I change my shipping address?",
        "What's your return policy?"
    ]
)

if __name__ == "__main__":
    iface.launch()