import gradio as gr
import os

# ================= API KEYS =================
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
HF_TOKEN = os.getenv("HF_TOKEN")

# ================= SYSTEM PROMPT =================
SYSTEM_PROMPT = """
You are Rizer AI, an expert personal finance coach for Indian users.

STRICT RULES:
- Only answer in finance context
- SIP means Systematic Investment Plan (NOT VoIP)
- Use Indian context (₹, SEBI, RBI, mutual funds)
- Keep answers simple and structured

FORMAT:
1. Definition
2. Key Points
3. Example (India)
4. Tip (practical advice)

Avoid technical jargon. Be clear and helpful.
"""

# ================= PREPROCESS =================
def preprocess(question):
    q = question.lower()

    if "sip" in q:
        return question + " (in finance, systematic investment plan)"
    
    return question


# ================= GROQ =================
def ask_groq(question):
    from groq import Groq
    client = Groq(api_key=GROQ_API_KEY)

    chat = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question}
        ],
        temperature=0.4,   # more accurate
        top_p=0.9,
    )
    return chat.choices[0].message.content


# ================= GEMINI =================
def ask_gemini(question):
    import google.generativeai as genai
    genai.configure(api_key=GEMINI_API_KEY)

    model = genai.GenerativeModel("gemini-pro")

    response = model.generate_content(
        SYSTEM_PROMPT + "\nUser: " + question
    )
    return response.text


# ================= HF FALLBACK =================
from transformers import pipeline
from huggingface_hub import login

if HF_TOKEN:
    login(token=HF_TOKEN)

hf_model = pipeline(
    "text-generation",
    model="rohith2006345/rizer-finance-ai",
    token=HF_TOKEN
)

def ask_hf(question):
    prompt = f"""
You are a finance expert.

IMPORTANT:
- SIP = Systematic Investment Plan
- Answer only in finance context

Question: {question}
Answer:
"""
    result = hf_model(
        prompt,
        max_new_tokens=200,
        temperature=0.5,
        top_p=0.9,
        repetition_penalty=1.2,
        do_sample=True,
        pad_token_id=hf_model.tokenizer.eos_token_id
    )
    return result[0]["generated_text"].split("Answer:")[-1].strip()


# ================= KEYWORD FILTER =================
FINANCE_KEYWORDS = [
    "money","invest","sip","mutual fund","stock","nse","bse","tax",
    "budget","save","fd","upi","bank","credit","loan","insurance",
    "nifty","sensex","sebi","rbi","rupee","finance","salary"
]

def is_finance_question(q):
    return any(w in q.lower() for w in FINANCE_KEYWORDS)


# ================= MAIN FUNCTION =================
def ask_rizer(question):
    if not question.strip():
        return "Please ask a question 😊"

    if not is_finance_question(question):
        return "I'm Rizer AI 💰 — ask me about investing, saving, SIP, stocks, etc."

    question = preprocess(question)

    # 1️⃣ GROQ (BEST)
    if GROQ_API_KEY:
        try:
            return ask_groq(question)
        except Exception as e:
            print("Groq failed:", e)

    # 2️⃣ GEMINI
    if GEMINI_API_KEY:
        try:
            return ask_gemini(question)
        except Exception as e:
            print("Gemini failed:", e)

    # 3️⃣ HF fallback
    return ask_hf(question)


# ================= UI =================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Rizer AI 💰")
    gr.Markdown("### Your smart personal finance coach for Indian youth 🇮🇳")

    with gr.Row():
        question = gr.Textbox(
            placeholder="Ask about SIP, stocks, saving, tax...",
            label="Your Question"
        )

    with gr.Row():
        submit = gr.Button("Submit 🚀")
        clear = gr.Button("Clear ❌")

    answer = gr.Textbox(label="Rizer AI Answer", lines=10)

    submit.click(ask_rizer, inputs=question, outputs=answer)
    clear.click(lambda: ("", ""), None, [question, answer])


demo.launch()