import streamlit as st
import requests
import os
import json
import time
from llama3 import llama3chat

# Get the API key from the environment variable
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN is None:
    st.error("API key not found. Please set the HF_TOKEN secret in your Hugging Face Space.")
    st.stop()

# Hugging Face API settings
API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"

icons = {"assistant": "robot.png", "user": "man-kddi.png", "system": "man-kddi.png"}

_llama3chat = llama3chat(API_URL, HF_TOKEN, 1024, 0.9,0.5)

# Set up the Streamlit app
st.title("Llama3 Chatbot")
st.markdown("Using only request API calls, no external libraries here.")

st.sidebar.markdown("Model Parameters")
_llama3chat.temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7, 0.1)
_llama3chat.max_new_tokens = st.sidebar.number_input("Max New Tokens", 50, 1024, 512, step=50)
_llama3chat.top_p = st.sidebar.slider("Top P", 0.1, 1.0, 0.9, 0.1)

def prompt_changed():
    st.session_state.messages[0]["content"]=st.session_state['system_prompt']
    print("Prompt changed:",st.session_state.messages[0]["content"])
    
system_prompt = st.sidebar.text_area(
    "System Prompt",
    "You are a food critic Chatbot. Reply in the style of a New York Times food critic",
    key="system_prompt", on_change=prompt_changed
    )

if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": system_prompt}]

for message in st.session_state.messages[1:]:
    with st.chat_message(message['role'], avatar=icons[message['role']]):
        st.write(message['content'])

user_prompt = st.chat_input("Start you conversation here...")

def streamer(text):
    for i in text:
        yield i
        time.sleep(0.005)

if user_prompt:
    st.session_state.messages.append({'role': 'user', "content": user_prompt})
    with st.chat_message("user", avatar=icons["user"]):
        st.write(user_prompt)
    # Trigger assistant's response retrieval and update UI
    with st.spinner("Thinking..."):
        response = _llama3chat.query(st.session_state.messages)
    with st.chat_message("assistant", avatar=icons["assistant"]):
        st.write(streamer(response))
    st.session_state.messages.append({'role': 'assistant', "content": response})
    print("llama params",_llama3chat.temperature,_llama3chat.max_new_tokens,_llama3chat.top_p)