import streamlit as st
import requests
import base64

# Function to encode image in base64
def encode_image(image):
    return base64.b64encode(image.read()).decode('utf-8')

# Function to send data to the GPT-4 Vision API
def analyze_images(images, use_case, api_key):
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}

    prompt_text = (
        f"This is a prototype for {use_case}.\n"
        "Review each screenshot carefully, focusing on different aspects of usability.\n"
        "Evaluate each screenshot based on Nielsen's 10 usability heuristics:\n"
        "- Visibility of system status\n"
        "- Match between system and the real world\n"
        "- User control and freedom\n"
        "- Consistency and standards\n"
        "- Error prevention\n"
        "- Recognition rather than recall\n"
        "- Flexibility and efficiency of use\n"
        "- Aesthetic and minimalist design\n"
        "- Help users recognize, diagnose, and recover from errors\n"
        "- Help and documentation\n"
        "Create a table with the following columns:\n"
        "Heuristic evaluated, Comments (Provide specific observations and suggestions for improvement)"
    )

    messages_content = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt_text}
            ]
        }
    ]

    for image in images:
        base64_image = encode_image(image)
        messages_content[0]["content"].append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
        })

    data = {
        "model": "gpt-4-turbo",
        "messages": messages_content,
        "max_tokens": 1024  # Adjusted max_tokens to give more room for detailed analysis
    }

    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=data)
    return response.json()

# Streamlit UI setup
st.title("Auto Heuristic Evaluation")
uploaded_files = st.file_uploader("Upload images (up to 5)", accept_multiple_files=True, type=['jpg', 'jpeg', 'png'])
use_case = st.text_input("Use Case - Describe the use case for your prototype")
api_key = st.text_input("OpenAI API Key", type="password")

if st.button("Analyze Images"):
    if not uploaded_files:
        st.warning("Please upload at least one image.")
    elif not use_case:
        st.warning("Please enter the use case description.")
    elif not api_key:
        st.warning("Please enter the OpenAI API key.")
    else:
        with st.spinner("Analyzing images..."):
            analysis_result = analyze_images(uploaded_files, use_case, api_key)
            if 'error' in analysis_result:
                st.error(f"Analysis failed: {analysis_result.get('error', {}).get('message', 'Unknown error')}")
            else:
                st.success("Analysis complete!")
                if 'choices' in analysis_result and analysis_result['choices']:
                    content = analysis_result['choices'][0]['message']['content']
                    st.markdown(content)
                else:
                    st.warning("No detailed results available.")