import streamlit as st
import os

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

# Azure Blob Storage connection strings

# Container names
FIRST_CONTAINER_NAME = "source"
SECOND_CONTAINER_NAME = "target"

# Initialize Azure Blob Service Clients
blob_service_client = BlobServiceClient.from_connection_string(os.environ["AZURE_STORAGE_CONNECTION_STRING"])

# Function to upload file to Azure Storage
def upload_to_azure(blob_service_client, container_name, file, file_name):
    container_client = blob_service_client.get_container_client(container_name)
    container_client.upload_blob(name=file_name, data=file, overwrite=True)

# Function to download file from Azure Storage
def download_from_azure(blob_service_client, container_name, file_name):
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(blob=file_name)
    file_content = blob_client.download_blob().readall()
    return file_content

# Function to delete file from Azure Storage
def delete_from_azure(blob_service_client, container_name, file_name):
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(blob=file_name)
    blob_client.delete_blob()

key = os.environ["AZURE_AI_TRANSLATOR_KEY"]
endpoint = os.environ["AZURE_AI_ENDPOINT_URL"]

sourceUri = "https://cbdtranslation.blob.core.windows.net/source"
targetUri = "https://cbdtranslation.blob.core.windows.net/target"

# Initialize a new instance of the DocumentTranslationClient
client = DocumentTranslationClient(endpoint, AzureKeyCredential(key))

def translate(lang_id, lang_name):
    poller = client.begin_translation(sourceUri, targetUri, lang_id)
    result = poller.result()
    st.write(
        'Total documents: {}'.format(
            poller.details.documents_total_count
        )
    )

    for document in result:
        if document.status == 'Succeeded':
            st.success('Translated to language: {}\n'.format(lang_name))
        else:
            st.error(
                'Error Code: {}, Message: {}\n'.format(
                    document.error.code, document.error.message
                )
            )

# Streamlit UI
st.title("Azure Translation Tools")

# Step 1: Upload File
uploaded_file = st.file_uploader("Upload a file to start the process")

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_file:
    submit = st.button("Get Result", key='submit')

if uploaded_file and submit:
    file_name = uploaded_file.name
    file_content = uploaded_file.read()

    upload_to_azure(blob_service_client, FIRST_CONTAINER_NAME, file_content, file_name)

    st.info("Performing translations on the file...")
    translate(lang_id, lang_name)

    # Upload file to second storage for simulation
    downloaded_file_content = download_from_azure(blob_service_client, SECOND_CONTAINER_NAME, file_name)

    # Step 5: Delete both files
    delete_from_azure(blob_service_client, FIRST_CONTAINER_NAME, file_name)
    delete_from_azure(blob_service_client, SECOND_CONTAINER_NAME, file_name)

    # Allow user to download the file
    st.download_button(
        label="Download the Processed File",
        data=downloaded_file_content,
        file_name=f"{lang_name}-translated-{file_name}",
        mime="application/octet-stream"
    )