import streamlit as st
import io
import PyPDF2
from google import genai
from google.genai import types
import os

# Sử dụng biến môi trường GOOGLE_API_KEY đã được thiết lập
client = genai.Client()

api_key = os.getenv('GOOGLE_API_KEY')
if not api_key:
    st.error("Biến môi trường GOOGLE_API_KEY chưa được thiết lập!")
else:
    st.write("API key đã được cấu hình.")

st.title("Gemini 2.5 Pro Chatbot")
st.write("Chatbot sử dụng Google Gen AI SDK (Gemini Developer API)")

# Nhập nội dung chat từ người dùng
user_text = st.text_input("Nhập nội dung chat:")

# Cho phép upload file với các định dạng: jpg, jpeg, png, mp4, mp3, wav, pdf
uploaded_files = st.file_uploader(
    "Tải lên file (hỗ trợ: jpg, jpeg, png, mp4, mp3, wav, pdf):",
    type=["jpg", "jpeg", "png", "mp4", "mp3", "wav", "pdf"],
    accept_multiple_files=True
)

# Tạo danh sách contents để gửi lên model
contents = []

if user_text:
    contents.append(user_text)

if uploaded_files:
    for file in uploaded_files:
        file_bytes = file.read()
        file_ext = file.name.split('.')[-1].lower()
        if file_ext == "pdf":
            # Trích xuất nội dung văn bản từ PDF
            try:
                pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_bytes))
                pdf_text = ""
                for page in pdf_reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        pdf_text += page_text + "\n"
                if pdf_text:
                    contents.append(pdf_text)
                else:
                    st.warning(f"Không thể trích xuất nội dung từ {file.name}.")
            except Exception as e:
                st.error(f"Lỗi khi xử lý PDF {file.name}: {e}")
        else:
            # Xác định mime type cho các file khác
            if file_ext in ["jpg", "jpeg"]:
                mime_type = "image/jpeg"
            elif file_ext == "png":
                mime_type = "image/png"
            elif file_ext == "mp4":
                mime_type = "video/mp4"
            elif file_ext == "mp3":
                mime_type = "audio/mpeg"
            elif file_ext == "wav":
                mime_type = "audio/wav"
            else:
                mime_type = "application/octet-stream"
            # Tạo Part từ bytes
            part = types.Part.from_bytes(data=file_bytes, mime_type=mime_type)
            contents.append(part)

if st.button("Gửi"):
    if not contents:
        st.warning("Vui lòng nhập nội dung chat hoặc upload file.")
    else:
        with st.spinner("Đang xử lý yêu cầu..."):
            try:
                # Gọi API generate_content theo cách mẫu trong tài liệu
                response = client.models.generate_content(
                    model='gemini-2.5-flash-preview-04-17',
                    contents=contents
                )
                st.write("**Phản hồi từ Gemini:**")
                st.write(response.text)
            except Exception as e:
                st.error(f"Có lỗi xảy ra: {e}")