import uuid
import gradio as gr
import boto3
import os
from PIL import Image

hf_token = os.environ.get("HF_TOKEN")
tasomi_model = os.environ.get("TASOMI_MODEL")
s3_bucket = os.environ.get("IMCAPSULE_EXTERNAL_BUCKET")

m=gr.load(tasomi_model, src="models",hf_token=hf_token)

def detect_language(text):
    # Initialize the Comprehend client
    comprehend = boto3.client('comprehend')
    # Call DetectDominantLanguage API
    response = comprehend.detect_dominant_language(Text=text)
    # Extract the detected language
    detected_language = response['Languages'][0]['LanguageCode']
    return detected_language

def upload_file(file_path):
    s3 = boto3.client('s3')
    key='public/' + os.path.basename(file_path)
    # Upload the file
    s3.upload_file(file_path, s3_bucket, key)
    s3_path = f"{s3.meta.endpoint_url}/{s3_bucket}/{key}"
    print(f'File uploaded successfully. s3_path= {s3_path} ')
    return s3_path

def translate_and_predict(text):
    """Translates text and generates an image using Stable Diffusion."""
    source_language_code=detect_language(text)
    # Translate input text
    translate = boto3.client('translate')
    target_language_code="en"
    print("source_language_code:" , source_language_code)
    translation_result = translate.translate_text(Text=text, SourceLanguageCode=source_language_code, TargetLanguageCode=target_language_code)
    translated_text = translation_result.get('TranslatedText')
    print("translated_text:" , translated_text)
    gen_image=m(translated_text)
    webp_image = Image.open(gen_image)
    converted_image_path = f"./tasomi-{uuid.uuid4().hex}.jpg"
    webp_image.save(converted_image_path, "JPEG")
    s3_path = upload_file(converted_image_path)
    os.remove(converted_image_path)
    return s3_path,s3_path

# Define the Gradio interface
iface = gr.Interface(
    fn=translate_and_predict,  # Since we're only displaying an image, there's no function to call
    inputs=["text"],  # Input: text, source language, target language
    outputs=["image","text"],  # Output: generated image
    title="Tasomi Images",
)

# Launch the interface
iface.launch()