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

def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX'):
    if translation_model is None: return {"error": "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]; return {"translated_text": translation}

def translation_api(text):
    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 = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
    if "error" in output: return jsonify({"error": output["error"]}), 500
    return output