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

# 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})

    response = client.chat.completions.create(
        model = "llama-3.1-70b-versatile",
        messages = history_list,
        temperature = 1.0,
        max_tokens=4000,
        stream = True   
    )

    partialMessage = ""
    chunkCount = 0
    for chunk in response:
        chunkContent = chunk.choices[0].delta.content
        if chunkContent:
            chunkCount+=1
            partialMessage= partialMessage + chunkContent
            yield partialMessage

    pprint(f"[tokens = {chunkCount}] {message}")        

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)