import streamlit as st
import os
import json
import faiss
import numpy as np
import google.generativeai as genai
from sentence_transformers import SentenceTransformer
import requests
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
SERPER_API_KEY = os.getenv("SERPER_API_KEY")

# Configure Gemini
genai.configure(api_key=GOOGLE_API_KEY)

# Load local knowledge base
with open("data.txt", "r", encoding="utf-8") as f:
    snippets = [line.strip() for line in f if line.strip()]

# Initialize embedding model and FAISS
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = embed_model.encode(snippets)
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(np.array(embeddings))

def search_serper(query):
    url = "https://google.serper.dev/search"
    headers = {"X-API-KEY": SERPER_API_KEY}
    payload = {"q": query}
    res = requests.post(url, headers=headers, json=payload)
    return res.json().get("organic", [])[:3]

def ask_gemini(user_input, predicted_condition, local_docs, web_results):
    prompt = f"""
⚠️ *Disclaimer*: This response is not a substitute for professional medical advice.

**Patient Input**: {user_input}

**Likely Condition(s)**: {predicted_condition}

**Local Medical Knowledge**:
{chr(10).join(f"- {doc}" for doc in local_docs)}

**Web Evidence**:
{chr(10).join(f"- {r['title']}: {r['snippet']}" for r in web_results)}

Format the response using bullet points and keep it within 250 words. Cite sources at the end by title or snippet name. keep *Disclaimer*
"""
    model = genai.GenerativeModel("gemini-2.0-flash-001")
    response = model.generate_content(prompt)
    return response.text

def predict_condition(user_input):
    input_lower = user_input.lower()
    if "sugar" in input_lower or "glucose" in input_lower:
        return "Diabetes or Hypoglycemia"
    elif "chest" in input_lower or "pain" in input_lower:
        return "Myocardial Infarction or Angina"
    elif "urine" in input_lower or "swelling" in input_lower:
        return "AKI or CKD"
    elif "heartbeat" in input_lower:
        return "Arrhythmia"
    return "Undetermined – please consult a doctor"

def get_local_knowledge(user_input, top_k=3):
    q_emb = embed_model.encode([user_input])
    D, I = index.search(np.array(q_emb), k=top_k)
    return [snippets[i] for i in I[0]]

# Streamlit UI
st.title("🩺 Patient-Safety–Aware Chatbot")
st.markdown("Enter symptoms to receive safe, evidence-informed first-aid guidance.")

user_input = st.text_area("Describe symptoms here")

if st.button("Get First-Aid Guidance") and user_input:
    with st.spinner("Analyzing..."):
        condition = predict_condition(user_input)
        local_hits = get_local_knowledge(user_input)
        web_hits = search_serper(user_input)
        response = ask_gemini(user_input, condition, local_hits, web_hits)

    st.markdown("---")
    st.subheader("🛟 First-Aid Guidance")
    st.markdown(response)