import streamlit as st
import google.generativeai as genai
from db.mongodb import save_file_to_mongodb, fetch_file_from_mongodb, detect_file_type
from powerpoint.xml_handling import (
    extract_text_from_xml, update_xml_with_translated_text_mongodb, ppt_to_xml_mongodb
)
from translate.translator import translate_text_dict
from powerpoint.pptx_object import create_translated_ppt
from excel.excel_translate import translate_xlsx, translate_csv
from word.word_translate import translate_docx_from_mongodb
import dotenv
import os

dotenv.load_dotenv(".env")

genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

st.title("Translate Your File Easily! 🌍")

uploaded_file = st.file_uploader("📂 Chọn file để dịch")
target_lang = st.selectbox("🌐 Chọn ngôn ngữ", ["english", "vietnamese"])

def process_file(file, file_type):
    progress_bar = st.progress(0)
    file_id = save_file_to_mongodb(uploaded_file=file, db_name=file_type.lower(), collection_name="root_file")
    progress_bar.progress(20)
    st.write(f"📂 File ID: {file_id}")

    if file_type == "PPTX":
        xml_file_id = ppt_to_xml_mongodb(file_id)
        progress_bar.progress(40)
        text_dict = extract_text_from_xml(file_id=xml_file_id)
        translated_dict = translate_text_dict(text_dict, target_lang=target_lang)
        progress_bar.progress(60)
        final_xml_id = update_xml_with_translated_text_mongodb(xml_file_id, translated_dict)
        final_id = create_translated_ppt("pptx", file_id, final_xml_id, "final_file")
    elif file_type == "Excel":
        final_id = translate_xlsx(file_id = file_id, target_lang = target_lang)
    elif file_type == "CSV":
        final_id = translate_csv(file_id = file_id, target_lang = target_lang)
    elif file_type == "Word":
        final_id = translate_docx_from_mongodb(file_id, target_lang)
    else:
        st.error("❌ Loại file không hỗ trợ!")
        return

    progress_bar.progress(80)
    st.write("✅ File đã được dịch xong!")
    file_io, file_name = fetch_file_from_mongodb(file_type.lower(), "final_file", final_id)
    progress_bar.progress(100)

    if file_io:
        st.download_button("⬇️ Tải file về", data=file_io.getvalue(), file_name=file_name)
    else:
        st.error("❌ Không thể tải xuống file. Vui lòng thử lại!")

if uploaded_file and st.button("🚀 Upload và dịch ngay!"):
    file_type = detect_file_type(uploaded_file)
    st.write(f"🔍 Loại file phát hiện: {file_type}")
    process_file(uploaded_file, file_type)