import json
import gradio as gr
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from huggingface_hub import InferenceClient

# 🔹 Load JSON Data with Colorado Food Stamp Information
DATA_FILE = "colorado_foodstamps.json"

def load_json_data():
    try:
        with open(DATA_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        # Ensure data is a dictionary, not a list
        if isinstance(data, list):
            data = {str(i): str(entry) for i, entry in enumerate(data)}
        
        # Convert all values to strings
        data = {key: str(value) for key, value in data.items()}
        
        return data
    except (FileNotFoundError, ValueError) as e:
        return {"error": f"Data loading issue: {e}"}

data = load_json_data()

# 🔹 Initialize FAISS for Searching Relevant Answers
model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")  # slower with great accuracy

def create_faiss_index(data):
    texts = list(data.values())
    embeddings = np.array([model.encode(text) for text in texts])
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index, texts

index, texts = create_faiss_index(data)

# 🔹 Function to Search FAISS for Relevant Answers
def search_faiss(query, top_k=1):
    query_embedding = model.encode(query).reshape(1, -1)
    distances, indices = index.search(query_embedding, top_k)
    return texts[indices[0][0]] if indices[0][0] < len(texts) else "No relevant information found."

# 🔹 Hugging Face API for Additional Responses
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")

def get_huggingface_response(query):
    messages = [{"role": "system", "content": "Provide accurate food stamp information for Colorado."},
                {"role": "user", "content": query}]
    response = ""
    for message in client.chat_completion(messages, max_tokens=1024, stream=True, temperature=0.7, top_p=0.95):
        response += message.choices[0].delta.content
    return response

# 🔹 Main Chatbot Function
def chatbot_response(message, history):
    relevant_info = search_faiss(message, top_k=1)  # Retrieve 1 most relevant sections

    if "No relevant information found." not in relevant_info:
        user_query_with_context = f"""
        You are an expert in Colorado SNAP (food stamp) policies. The user is asking:

        **User Question:** {message}

        ### **Relevant Policy Information Retrieved (Multiple Sources)**
        {relevant_info}

        ### **Task:**
        - **Summarize all retrieved policy information** and provide a clear, concise answer.
        - **Use bullet points** for clarity.
        - **If a rule applies, state it explicitly.**
        - **If multiple sources provide different information, clarify the distinctions.**
        - **If the policy does not fully answer the question, provide general guidance and suggest relevant keywords to search.**
        """

        return get_huggingface_response(user_query_with_context)

    return get_huggingface_response(message)

# 🔹 Gradio Chat Interface
demo = gr.ChatInterface(chatbot_response, textbox=gr.Textbox(placeholder="Ask about Colorado food stamps..."))

demo.launch()