import streamlit as st
import os
import json
from zipfile import ZipFile
from io import BytesIO

# Initialized Modules
from modules.image_process import image_process
from modules.audio_process import audio_process
from modules.transformation import json_to_dms_excel
# from modules.saving import saving_json, saving_excel


# Streamlit app
st.title("TOOL NHẬP ĐƠN HÀNG TỰ ĐỘNG")

# Add a text input for employee code
employee_code = st.text_input("Nhập mã nhân viên của bạn:", placeholder="Administrator", max_chars=20)

# Add a button to confirm the input
if st.button("Xác nhận mã nhân viên"):
    if employee_code == "":
        st.error("Vui lòng nhập mã nhân viên trước khi xác nhận.")
    else:
        st.success(f"Mã nhân viên của bạn là: {employee_code}")

st.write("Tải lên folder chứa hình ảnh đơn hàng (dưới dạng file Zip).")

# File uploader for ZIP file
uploaded_file = st.file_uploader("Tải file Zip chứa hình ảnh đơn hàng", type=["zip"])

if uploaded_file is not None:
    # Extract ZIP file
    with ZipFile(uploaded_file) as zip_file:
        zip_file.extractall("temp_invoices")

        # Collect image and audio files
        all_files = zip_file.namelist()
        # Image files
        image_files = [os.path.join("temp_invoices", file) for file in all_files if file.lower().endswith((".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"))]
        # Audio files
        audio_files = [os.path.join("temp_invoices", file) for file in all_files if file.lower().endswith((".mp3", ".wav", ".m4a", ".flac"))]

    if not image_files and not audio_files:
        st.error("Không có file hình ảnh và/hoặc file âm thanh hợp lệ trong thư mục file Zip.")
    else:
        st.write(f"Tìm thấy file {len(image_files)} ảnh và {len(audio_files)} file âm thanh trong thư mục.")

        order_counter = 1 # Initialize the order ID counter
        results = {}

        # Process image files
        for i, image_path in enumerate(image_files, start=1):
            st.write(f"Processing image{i}/{len(image_files)}: {os.path.basename(image_path)}")
            try:
                content = image_process(image_path= image_path, order_id= order_counter)
                results[f"image_file_{i}"] = content
                order_counter += 1
            except Exception as e:
                st.error(f"Error processing {image_path}: {e}")

        # Process audio files
        for i, audio_path in enumerate(audio_files, start=1):
            st.write(f"Processing audio{i}/{len(audio_files)}: {os.path.basename(audio_path)}")
            try:
                content = audio_process(audio_path= audio_path, order_id= order_counter)
                results[f"audio_file_{i}"] = content
                order_counter += 1
            except Exception as e:
                st.error(f"Error processing {audio_path}: {e}")

        # Display the results
        st.write("OCR Results:")
        st.json(results)

        # Allow user to download the results as a JSON file
        json_data = json.dumps(results, indent=4)
        st.download_button(
            label="Tải file kết quả ở dạng JSON",
            data=json_data,
            file_name="invoice_results.json",
            mime="application/json"
        )
        # Allow user to download the results as Excel file
        # Convert results to Excel format
        df = json_to_dms_excel(results, employee_code)

        # Allow user to preview excel file before downloading
        st.subheader("Xem trước dữ liệu Excel:")
        st.dataframe(df.head())

        excel_buffer = BytesIO()
        df.to_excel(excel_buffer, index=False)
        excel_buffer.seek(0)
        st.download_button(
            label="Tải file theo template Excel từ DMS",
            data=excel_buffer,
            file_name="invoice_results.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        )