import streamlit as st
import google.generativeai as genai
from openai import OpenAI
from huggingface_hub import InferenceClient
import time
import os

GEMINI_MODEL_ID = "gemini-2.0-flash-exp" 
gemini_api_key = os.getenv("GEMINI_API_KEY")
model_id = GEMINI_MODEL_ID
genai.configure(api_key=gemini_api_key)

gpt_api_key = os.getenv("OPENAI_API_KEY")
gptclient = OpenAI(api_key=gpt_api_key)

LLAMA_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
hf_token = os.getenv("HF_TOKEN")
llamaclient = InferenceClient(api_key=hf_token)


def ask_gemini(prompt):
    """Generates Gemini response to the  prompt."""
    try:
        model = genai.GenerativeModel(GEMINI_MODEL_ID)
        chat = model.start_chat()
            
        # Record the time before sending the prompt
        start_time = time.time()

        # Send a text prompt to Gemini API
        response = chat.send_message(
            [
                prompt
            ],
        )

        # Record the time after receiving the response
        end_time = time.time()

        # Calculate the elapsed time
        elapsed_time = end_time - start_time

        # Format the elapsed time as a string
        elapsed_time_str = f"{elapsed_time:.2f} seconds"

        return response.text, elapsed_time_str
   
    except Exception as e:
        st.error(f"An error occurred while generating AI response from Gemini: {e}")
        return None, None


def ask_gpt(prompt):
    try:

        # Record the time before sending the prompt
        start_time = time.time()

        completion = gptclient.chat.completions.create(
            model = "gpt-4o",
            messages = [{"role": "user", "content": prompt}],
            temperature=0.2,
            top_p=0.7,
            max_tokens=1024,
        )

       # Record the time after receiving the response
        end_time = time.time()

        # Calculate the elapsed time
        elapsed_time = end_time - start_time

        # Format the elapsed time as a string
        elapsed_time_str = f"{elapsed_time:.2f} seconds"

        return completion.choices[0].message.content, elapsed_time_str
        
    except Exception as e:
        st.error(f"Error processing GPT 4o response: {e}")
        return None, None

def ask_llama(prompt):
    try:
        messages = [
            {
                "role": "user",
                "content": prompt
            }
        ]

        # Record the time before sending the prompt
        start_time = time.time()

        completion = llamaclient.chat.completions.create(
            model=LLAMA_MODEL_ID, 
            messages=messages, 
            max_tokens=1024
        )

       # Record the time after receiving the response
        end_time = time.time()

        # Calculate the elapsed time
        elapsed_time = end_time - start_time

        # Format the elapsed time as a string
        elapsed_time_str = f"{elapsed_time:.2f} seconds"

        return completion.choices[0].message.content, elapsed_time_str
        
    except Exception as e:
        st.error(f"Error: {e}")
        return None, None