import os
import re
import json
import time
import requests
import gradio as gr
from google.auth.transport.requests import Request
from google.oauth2 import service_account


def get_token(function_url):
    cf_key = json.loads(os.environ.get("CF_KEY"))
    credentials = service_account.IDTokenCredentials.from_service_account_info(
        cf_key, target_audience=function_url
    )
    if credentials.valid is False:
        credentials.refresh(Request())
    token = credentials.token
    return token


def dubpro_english_transliteration(text):
    headers = {
        "Accept": "*/*",
        "Authorization": os.environ.get('DUBPRO_API_HEADER_KEY'),
        "Content-Type": "application/json"
    }
    payload = {
        "autoTranslation": True,
        "text": text
    }
    response = requests.post(f"https://api.dev.video.dubpro.ai/video-service/videos/re_translation", headers=headers, json=payload)
    result = response.json()
    return result["data"]["transliteratedText"]


def generate_rephrases(text, min_words, max_words):
    function_url = os.environ.get("REPHRASER_URL")
    assert function_url is not None
    TOKEN = get_token(function_url)
    
    headers = {
        "content-type": "application/json",
        "accept": "*/*",
        "Authorization": f"Bearer {TOKEN}"
    }

    retry = 0
    max_retries = 3

    payload = {
        "text": text,
        "min_num_words": min_words,
        "max_num_words": max_words
    }

    while retry < max_retries:
        try:
            response = requests.post(
                function_url,
                headers=headers,
                json=payload,
            )
            try:
                rephrased_text = response.json()["rephrased_text"]
            except:
                rephrased_text = response.text
            if rephrased_text == "ERROR":
                raise Exception("Error showed.")
            if response.status_code != 500:
                return rephrased_text, len(rephrased_text.split())
            else:
                time.sleep(5)
                retry += 1
        except:
            print(response.text)
            time.sleep(5)
            retry += 1
    return payload["text"], len(payload["text"].split())


with gr.Blocks() as demo:
    gr.Markdown("# Translator Assistance Tools")
    with gr.Tab("Rephraser Tool"):
        with gr.Row():
            rephrase_text = gr.Textbox(label="Input text", info="Please enter text.")
            with gr.Column():
                current_words_num = gr.Number(label="Current number of words", interactive=False)
                min_words_num = gr.Number(label="Minimum number of words")
                max_words_num = gr.Number(label="Maximum number of words")
        rephrase = gr.Button("Submit")
        with gr.Row():
            rephrased_text = gr.Textbox(label="Output text")
            word_count = gr.Textbox(label="Word count")
        rephrase.click(generate_rephrases, [rephrase_text, min_words_num, max_words_num], [rephrased_text, word_count])

    with gr.Tab("Transliteration"):
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label="Input text", info="Please enter English text.")
            output_text = gr.Textbox(label="Output text")
        transliterate = gr.Button("Submit")
        transliterate.click(dubpro_english_transliteration, input_text, output_text)

    rephrase_text.change(lambda x: len(x.split()), rephrase_text, current_words_num)
        
if __name__=="__main__":
    demo.launch(auth=(os.environ.get("USERNAME"), os.environ.get("PASSWORD")))