import streamlit as st
import requests
import json
import base64
from PIL import Image
import io

# Function to encode the image
def encode_image(image_file):
    return base64.b64encode(image_file.getvalue()).decode('utf-8')

# Set up the Streamlit page
st.set_page_config(page_title="LLM Chat Interface", page_icon="🤖")
st.title("LLM Chat Interface")

# Input for API endpoint and token
api_endpoint = st.sidebar.text_input("API Endpoint", "http://209.222.10.17:8000/v1/chat/completions")
api_token = st.sidebar.text_input("API Token", type="password")

# File uploader for image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Text input for user message
user_input = st.text_area("Enter your message:", height=100)

if st.button("Send"):
    if not api_token:
        st.error("Please enter an API token.")
    elif not user_input:
        st.error("Please enter a message.")
    else:
        # Prepare the message content
        content = [{"type": "text", "text": user_input}]
        
        # If an image is uploaded, add it to the content
        if uploaded_file is not None:
            # Encode and add the image
            image_base64 = encode_image(uploaded_file)
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}
            })

        # Prepare the request payload
        payload = {
            "model": "mistralai/Pixtral-12B-2409",
            "messages": [{"role": "user", "content": content}]
        }

        # Send the request to the API
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_token}"
        }

        with st.spinner("Waiting for response..."):
            try:
                response = requests.post(api_endpoint, headers=headers, json=payload)
                response.raise_for_status()
                result = response.json()
                
                # Display the response
                st.subheader("Response:")
                st.write(result['choices'][0]['message']['content'])
            except requests.exceptions.RequestException as e:
                st.error(f"An error occurred: {str(e)}")

# Display the uploaded image
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)