# llm_part.py
import os
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from datetime import datetime


# Initialize LLM models
gemeni_key=os.getenv("gemeni_key")
llama_key=os.getenv("llama_key")
llm_1 = ChatGoogleGenerativeAI(model="gemini-pro", api_key=gemeni_key)
llm_2 = ChatGroq(model="llama3-groq-70b-8192-tool-use-preview", groq_api_key=llama_key)

# Prayer times for the current day only (Fajr, Dhuhr, Asr, Maghrib, Isha)
prayer_times = {
    "Fajr": "04:30",
    "Dhuhr": "11:45",
    "Asr": "15:30",
    "Maghrib": "17:45",
    "Isha": "19:00"
}

# Define the next prayer function
def get_next_prayer(current_time):
    # Retrieve today's prayer times
    times = list(prayer_times.values())
    prayer_names = list(prayer_times.keys())

    # Current time in comparable format
    current_hour = current_time.hour
    current_minute = current_time.minute

    for prayer, time_str in zip(prayer_names, times):
        prayer_hour, prayer_minute = map(int, time_str.split(':'))
        if (prayer_hour > current_hour) or (prayer_hour == current_hour and prayer_minute > current_minute):
            return prayer, time_str  # Return both prayer name and time
        
    #return "Isha", times[-1] 
    return prayer_times, times[-1] 


templates = {
    "Learning": "You just completed a learning task: {task_name}. You studied for {hours} hours and {minutes} minutes. "
                "Can you suggest improvements to my learning process? What should I learn next?",
    
    "Gym": "You completed a gym task: {task_name}. It lasted {hours} hours and {minutes} minutes. "
           "How can I optimize my workout? What should I focus on next time in the gym?",
    
    "Personal": "I just completed a personal task: {task_name}. It took me {hours} hours and {minutes} minutes. "
                "What advice can you give for better time management or improvement?",
    
    "Work": "I completed a work-related task: {task_name}, which lasted {hours} hours and {minutes} minutes. "
            "Can you suggest ways to improve my work efficiency? What should I work on next?",
    
    "Prayer": "You completed a prayer task: {task_name}. It took you {hours} hours and {minutes} minutes. "
              "What practical steps can I take to deepen my spiritual connection during Salah (prayer)?"
              "Please provide specific traditions and teachings for the next prayer, **{next_prayer}**, which is at **{next_time}**. "
              "Format your response as bullet points for clarity."
}

# Function to select a template based on the task category
def get_template(category):
    return templates.get(category, "You just completed the task: {task_name}. It took {hours} hours and {minutes} minutes. "
                                   "What advice can you give? What is the next most productive task?")

# Template function to generate advice and next task recommendations for each task
def get_task_advice(task_name, task_category, hours, minutes):
    # Get the template based on the category
    template_str = get_template(task_category)

    # Determine the next prayer and its time
    current_time = datetime.now()
    next_prayer, next_time = get_next_prayer(current_time)

    # Create the prompt
    prompt_template = PromptTemplate(
        input_variables=["task_name", "hours", "minutes", "next_prayer", "next_time"],
        template=template_str
    )
    
    prompt = prompt_template.format(task_name=task_name, hours=hours, minutes=minutes, 
                                     next_prayer=next_prayer, next_time=next_time)

    # Set up LLMChain for both models
    chain_1 = LLMChain(llm=llm_1, prompt=prompt_template)
    chain_2 = LLMChain(llm=llm_2, prompt=prompt_template)

    try:
        # Run the LLMChain to get the responses
        response_llm1 = chain_1.run({"task_name": task_name, "hours": hours, "minutes": minutes, 
                                       "next_prayer": next_prayer, "next_time": next_time})
    except Exception as e:
        response_llm1 = f"Error generating response from Gemini: {str(e)}"
    
    try:
        response_llm2 = chain_2.run({"task_name": task_name, "hours": hours, "minutes": minutes, 
                                       "next_prayer": next_prayer, "next_time": next_time})
    except Exception as e:
        response_llm2 = f"Error generating response from Llama: {str(e)}"
    
    return response_llm1, response_llm2