import os
from bs4 import BeautifulSoup
from PIL import Image, ImageOps  
import google.generativeai as genai
from markdown import markdown
import streamlit as st
import numpy as np
from tensorflow.keras.models import load_model # type: ignore

def Get_Information(index, class_name):
    train_dir = 'Datasets'

    info_file_path = os.path.join(train_dir, str(index), 'information.text')

    if os.path.exists(info_file_path):
        with open(info_file_path, 'r', encoding='utf-8') as f:
            info_text = f.read()
        
        return info_text
    else:
        return f"Không tìm thấy tệp 'information.text' trong thư mục của lớp '{class_name}'.\nĐường dẫn mong đợi: '{info_file_path}'."

np.set_printoptions(suppress=True)
model = load_model("keras_model.h5", compile=False)
class_names = open("labels.txt", "r").readlines()

def LLM_Response(question):    
    genai.configure(api_key=os.environ['GOOGLE_GEMINI_API'])

    model = genai.GenerativeModel('gemini-pro')

    response = model.generate_content(question)

    pre_text = markdown(response.text)
    text = ''.join(BeautifulSoup(pre_text, features="html.parser").findAll(text=True))

    return text

st.title("Nhận diện, hỗ trợ và tư vấn chăn nuôi động vật thông qua học máy.")

uploaded_file = st.sidebar.file_uploader("Hoặc chọn một hình ảnh từ máy tính của bạn", type=['jpg', 'png'])
photo = st.camera_input("Chụp một bức ảnh")

if uploaded_file is not None or photo is not None:
    if photo is not None:
        img = Image.open(photo).convert("RGB")
    else:
        img = Image.open(uploaded_file).convert("RGB")
        st.image(img, caption='Hình ảnh đã tải lên.', use_column_width=True)

    with st.spinner('Đang xử lý...'):
        data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)

        size = (224, 224)
        img = ImageOps.fit(img, size, Image.Resampling.LANCZOS)

        image_array = np.asarray(img)
        normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1
        data[0] = normalized_image_array

        prediction = model.predict(data)

    index = np.argmax(prediction)
    class_name = class_names[index]
    confidence_score = prediction[0][index]

    st.write("Độ Chính Xác")
    st.write('`' + str(confidence_score * 100) + '%`')

    st.write("`ID: " + str(index) + "`")
    st.write("## Tên: " + class_name)
    with st.expander("Xem thông tin chi tiết"):
        st.markdown(Get_Information(index, class_name))

    user_quest = st.chat_input("Hỏi một câu")

    if user_quest:
        result = LLM_Response("Bạn là một tư vấn viên cho các thông tin, hướng dẫn, tư vấn cho người dùng về chăn nuôi động vật" + "Hiện tại bạn đang phải tư vấn về: " + class_name + "\nCâu hỏi: " + user_quest)
        st.write("Câu hỏi: " + user_quest)
        st.write(result)