from flask import jsonify, send_file, request
from main import *

def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX', output_path="output_translation.txt"):
    if translation_model is None:
        return "Translation model not initialized."

    encoded_text = translation_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
    generated_tokens = translation_model.generate(input_ids=encoded_text['input_ids'], attention_mask=encoded_text['attention_mask'], forced_bos_token_id=translation_model.config.lang_code_to_id[target_language_code])
    translation = translation_model.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

    with open(output_path, "w") as file:
        file.write(translation)
    return output_path

def translation_api():
    data = request.get_json()
    text = data.get('text')
    target_lang = data.get('target_lang', 'es')
    source_lang = data.get('source_lang', 'en')
    if not text:
        return jsonify({"error": "Text is required"}), 400
    output_file = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
    if output_file == "Translation model not initialized.":
        return jsonify({"error": "Translation failed"}), 500
    return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output_translation.txt")