import os
import gradio as gr
import spaces
from openai import OpenAI

# Read the key once from the Space secret at startup
NV_API_KEY = os.environ.get("NVIDIA_API_KEY", "")

@spaces.GPU(duration=30)
def categorize_expense(item_input):
    """Calls the Nvidia API to categorize the user's item."""
    if not item_input.strip():
        return "Please enter an item to categorize."

    if not NV_API_KEY:
        return "⚠️ Error: NVIDIA_API_KEY secret is not set on this Space. Add it in Settings → Variables and secrets."

    try:
        # Setup client with Nvidia's base URL
        client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key=NV_API_KEY
        )

        system_prompt = (
            "You are an expert expense tracking assistant. Categorize the provided item into exactly "
            "one of the following categories:\n"
            "- Commodity (Groceries, food, chai, snacks, etc.)\n"
            "- Stationary (Books, pens, modules, copies, etc.)\n"
            "- Personal Care & Hygiene\n"
            "- Tech & Digital Services (Mobile recharge, internet, hosting, domains, etc.)\n"
            "- Utilities & Living (Rent, electricity, water, fuel, transit, etc.)\n"
            "- Health & Miscellaneous (Medicines, clothes, entertainment, medical, etc.)\n\n"
            "Respond with ONLY the exact name of the category. Do not include punctuation, explanations, or introductory text."
        )

        completion = client.chat.completions.create(
            model="meta/llama-3.1-70b-instruct",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Item: {item_input}"}
            ],
            temperature=0.1,
            max_tokens=20
        )

        category = completion.choices[0].message.content.strip()
        return f"**Item:** {item_input.title()}\n\n**Category:** {category}"

    except Exception as e:
        return f"❌ An error occurred: {str(e)}"

# Define the Gradio interface
with gr.Blocks(title="AI-Powered Expense Categorizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 AI-Powered Expense Categorizer")
    gr.Markdown("Type any product or expense to dynamically categorize it using Llama 3.1 via Nvidia's API.")

    with gr.Row():
        # Input fields
        with gr.Column():
            item_input = gr.Textbox(
                label="Enter an item",
                placeholder="e.g., recharge, medicine, rent, clothes..."
            )
            submit_btn = gr.Button("Categorize", variant="primary")

        # Output display
        with gr.Column():
            output_display = gr.Markdown(label="Result")

    # Trigger the function on button click or when pressing 'Enter' in the text box
    submit_btn.click(
        fn=categorize_expense,
        inputs=[item_input],
        outputs=output_display
    )
    item_input.submit(
        fn=categorize_expense,
        inputs=[item_input],
        outputs=output_display
    )

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