import io
import os
import streamlit as st
import requests
import zipfile

from azure.storage.blob import BlobServiceClient
from azure.core.credentials import AzureKeyCredential
from azure.ai.translation.document import DocumentTranslationClient
from dotenv import load_dotenv

load_dotenv()

# Streamlit UI
st.title("Azure Translation Tools")
uploaded_files = st.file_uploader("Upload files to start the process", accept_multiple_files=True)

# Language options
langs = (
    'id - Indonesian',
    'en - English',
    'es - Spanish',
    'zh - Chinese',
    'ar - Arabic',
    'fr - French',
    'ru - Russian',
    'hi - Hindi',
    'pt - Portuguese',
    'de - German',
    'ms - Malay',
    'ta - Tamil',
    'ko - Korean',
    'th - Thai',
)

lang = st.selectbox('Target language selection:', langs, key='lang')
lang_id = lang.split()[0]
lang_name = lang.split()[-1]

if uploaded_files:
    submit = st.button("Get Result", key='submit')

if uploaded_files and submit:

    # Create a zip file in memory
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
        # Add progress bar
        progress_bar = st.progress(0)
        for idx, uploaded_file in enumerate(uploaded_files):
            file_name = uploaded_file.name
            file_content = uploaded_file.read()

            headers = {
                "Ocp-Apim-Subscription-Key": os.environ["AZURE_AI_TRANSLATOR_KEY"],
            }

            files = {
                "document": (file_name, file_content, "ContentType/file-extension"),
            }

            url = f"{os.environ["AZURE_AI_ENDPOINT_URL"]}/translator/document:translate?targetLanguage={lang_id}&api-version={os.environ["AZURE_AI_API_VERSION"]}"
            response = requests.post(url, headers=headers, files=files)
            
            if response.status_code == 200:
                # Add translated file to zip
                zip_file.writestr(f"{lang_name}-translated-{file_name}", response.content)
                st.success(f"Successfully translated: {file_name}")
            else:
                st.error(f"Failed to translate {file_name} with status code {response.status_code}: {response.text}")
            
            # Update progress bar
            progress = (idx + 1) / len(uploaded_files)
            progress_bar.progress(progress)

    # Allow user to download all files as zip
    st.download_button(
        label="Download All Translated Files",
        data=zip_buffer.getvalue(),
        file_name=f"{lang_name}-translated-files.zip",
        mime="application/zip"
    )