import gradio
import gradio as gr

from download import download_generator
import anvil.media
import os
import dotenv
from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE

from utils.apis import render_api_elements
from utils.utils import get_args

dotenv.load_dotenv()

anvil.server.connect(os.environ.get('ANVIL_UPLINK_KEY'))
queue_placeholder = None

args = get_args()
gradio_share: bool = args.get("public")
model_size: str = args.get("model")
preload_model: str = args.get("preload")


LANG_CHOICES = sorted([x.capitalize() for x in LANGUAGES.values()])
LANG_CHOICES.insert(0, "Autodetect")


url_input = gr.Textbox(label="Youtube/Twitter/etc video URL (supports many services)", value='https://twitter.com/starsonxh/status/1552945347194142720', lines=1, elem_id="url_input")
# download_status = gr.Textbox(label="Status:", value='', lines=1, elem_id="download_status")
download_status = gr.Checkbox(label="", elem_id="download_status", interactive=False)
translate_action = gr.Checkbox(label="Auto translate to english", elem_id='translate_toggle', interactive=True, value=True)
init_video = gr.Video(label="Upload video manually", visible=True, interactive=True, mirror_webcam=False)
init_audio = gr.Audio(label="Downloaded audio", visible=False)
output_text = gr.Textbox(label="Output text", lines=5, visible=False, max_lines=10, interactive=True, elem_id="output_text")
sub_video = gr.Video(label="Subbed video", visible=False, mirror_webcam=False)


def predownload(url, translate_action, source_language):
  files = []
  for response in download_generator(url, translate_action, source_language):
    updates_object = {}
    updates_object[download_status] = gr.update(label=f"{response.get('message')}")
    meta = response.get('meta')

    if 'video' in response:
      updates_object[init_video] = gr.update(visible=True, value=response["video"],
                                             label=f"Init Video: {meta['id']}.{meta['ext']}")
      updates_object[init_audio] = gr.update(visible=True, value=response["audio"],
                                             label=f"Extracted audio : {meta['id']}.mp3")
      files.append(response["video"])
      files.append(response["audio"])
    if 'whisper_result' in response:
      updates_object[output_text] = gr.update(value=response['whisper_result'].get('srt'), visible=True,
                                              label=f"Subtitles transcribed from {response['whisper_result'].get('language')} (detected language)")
    if 'srt_path' in response:
      files.append(response["srt_path"])

    if 'sub_video' in response:
      updates_object[sub_video] = gr.update(visible=True, value=response["sub_video"],
                                            label=f"Subbed video: {meta['id']}_translated.mp4")
      files.append(response["sub_video"])

    updates_object[output_file] = gr.update(value=files, visible=len(files) > 0, label=f"Output Files")
    yield updates_object

def correct_subtitles(url, output_text):
  for response in download_generator(url, corrected_subtitles=output_text):
    updates_object = {}
    updates_object[download_status] = gr.update(label=f"STATUS: {response.get('message')}")
    if 'sub_video' in response:
      updates_object[sub_video] = gr.update(visible=True, value=response["sub_video"],
                                            label=f"Corrected subtitles")
    yield updates_object


subtitled_video = False

with gr.Blocks(css='@import "file=static/css/main.css";', theme='darkpeach', title='Vid Translator Studio') as demo:
  gr.HTML('<h1 class="main-title">VidTranslator Studio 0.1</h1>')
  gr.HTML("<h2 class='secondary'>Automatic social media video translation from 99 languages</h2>")

  with gr.Row(elem_id="input_row"):
    with gr.Group() as group:
      url_input.render()
      action_btn = gr.Button(elem_id='submit', variant='primary', value="Translate")
      gr.StatusTracker()
      with gr.Row(elem_id="second_row"):
        source_language = gr.Dropdown(choices=LANG_CHOICES,
                                      label="Source Language",
                                      value='Autodetect',
                                      interactive=True, elem_id="source_language")
        download_status.render()
        translate_action.render()

  with gr.Row():
    with gr.Column():
      init_video.render()
      init_audio.render()
      with gr.Row():
          output_file = gr.Files(label="Output Files", visible=False)

    with gr.Column():
      output_text.render()
      correct_btn = gr.Button("Correct subtitles")

    with gr.Column():
      sub_video.render()


  outputs = [download_status, init_video, init_audio, output_text, sub_video, output_file ]
  inputs = [url_input, translate_action, source_language]
  action_btn.click(fn=predownload, inputs=inputs, outputs=outputs, api_name='predownload')
  url_input.submit(fn=predownload, inputs=inputs, outputs=outputs)

  correct_btn.click(fn=correct_subtitles, inputs=[url_input, output_text], outputs=[download_status, output_text, sub_video])

  translate_action.change(fn=lambda x: {action_btn: gr.update(value=f"Translate" if x else "Transcribe")},
                          inputs=[translate_action], outputs=[action_btn])

  gr.HTML("""<div class='footer'>
    <div class="relative">
  <div class="absolute inset-0 flex items-center" aria-hidden="true">
    <div class="w-full border-t border-gray-300"></div>
  </div>
  <div class="relative flex justify-center">
    <span class="bg-white px-2 text-sm text-gray-500">Continue</span>
  </div>
</div>

     <a href='https://twitter.com/altryne/'><img src='https://img.shields.io/badge/%40altryne-follow-green' alt=''></a>
     </div>""")

  def init_video_manual_upload(url, init_video):
    print(url)
    print(init_video)

  init_video.change(fn=init_video_manual_upload, inputs=[url_input, init_video], outputs=[])

  # Render imported buttons for API bindings
  render_api_elements(url_input,download_status, output_text, sub_video, output_file)

queue_placeholder = demo.queue()


@anvil.server.callable
def temp():
  return 'temp worked'

if __name__ == "__main__":
  gradio.close_all()
  port = os.environ.get('SERVER_PORT', 8111)
  demo.launch(show_error=True, debug=True, share=gradio_share,server_port=int(port), favicon_path='fonts/icon.png')