app.py

from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
import google.generativeai as genai
import firebase_admin
from firebase_admin import credentials, firestore
import json, os
from datetime import datetime

# =========================
# 🔑 CONFIG
# =========================

# Configure Gemini API Key
os.environ["GOOGLE_API_KEY"] = "AIzaSyBzSyySlN8ELAY3cCZp3tzWTC3YHua1Ej0"
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
model = genai.GenerativeModel("gemini-2.5-flash")

# =========================
# 🔥 FIREBASE SETUP (with secret)
# =========================

# Try loading Firebase key from Hugging Face Secrets first
firebase_key_content = os.getenv("FIREBASE_KEY")

if firebase_key_content:
    print("✅ Using Firebase key from secret environment variable")
    cred_dict = json.loads(firebase_key_content)
    cred = credentials.Certificate(cred_dict)
else:
    print("⚠️ Using local firebase_key.json (for local dev only)")
    cred = credentials.Certificate("firebase_key.json")

firebase_admin.initialize_app(cred)
db = firestore.client()

# =========================
# ⚡ FASTAPI SETUP
# =========================

app = FastAPI(
    title="🌾 Smart Farmer Advisor (Safe Memory)",
    description="Gemini AI with persistent memory stored in user subcollection (safe for frontend)",
    version="3.1"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =========================
# 🧠 MEMORY MANAGEMENT (SAFE)
# =========================

def get_user_chat_history(uid: str, session_id: str = "session_001"):
    """Retrieve user's chat history from subcollection."""
    chat_doc = db.collection("users").document(uid).collection("chatbot_history").document(session_id)
    doc = chat_doc.get()
    if doc.exists:
        return doc.to_dict().get("history", [])
    return []


def save_user_chat_history(uid: str, history: list, session_id: str = "session_001"):
    """Save chat history under user's subcollection."""
    chat_doc = db.collection("users").document(uid).collection("chatbot_history").document(session_id)
    chat_doc.set({
        "history": history,
        "last_updated": datetime.utcnow()
    })

# =========================
# 🌾 GEMINI LOGIC
# =========================

def get_farmer_schemes(query: str, uid: str, session_id: str = "session_001"):
    """Use Gemini with per-user memory stored in subcollection."""
    history = get_user_chat_history(uid, session_id)
    chat = model.start_chat(history=history)

    prompt = f"""
    You are an Indian agriculture policy expert.
    Continue assisting this farmer based on previous chat context.
    Always reply in JSON format as:
    [
      {{
        "name": "Scheme Name",
        "ministry": "Responsible Ministry/Department",
        "benefits": "Key features and benefits",
        "eligibility": "Who can apply",
        "link": "Official or reliable info link"
      }}
    ]
    Query: {query}
    """

    response = chat.send_message(prompt)

    # Update chat memory
    history.append({"role": "user", "parts": [query]})
    history.append({"role": "model", "parts": [response.text]})
    save_user_chat_history(uid, history, session_id)

    # Parse JSON safely
    text = response.text.strip()
    try:
        return json.loads(text)
    except:
        start = text.find('[')
        end = text.rfind(']')
        try:
            return json.loads(text[start:end + 1])
        except:
            return {"raw_response": text}

# =========================
# 🌐 ROUTES
# =========================

@app.get("/schemes")
def get_schemes(
    query: str = Query(..., description="User's query about farmer schemes"),
    uid: str = Query(..., description="Firebase UID of the user"),
    session_id: str = Query("session_001", description="Optional chat session ID")
):
    """
    Example:
    /schemes?query=crop+insurance+in+Punjab&uid=UID_12345
    """
    data = get_farmer_schemes(query, uid, session_id)
    return {"uid": uid, "session_id": session_id, "query": query, "schemes": data}


@app.get("/reset_chat")
def reset_chat(uid: str = Query(...), session_id: str = Query("session_001")):
    db.collection("users").document(uid).collection("chatbot_history").document(session_id).delete()
    return {"message": f"Chat history cleared for {uid} ({session_id})"}

# =========================
# 💻 RUN SERVER
# =========================

if __name__ == "__main__":
    import uvicorn
    print("🚜 Smart Farmer Advisor (Safe Memory) running at http://127.0.0.1:7860")
    uvicorn.run(app, host="0.0.0.0", port=7860)