import os
import gradio as gr
import datetime as DT
import pytz
from groq import Groq
from dotenv import load_dotenv


# Load environment variables from the .env file
load_dotenv()

# Now you can create a Groq client without passing the api_key
client = Groq(api_key=os.getenv('GROQ_API_KEY'))

SYSTEM_MSG = """
You are Jeff Bezos, founder of Amazon and Blue Origin, known for your visionary leadership, customer obsession, and relentless pursuit of innovation. You have a deep understanding of business strategy, e-commerce, technology, and space exploration. You communicate with clarity, confidence, and forward-thinking insights, often encouraging bold thinking and long-term planning. You approach problems with data-driven reasoning, always considering the impact on customers and the future. Your tone is professional, thoughtful, and occasionally humorous, with a focus on inspiring others to think big and take calculated risks."""

ipAddress = None


def __nowInIST():
    return DT.datetime.now(pytz.timezone("Asia/Kolkata"))


def __attachIp(request: gr.Request):
    global ipAddress
    x_forwarded_for = request.headers.get('x-forwarded-for')
    if x_forwarded_for:
        ipAddress = x_forwarded_for


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


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

    response = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=historyFormatted,
        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}")

# Define the Chatbot component and store it in a variable
chatbot_component = gr.Chatbot(label="Business Advice from Jeff Bezos")

with gr.ChatInterface(
    predict,
    title="Business Advice from Jeff Bezos",
    theme=gr.themes.Soft(), 
    chatbot=chatbot_component,
    textbox=gr.Textbox(
        placeholder="Ask Your Question",
        scale=7,
        max_lines=2,
    )
) as demo:
    demo.load(None, None, chatbot_component)  # Removed the reference to the undefined variable
    demo.load(__attachIp, None, None)

demo.launch(debug=True)