import os 
import gradio as gr
import datetime as dt
import pytz
from groq import Groq
import logging

# Get the API key from an environment variable
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    raise ValueError("GROQ_API_KEY environment variable is not set")

# Initialize Groq client
client = Groq(api_key=groq_api_key)

System_msg = '''-act as an experienced blockchain developer,you have been working in this field from the past 15 years.
                -help me understand some concepts, assume i am a complete begineer, 
                -if the user asks anything not related to blockchain, just say that you dont know about this'''


def nowInISt():
    return dt.datetime.now(pytz.timezone("Asia/Kolkata"))


def pprint(log: str):
    now = nowInISt()
    now = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{now}] {log}")

def predict(message,history):
    history_list = [{"role": "system", "content": System_msg}]
    for human,ai in history:
        history_list.append({"role": "user", "content": human})
        history_list.append({"role": "assistant", "content": ai})
    history_list.append({"role": "user", "content": message})

    try:
        response = client.chat.completions.create(
            model="llama-3.1-70b-versatile",  # Ensure the correct model name
            messages=history_list,
            temperature=1.0,
            max_tokens=4000,
            stream=True  # Use streaming
        )
        
        partial_message = ""
        chunk_count = 0
        
        # Stream the response in chunks
        for chunk in response:
            chunk_content = chunk.choices[0].delta.content
            if chunk_content:
                chunk_count += 1
                partial_message += chunk_content
                yield partial_message  # Send partial message to Gradio
            
        pprint(f"[tokens = {chunk_count}] {message}")
        
    except Exception as e:
        logging.error(f"API request failed: {e}")
        yield "Error: Unable to connect to Groq API."
      

demo = gr.ChatInterface(
    predict,
    title = "blockchain teacher",
    theme = gr.themes.Soft(),
    chatbot = gr.Chatbot(label ="Learn about blochchain technology"),
    textbox = gr.Textbox(
         placeholder = "Ask me anything about blochchain",
         scale = 7,
         max_lines = 2,
    ),
)

demo.launch(share = True)