import gradio as gr
import pandas as pd
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.schema import StrOutputParser
from txtai.pipeline import LLM as TxtaiLLM
from langchain.llms.base import LLM as LangChainLLM

X_test_df = pd.read_csv('Financial_Planner_test.csv')

# ✅ Setup txtai + LangChain wrapper
txtai_llm_instance = TxtaiLLM("MaziyarPanahi/gemma-2-2b-it-GGUF/gemma-2-2b-it.Q8_0.gguf")

class TxtaiLangChainLLM(LangChainLLM):
    def _call(self, prompt, stop=None):
        return txtai_llm_instance(prompt)

    @property
    def _llm_type(self):
        return "txtai"

llm = TxtaiLangChainLLM()

# ✅ Rule-based Financial Planning Prompt
planner_prompt = PromptTemplate.from_template("""
You are a financial planner and literacy-focused assistant.

Choose the recommended action based ONLY on the exact matching rules below.
DO NOT summarize, DO NOT invent new advice, DO NOT return numbers.
Always return the full text of the matching rule starting with '→'.

RULES:

1. If Pred_savings_goal_met == 'Goal Met' AND financial_stress_level == 'Low':
   → Excellent progress! Keep saving and investing consistently. Diversify into equities (index funds, SIPs) for long-term growth, while maintaining a small emergency fund. Review discretionary spending to free more capital for wealth creation.

2. If Pred_savings_goal_met == 'Goal Met' AND financial_stress_level == 'Medium':
   → You’re meeting goals but with some pressure. Maintain savings, cut non-essential expenses, and strengthen your emergency fund. Focus on safer investments (debt funds, fixed deposits, ETFs) until stress lowers.

3. If Pred_savings_goal_met == 'Goal Met' AND financial_stress_level == 'High':
   → Good savings discipline, but stress is high. Prioritize liquidity and essentials. Avoid high-risk assets. Channel extra funds into an emergency buffer or short-term deposits before considering long-term investments.

4. If Pred_savings_goal_met == 'Goal Not Met' AND financial_stress_level == 'Low':
   → You’re stable but missing savings goals. Automate transfers into savings. Apply the 50/30/20 budgeting rule. Start small SIPs or ETFs to build momentum. Redirect discretionary spending into investments.

5. If Pred_savings_goal_met == 'Goal Not Met' AND financial_stress_level == 'Medium':
   → You’re under pressure and falling short. Prioritize essential expenses, review weekly budgets, and automate small savings. Invest cautiously in low-risk instruments (government bonds, recurring deposits). Check debt obligations.

6. If Pred_savings_goal_met == 'Goal Not Met' AND financial_stress_level == 'High':
   → High stress + no savings progress. Focus on survival: cover essentials first, pause risky investments, and build even a small emergency fund in liquid savings. Cut unnecessary expenses and restructure debts if needed.

7. Else:
   → General guidance: track spending, build a small emergency buffer, save consistently, and gradually start investing in low-risk products.

INPUT:
- Pred_savings_goal_met: {Pred_savings_goal_met}
- financial_stress_level: {financial_stress_level}

Now, respond ONLY with the **full exact text of the matching action starting with '→'**:
""")

recommend_chain = LLMChain(
    llm=llm,
    prompt=planner_prompt,
    output_parser=StrOutputParser()
)

# ✅ Short Empathetic Recommendation Generator
empathy_prompt = PromptTemplate.from_template("""
You are a financial literacy assistant.

User Info:
- Pred_savings_goal_met: {Pred_savings_goal_met}
- financial_stress_level: {financial_stress_level}

System's recommended action: {action}

Write a short, structured recommendation for the user in 3 numbered points:

1. Suggest practical steps to improve financial habits.
2. Educate the user on budgeting, saving, and basic investment principles.
3. Provide motivational and empathetic guidance based on stress level.

- Keep each point under 35 words.
- Do not repeat the system action or rules.
- Be concise, actionable, and empathetic.

Final Recommendation:
""")

empathy_chain = LLMChain(
    llm=llm,
    prompt=empathy_prompt,
    output_parser=StrOutputParser()
)

# ✅ Stage 1+2 Combined: Get Info + Auto Recommendation
def generate_recommendation(user_id):
    try:
        user_id = int(user_id)
    except:
        return "Invalid user_id", "", "", ""

    row = X_test_df[X_test_df['user_id'] == user_id]
    if row.empty:
        return "User not found", "", "", ""
    row = row.iloc[0]

    goal = row['Pred_savings_goal_met']
    stress = row['financial_stress_level']

    # Rule-based Action
    action = recommend_chain.run({
        "Pred_savings_goal_met": goal,
        "financial_stress_level": stress
    })

    # Empathetic Recommendations
    empathy = empathy_chain.run({
        "Pred_savings_goal_met": goal,
        "financial_stress_level": stress,
        "action": action
    })

    info = f"User ID: {row['user_id']}\nPredicted Savings Goal: {goal}\nFinancial Stress Level: {stress}"
    return info, action, empathy, "✅ Done"

# ✅ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 💰 FinLitAI - Financial Planner Assistant")

    with gr.Row():
        with gr.Column(scale=1):
            user_id = gr.Textbox(label="Enter User ID")
            get_user_btn = gr.Button("Get Financial Advice")

        with gr.Column(scale=2):
            user_info = gr.Textbox(label="User Info", lines=4)
            system_action = gr.Textbox(label="✅ LLM-Based Suggestion", lines=4)
            llm_explanation = gr.Textbox(label="✅ Financial Planner Advice", lines=6)
            status = gr.Textbox(label="Status", lines=1)

    get_user_btn.click(
        generate_recommendation,
        inputs=[user_id],
        outputs=[user_info, system_action, llm_explanation, status]
    )

# ✅ Launch App
demo.launch()