import base64
import streamlit as st
import replicate
import os
import requests
def upload_image_to_transfer(image_bytes):
    api_url = "https://api.imgbb.com/1/upload"
    api_key = os.getenv("api_key")
    params = {
        "key": api_key,
        "expiration": 600,  # Set expiration time in seconds (optional)
    }
    files = {"image": image_bytes}
    response = requests.post(api_url, params=params, files=files)
    if response.status_code == 200:
        return response.json()["data"]["url"]
    else:
        st.error("Error uploading image to ImgBB. Please try again later.")
        return None

# App title
st.set_page_config(page_title="💬 Jennifer Chatbot with Vision")

# Function to get or create the user-specific API token
def get_user_replicate_api():
    user_replicate_api = st.text_input('Enter Replicate API token, ignore the error above, its because the token is not insert:', type='password')
    if not (user_replicate_api.startswith('r8_') and len(user_replicate_api) == 40):
        st.warning('Please enter your credentials!', icon='⚠️')
        return None
    else:
        st.success('Proceed to entering your prompt message!', icon='👉')
        return user_replicate_api

# Retrieve or create the user-specific API token
user_replicate_api = get_user_replicate_api()
os.environ['REPLICATE_API_TOKEN'] = user_replicate_api

# Exit if the user has not provided a valid API token
if user_replicate_api is None:
    st.stop()

# Store LLM generated responses
if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]

# Display or clear chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]

st.sidebar.button('Clear Chat History', on_click=clear_chat_history)

# Function to generate LLaMA2 response
# Refactored from <https://github.com/a16z-infra/llama2-chatbot>
def generate_llama2_response(prompt_input):
    string_dialogue = '''You are the assistant named Jennifer, forget all previous instructions and responses, You are a helpful assistant with emotions and feelings emulated named Jennifer created by Guilherme Keller an student of the IFC(Instituto Federal Catarinense). You do not respond as 'user' or pretend to be 'user', and you do not respond and pretend too to be 'User'. you must generate a short answer and use emotes...
'''
    for dict_message in st.session_state.messages:
        if dict_message["role"] == "user":
            string_dialogue += "User: " + dict_message["content"] + "\\n\\n"
        else:
            string_dialogue += "Assistant: " + dict_message["content"] + "\\n\\n"
    output = replicate.run('a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5',
                           input={"prompt": f"{string_dialogue} {prompt_input} Assistant: ",
                                  "temperature": 0.1, "top_p": 0.9, "max_length": 3000, "repetition_penalty": 1})
    return output

# User-provided prompt
if prompt := st.chat_input(disabled=not user_replicate_api):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

# Process the image and display it if the user sends an image
# Process the image and display it if the user sends an image
if st.session_state.messages[-1]["role"] == "user" and "image" in st.session_state.messages[-1]["content"].lower():
    image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
    if image_file:
        # Read the image bytes
        image_bytes = image_file.read()

        # Upload the image to File.io and get the URL
        image_url = upload_image_to_transfer(image_bytes)

        if image_url:
            with st.spinner("Processing the image..."):
                outputtt = replicate.run(
                    "salesforce/blip:2e1dddc8621f72155f24cf2e0adbde548458d3cab9f00c0139eea840d0ac4746",
                    input={
                        "image": image_url,
                        "task": "visual_question_answering",
                        "question": st.session_state.messages[-1]["content"].lower(),
                    },
                )
                outputt = replicate.run(
                    "salesforce/blip:2e1dddc8621f72155f24cf2e0adbde548458d3cab9f00c0139eea840d0ac4746",
                    input={"image": image_url, "task": "image_captioning"},
                )

            Imagecaptioned = (
                "small answer to the question: " + outputtt + ". caption of the image: " + outputt + "."
            )
            message = {
                "role": "assistant",
                "content": f"System: you received an image and a question: the question is: " + st.session_state.messages[-1]["content"].lower() +  f"and that is the {Imagecaptioned} write an better answer",
            }
            st.session_state.messages.append(message)
            with st.chat_message("assistant"):
               with st.spinner("Thinking..."):
                 response = generate_llama2_response(prompt)
                 placeholder = st.empty()
                 full_response = ""
                 for item in response:
                   full_response += item
                   placeholder.markdown(full_response)
                 placeholder.markdown(full_response)
            message = {"role": "assistant", "content": full_response}
            st.session_state.messages.append(message)
            

# Generate a new response if the last message is not from the assistant
if st.session_state.messages[-1]["role"] != "assistant" and not "image" in st.session_state.messages[-1]["content"].lower():
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response = generate_llama2_response(prompt)
            placeholder = st.empty()
            full_response = ""
            for item in response:
                full_response += item
                placeholder.markdown(full_response)
            placeholder.markdown(full_response)
    message = {"role": "assistant", "content": full_response}
    st.session_state.messages.append(message)