import os
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# ✅ Correct IBM Granite model
model_id = "ibm-granite/granite-3.3-2b-instruct"
token = os.getenv("HF_TOKEN")  # Load from Hugging Face Secrets

# ✅ Load model and tokenizer using the updated syntax
tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    token=token,
    device_map="auto",
    torch_dtype=torch.float32
)

def query_granite(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Pages
def home():
    with gr.Blocks() as demo:
        gr.Markdown("# 🏥 Welcome to HealthAI")
        gr.Markdown("Your intelligent healthcare assistant using IBM Granite 3.3-2B Instruct.")
    return demo

def symptoms_app():
    def identify(symptom):
        return query_granite(f"What illness could cause: {symptom}?")
    with gr.Blocks() as demo:
        gr.Markdown("## 🩺 Symptom Identifier")
        symptom = gr.Textbox(label="Enter your symptom")
        output = gr.Textbox(label="AI Diagnosis")
        btn = gr.Button("Analyze")
        btn.click(identify, inputs=symptom, outputs=output)
    return demo

def remedies_app():
    def get_remedies(issue):
        return query_granite(f"What are home remedies for {issue}?")
    with gr.Blocks() as demo:
        gr.Markdown("## 🌿 Home Remedies")
        issue = gr.Textbox(label="What are you suffering from?")
        output = gr.Textbox(label="Suggested Remedy")
        btn = gr.Button("Suggest")
        btn.click(get_remedies, inputs=issue, outputs=output)
    return demo

def diet_app():
    def suggest(goal):
        return query_granite(f"Suggest a diet for: {goal}")
    with gr.Blocks() as demo:
        gr.Markdown("## 🥗 Diet Suggestions")
        goal = gr.Textbox(label="Your health goal")
        output = gr.Textbox(label="Diet Plan")
        btn = gr.Button("Get Plan")
        btn.click(suggest, inputs=goal, outputs=output)
    return demo

def mental_app():
    def tip(topic):
        return query_granite(f"Mental health advice about: {topic}")
    with gr.Blocks() as demo:
        gr.Markdown("## 🧠 Mental Wellness")
        topic = gr.Textbox(label="Enter mental health topic")
        output = gr.Textbox(label="Wellness Tip")
        btn = gr.Button("Get Tip")
        btn.click(tip, inputs=topic, outputs=output)
    return demo

def faqs_app():
    with gr.Blocks() as demo:
        gr.Markdown("## ❓ FAQs")
        gr.Markdown("**Q1:** What is HealthAI?")
        gr.Markdown("**A:** It's an AI assistant to help with health-related queries using IBM Granite 3.3-2B.")
    return demo

# App configuration
pages = [
    gr.Page(title="🏠 Home", path="/", block=home),
    gr.Page(title="🩺 Symptoms", path="/symptoms", block=symptoms_app),
    gr.Page(title="🌿 Remedies", path="/remedies", block=remedies_app),
    gr.Page(title="🥗 Diet", path="/diet", block=diet_app),
    gr.Page(title="🧠 Mental Wellness", path="/mental", block=mental_app),
    gr.Page(title="❓ FAQs", path="/faqs", block=faqs_app),
]

app = gr.App(pages=pages)
app.launch()