import streamlit as st
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
from dotenv import load_dotenv
import os
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain_groq import ChatGroq

# Load environment variables
load_dotenv()
os.environ['GROQ_API_KEY'] = os.getenv("GROQ_API_KEY")

groq_api_key = os.getenv("GROQ_API_KEY")

# Initialize Chatbot
chat = ChatGroq(groq_api_key=groq_api_key, model_name="llama-3.3-70b-versatile")

if 'flow_messages' not in st.session_state:
    st.session_state['flow_messages'] = [
        SystemMessage(content=(
            "You are a highly intelligent assistant specializing in food safety and hygiene. "
            "You help users interpret food contamination results, provide safe food-handling practices, "
            "and answer questions related to food quality and safety."
        ))
    ]

# Function to get chatbot response
def get_response(question):
    st.session_state['flow_messages'].append(HumanMessage(content=question))
    answer = chat(st.session_state['flow_messages'])
    st.session_state['flow_messages'].append(AIMessage(content=answer.content))
    return answer.content

# Load the trained food contamination model
model = tf.keras.models.load_model('model.h5')

# Preprocess uploaded image
def preprocess_image(img, target_size=(224, 224)):
    img = img.resize(target_size)  # Resize to match the model's input size
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array = img_array / 255.0  # Normalize to [0, 1]
    return img_array

# Streamlit app setup
st.set_page_config(page_title="Food Safety Detection and Chatbot", layout="wide")
st.title("Food Safety Detection and AI Assistant")
st.write("Upload an image to determine if food is safe or contaminated.")

# Image classification section
st.header("Food Contamination Detection")
uploaded_file = st.file_uploader("Upload an image of food", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    try:
        img = Image.open(uploaded_file)
        st.image(img, caption="Uploaded Image")
        img_array = preprocess_image(img)
        prediction = model.predict(img_array)
        if prediction[0] > 0.5:
            st.success("Prediction: Good Food")
        else:
            st.error("Prediction: Contaminated Food")
    except Exception as e:
        st.error(f"An error occurred while processing the image: {e}")

# Sidebar Chatbot
st.sidebar.title("Chatbot Assistant")

with st.sidebar:
    if 'chat_history' not in st.session_state:
        st.session_state['chat_history'] = []

    st.write("Ask me anything related to food safety and hygiene:")
    user_input = st.text_input("Your Question", key="sidebar_input")
    if st.button("Send", key="sidebar_send"):
        if user_input.strip():
            response = get_response(user_input)
            st.session_state['chat_history'].append((user_input, response))

    if st.button("Clear Chat"):
        st.session_state['chat_history'] = []
        st.session_state['flow_messages'] = [
            SystemMessage(content=(
                "You are a highly intelligent assistant specializing in food safety and hygiene. "
                "You help users interpret food contamination results, provide safe food-handling practices, "
                "and answer questions related to food quality and safety."
            ))
        ]

    st.write("Chat History:")
    for question, answer in st.session_state['chat_history']:
        st.markdown(f"**You:** {question}")
        st.markdown(f"**Bot:** {answer}")