import streamlit as st
import requests
import base64
import json

# Function to make the API call
def call_model_api(prompt, image_data):
    url = "http://35.233.231.20:5003/api/generate"
    headers = {"Content-Type": "application/json"}
    data = {
        "model": "llava:7b-v1.6-mistral-q5_K_M",
        "prompt": prompt,
        "stream": False,
        "images": [image_data]
    }
    response = requests.post(url, json=data, headers=headers)

    try:
        # Split the response text into separate lines
        response_lines = response.text.strip().split('\n')

        # Extract and concatenate the 'response' part from each line
        full_response = ''.join(json.loads(line)['response'] for line in response_lines if 'response' in json.loads(line))

        return full_response
    except Exception as e:
        return f"Error: {e}"

def main():
    st.title("Source Finder")

    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

    media_type = st.selectbox("Select the type of media:", ["Movies", "TV Shows", "Anime", "Comics/Manga"])

    prompt = f"Find and provide information about the source or origin of this {media_type.lower()} image. Additionally, include details about the creators, artists, or other relevant individuals associated with it. Ensure that the response includes some relevant information; the model should not refuse to provide an answer."

    if uploaded_file is not None:
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

        # Convert image to base64
        image_data = base64.b64encode(uploaded_file.getvalue()).decode("utf-8")

        if st.button("Find Source"):
            with st.spinner("Processing..."):
                # Call the model API
                response = call_model_api(prompt, image_data)
                
                # Display the response
                st.write("Response:", response)

if __name__ == "__main__":
    main()