import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
import torch
import re

# Initialize models and client
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")

# Load Medicaid/health info
with open("medicaid_info.txt", "r", encoding="utf-8") as file:
    health_info = file.read()

# Load clinic info
with open("info.txt", "r", encoding="utf-8") as file:
    clinic_data = file.read().lower()

# Preprocess info into chunks
def preprocess_chunks(text):
    chunks = [chunk.strip() for chunk in text.split('.') if chunk.strip()]
    return chunks

health_chunks = preprocess_chunks(health_info)
health_embeddings = embedding_model.encode(health_chunks, convert_to_tensor=True)

# Semantic search
def get_relevant_chunks(query, embeddings, text_chunks):
    query_embedding = embedding_model.encode(query, convert_to_tensor=True)
    query_embedding = query_embedding / query_embedding.norm()
    embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)
    similarities = torch.matmul(embeddings, query_embedding)
    top_indices = torch.topk(similarities, k=3).indices
    return [text_chunks[i] for i in top_indices]

# Clinic finder
def find_clinic_by_county(county):
    county = county.lower()
    sections = clinic_data.split("###")
    for section in sections:
        if county in section:
            lines = section.strip().split("\n")
            if lines and county in lines[0].lower():
                lines = lines[1:]
            return "\n".join(lines).strip()
    return "⚠️ Sorry, I couldn’t find clinics for that county. Check spelling or try a nearby county."

# Chatbot logic
def respond(message, history, name, focus_area):
    focus = focus_area[0] if focus_area else "general help"
    top_chunks = get_relevant_chunks(message, health_embeddings, health_chunks)
    context = "\n".join(top_chunks)

    messages = [
        {
            "role": "system",
            "content": (
                f"You are a friendly and supportive health advisor chatbot helping Washington residents, especially low-income individuals. "
                f"You're chatting with {name}. Speak clearly and kindly, and keep responses under 100 words. Focus on {focus}."
            )
        },
        {
            "role": "user",
            "content": f"Context:\n{context}\n\nQuestion: {message}"
        }
    ]
    
    response = client.chat_completion(messages, max_tokens=120)
    return response['choices'][0]['message']['content'].strip()

# Interface styling
title = "# 🩺 HealthHelpBot"
description = "### Helping Washington residents find care, understand Medicaid, and locate free clinics."

with gr.Blocks(theme="gstaff/soft-blue") as demo:
    with gr.Row():
        gr.Markdown(title)
    with gr.Row():
        gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Image("healthbot_avatar.png", width=250, show_label=False)
            user_name = gr.Textbox(label="👤 Your Name", placeholder="e.g. Maria")
            focus_area = gr.CheckboxGroup(["Medicaid", "Finding a clinic", "Insurance help"], label="What do you need help with?")
        with gr.Column(scale=2):
            gr.ChatInterface(fn=respond, additional_inputs=[user_name, focus_area], type="messages")

demo.launch()