import gradio as gr
from groq import ChatGroq
from llama_index import load_index_from_storage, ServiceContext, set_global_service_context
from llama_index.vector_stores.faiss import FaissVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.storage.storage_context import StorageContext
from PIL import Image

# Initialize Groq Vision LLM
vision_llm = ChatGroq(
    model_name="meta-llama/llama-4-scout-17b-16e-instruct",
    api_key="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"
)

# Load FAISS index from persisted directory
faiss_index_path = "faiss_store"
vector_store = FaissVectorStore.from_persist_dir(faiss_index_path)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = load_index_from_storage(storage_context)

# Setup embedding context
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en")
service_context = ServiceContext.from_defaults(embed_model=embed_model)
set_global_service_context(service_context)

# Setup query engine
query_engine = index.as_query_engine()

def multimodal_skin_rag(image, question):
    context = query_engine.query(question)
    user_prompt = f"""Based on the following skincare context, answer the user's question with reference to the image (if relevant):

Context:
{context}

Question:
{question}
"""
    messages = [
        {"role": "system", "content": "You are a skincare advisor who understands image-based inputs and medical-grade text."},
        {"role": "user", "content": user_prompt}
    ]
    response = vision_llm.chat(messages=messages, images=[image])
    return response.choices[0].message.content

demo = gr.Interface(
    fn=multimodal_skin_rag,
    inputs=[
        gr.Image(type="pil", label="Upload Skin Image"),
        gr.Textbox(label="Describe your skin concern or ask a question")
    ],
    outputs="text",
    title="SkinCare Assistant: FAISS + Groq LLM",
    description="Upload a skin image and ask any skincare-related question. This system retrieves relevant content using FAISS and answers using Groq's Vision LLM."
)

if __name__ == "__main__":
    demo.launch()