import os
import streamlit as st
from groq import Groq
import PyPDF2
import pdfplumber

# Initialize Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Function to process PDF report
def process_pdf(file):
    text = ""
    with pdfplumber.open(file) as pdf:
        for page in pdf.pages:
            text += page.extract_text()
    return text

# Function to query Groq API
def query_groq(prompt):
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.3-70b-versatile",
    )
    return chat_completion.choices[0].message.content

# Streamlit interface
st.set_page_config(page_title="vDoc", layout="wide")

st.title("🌟 vDoc: Your Virtual Doctor 🌟")
st.sidebar.header("User Input Options")

# Input Section
symptoms = st.sidebar.text_area("Describe your symptoms")
uploaded_files = st.sidebar.file_uploader(
    "Upload Medical Test Reports (PDF)", type=["pdf"], accept_multiple_files=True
)
medical_history = st.sidebar.text_area("Medical history or medications")

if st.sidebar.button("Submit"):
    st.subheader("Diagnosis")
    
    # Collect and process inputs
    input_data = []
    if symptoms:
        input_data.append(f"Symptoms: {symptoms}")
    if uploaded_files:
        for file in uploaded_files:
            report_text = process_pdf(file)
            input_data.append(f"Test Report: {report_text}")
    if medical_history:
        input_data.append(f"Medical History: {medical_history}")

    # Send inputs to Groq API for analysis
    combined_input = "\n".join(input_data)
    if combined_input:
        response = query_groq(f"Diagnose the following medical issue:\n{combined_input}")
        st.write(response)
    else:
        st.warning("Please provide at least one input for diagnosis.")

# Chatbot Section
st.subheader("🤖 Chat with vDoc")
chat_input = st.text_input("Your question or response", "")
if st.button("Ask"):
    if chat_input:
        chat_response = query_groq(chat_input)
        st.write(chat_response)
    else:
        st.warning("Please enter a question or response.")

# Additional Features
st.sidebar.subheader("Additional Features")
if st.sidebar.button("Explain Test Reports"):
    if uploaded_files:
        for file in uploaded_files:
            report_text = process_pdf(file)
            explanation = query_groq(f"Explain the following test report:\n{report_text}")
            st.subheader(f"Explanation for {file.name}")
            st.write(explanation)
    else:
        st.sidebar.warning("Upload test reports to get explanations.")

medicine_name = st.sidebar.text_input("Enter Medicine Name for Details")
if st.sidebar.button("Get Medicine Info"):
    if medicine_name:
        medicine_info = query_groq(f"Explain the use and details of the medicine: {medicine_name}")
        st.subheader(f"Details about {medicine_name}")
        st.write(medicine_info)
    else:
        st.sidebar.warning("Enter a medicine name to get details.")