import gradio as gr
from datetime import datetime
from enum import Enum
from pymongo import MongoClient
import PIL
import google.generativeai as genai
import re
from langchain_groq import ChatGroq
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

# MongoDB connection
client = MongoClient('mongodb://miniproject:miniproject@ac-yzyqqis-shard-00-00.iv8gz1x.mongodb.net:27017,ac-yzyqqis-shard-00-01.iv8gz1x.mongodb.net:27017,ac-yzyqqis-shard-00-02.iv8gz1x.mongodb.net:27017/?ssl=true&replicaSet=atlas-ayivip-shard-0&authSource=admin&retryWrites=true&w=majority')
db = client['Final_Year_Project']
collection = db['Patients']

# Langchain Groq API setup
groq_api_key = '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'
llm_agent = ChatGroq(api_key=groq_api_key, model='llama-3.3-70b-versatile', temperature=0.1)

# Set up the prompt template
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            f"You Are an AI Assistant which helps to manage the reminders for Patients. It is currently {datetime.now()}.",
        ),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)

# Enum for reminder status
class ReminderStatus(Enum):
    ACTIVE = "Active"
    COMPLETED = "Completed"
    CANCELLED = "Cancelled"

# Tool for saving user data to MongoDB
@tool
def save_user_data(user_id: str, patient_name: str, dr_name: str, prescription_date: datetime,
                   age: int, sex: str, medicines: list, notification_type: str = "Push",
                   notification_time: int = 30, status: ReminderStatus = ReminderStatus.ACTIVE):
    """
    Adds a new patient-related reminder to the MongoDB collection, allowing multiple medicines for Flutter notifications.

    Args:
        user_id: Identifier for the user creating the reminder.
        patient_name: Name of the patient.
        dr_name: Name of the doctor.
        prescription_date: ISO Format Date of the prescription (typically today) for example, "2000-10-31T01:30:00.000-05:00".
        age: Age of the patient.
        sex: Sex of the patient.
        medicines: List of medicines, where each medicine is a dictionary with 'name', 'dosage', 'frequency', and 'refill_date'.
        notification_type: Type of notification to send (e.g., "Push" for app notifications).
        notification_time: Time before the reminder should trigger (in minutes).
        status: Current status of the reminder (active, completed, etc.).

    Returns:
        The inserted document ID.
    """
    
    reminder = {
        "user_id": user_id,
        "patient_name": patient_name,
        "dr_name": dr_name,
        "prescription_date": prescription_date.isoformat(),
        "age": age,
        "sex": sex,
        "medicines": medicines,
        "notification": {
            "type": notification_type,
            "time_before": notification_time
        },
        "status": status.value,
        "created_at": datetime.now()
    }
    
    result = collection.insert_one(reminder)
    return result.inserted_id

tools = [save_user_data]
agent = create_tool_calling_agent(llm_agent, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools)

# API key for Google Generative AI
genai.configure(api_key="AIzaSyCltTyFKRtgLQBR9BwX0q9e2aDW9BwwUfo")

# Function to process the image and generate text using Google AI
def extract_and_store_text(file, user_id):
    try:
        # Open the image using PIL
        image = PIL.Image.open(file)

        # Generate text using Google Generative AI
        mod = genai.GenerativeModel(model_name="gemini-2.0-flash")
        text = mod.generate_content(contents=["Extract all text from this image, including medication names, dosages, instructions, patient information, and any other relevant medical details if this is a prescription. If the image is not a prescription, return only 'None'.", image])
        
        pattern = r'[_*]+\s*'
        data_str = re.sub(pattern, '', text.text)
        
        if not text.text:
            return "No relevant text found"
        
        else:
            res = agent_executor.invoke({"input": f'Add this Data to Mongo DB for user with user_id {user_id} DATA : {str(data_str)}'})
            ans = collection.find_one({'user_id': user_id}, {'_id': 0})
            return str(ans)
    
    except Exception as e:
        return str(e)

# Gradio interface
def gradio_interface(file, user_id):
    return extract_and_store_text(file, user_id)

# Define Gradio Inputs and Outputs
inputs = [
    gr.File(label="Upload Image"),
    gr.Textbox(label="User ID", value="101")  # Default user ID
]
outputs = gr.Textbox(label="Response")

# Create Gradio Interface
interface = gr.Interface(
    fn=gradio_interface, 
    inputs=inputs, 
    outputs=outputs,
    title="Patient Reminder System", 
    description="Extracts and stores patient prescription details from images."
)

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