import gradio as gr
import json
from difflib import Differ
import ffmpeg
import os
from pathlib import Path
import time
import aiohttp
import asyncio
from transformers import pipeline

# Set true if you're using huggingface inference API API https://huggingface.co/inference-api
API_BACKEND = True
MODEL = "facebook/wav2vec2-base-960h"

if API_BACKEND:
    from dotenv import load_dotenv
    import base64
    import asyncio
    load_dotenv(Path(".env"))

    HF_TOKEN = os.environ["HF_TOKEN"] = ""
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    API_URL = f'https://api-inference.huggingface.co/models/{MODEL}'

else:
    import torch
    from transformers import pipeline

    # is cuda available?
    cuda = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
    device = 0 if torch.cuda.is_available() else -1
    speech_recognizer = pipeline(
        task="automatic-speech-recognition",
        model=MODEL,
        tokenizer=MODEL,
        framework="pt",
        device=device,
    )

videos_out_path = Path("./videos_out")
videos_out_path.mkdir(parents=True, exist_ok=True)

samples_data = sorted(Path('examples').glob('*.json'))
SAMPLES = []
for file in samples_data:
    with open(file) as f:
        sample = json.load(f)
    SAMPLES.append(sample)
VIDEOS = list(map(lambda x: [x['video']], SAMPLES))

total_inferences_since_reboot = 0
total_cuts_since_reboot = 0


async def speech_to_text(video_file_path):
    """
    Takes a video path to convert to audio, transcribe audio channel to text and char timestamps

    Using https://huggingface.co/tasks/automatic-speech-recognition pipeline
    """
    global total_inferences_since_reboot
    if (video_file_path == None):
        raise ValueError("Error no video input")

    video_path = Path(video_file_path)
    try:
        # convert video to audio 16k using PIPE to audio_memory
        audio_memory, _ = ffmpeg.input(video_path).output(
            '-', format="wav", ac=1, ar='16k').overwrite_output().global_args('-loglevel', 'quiet').run(capture_stdout=True)
    except Exception as e:
        raise RuntimeError("Error converting video to audio")

    ping("speech_to_text")
    last_time = time.time()
    if API_BACKEND:
        # Using Inference API https://huggingface.co/inference-api
        # try twice, because the model must be loaded
        for i in range(10):
            for tries in range(4):
                print(f'Transcribing from API attempt {tries}')
                try:
                    inference_reponse = await query_api(audio_memory)
                    print(inference_reponse)
                    transcription = inference_reponse["text"].lower()
                    timestamps = [[chunk["text"].lower(), chunk["timestamp"][0], chunk["timestamp"][1]]
                                  for chunk in inference_reponse['chunks']]

                    total_inferences_since_reboot += 1
                    print("\n\ntotal_inferences_since_reboot: ",
                          total_inferences_since_reboot, "\n\n")
                    return (transcription, transcription, timestamps)
                except Exception as e:
                    print(e)
                    if 'error' in inference_reponse and 'estimated_time' in inference_reponse:
                        wait_time = inference_reponse['estimated_time']
                        print("Waiting for model to load....", wait_time)
                        # wait for loading model
                        # 5 seconds plus for certainty
                        await asyncio.sleep(wait_time + 5.0)
                    elif 'error' in inference_reponse:
                        raise RuntimeError("Error Fetching API",
                                           inference_reponse['error'])
                    else:
                        break
            else:
                raise RuntimeError(inference_reponse, "Error Fetching API")
    else:
        try:
            print(f'Transcribing via local model')
            output = speech_recognizer(
                audio_memory, return_timestamps="char", chunk_length_s=10, stride_length_s=(4, 2))

            transcription = output["text"].lower()
            timestamps = [[chunk["text"].lower(), chunk["timestamp"][0].tolist(), chunk["timestamp"][1].tolist()]
                          for chunk in output['chunks']]
            total_inferences_since_reboot += 1

            print("\n\ntotal_inferences_since_reboot: ",
                  total_inferences_since_reboot, "\n\n")
            return (transcription, transcription, timestamps)
        except Exception as e:
            raise RuntimeError("Error Running inference with local model", e)


async def cut_timestamps_to_video(video_in, transcription, text_in, timestamps, words_to_cut):
    """
    Given original video input, text transcript + timestamps,
    and edit text, cuts video segments into a single video
    """
    global total_cuts_since_reboot

    video_path = Path(video_in)
    video_file_name = video_path.stem
    if (video_in == None or text_in == None or transcription == None):
        raise ValueError("Inputs undefined")

    d = Differ()
    # compare original transcription with edit text
    diff_chars = d.compare(transcription, text_in)

    # Include additions in the filtered list
    filtered = list(filter(lambda x: x[0] != '-' and x[0] != '+', diff_chars))

    # Update grouping logic to handle additions and word cuts
    idx = 0
    grouped = {}
    word_cuts = []
    for (a, b) in zip(filtered, timestamps):
        if a[0] != '-':
            if idx in grouped:
                grouped[idx].append(b)
            else:
                grouped[idx] = []
                grouped[idx].append(b)
        elif a[0] == '-':
            idx += 1
        elif a[0] == '+':
            word_cuts.append(b)

    # after grouping, gets the lower and upper start and time for each group
    timestamps_to_cut = [[v[0][1], v[-1][2]] for v in grouped.values()]

    # Add word cut timestamps
    timestamps_to_cut.extend(word_cuts)

    between_str = '+'.join(
        map(lambda t: f'between(t,{t[0]},{t[1]})', timestamps_to_cut))

    if timestamps_to_cut:
        video_file = ffmpeg.input(video_in)
        video = video_file.video.filter(
            "select", f'({between_str})').filter("setpts", "N/FRAME_RATE/TB")
        audio = video_file.audio.filter(
            "aselect", f'({between_str})').filter("asetpts", "N/SR/TB")

        output_video = f'./videos_out/{video_file_name}.mp4'
        ffmpeg.concat(video, audio, v=1, a=1).output(
            output_video).overwrite_output().global_args('-loglevel', 'quiet').run()
    else:
        output_video = video_in

    tokens = [(token[2:], token[0] if token[0] != " " else None)
              for token in filtered]

    total_cuts_since_reboot += 1
    ping("video_cuts")
    print("\n\ntotal_cuts_since_reboot: ", total_cuts_since_reboot, "\n\n")
    return (tokens, output_video)


async def query_api(audio_bytes: bytes):
    """
    Query for Huggingface Inference API for Automatic Speech Recognition task
    """
    payload = json.dumps({
        "inputs": base64.b64encode(audio_bytes).decode("utf-8"),
        "parameters": {
            "return_timestamps": "char",
            "chunk_length_s": 10,
            "stride_length_s": [4, 2]
        },
        "options": {"use_gpu": False}
    }).encode("utf-8")
    async with aiohttp.ClientSession() as session:
        async with session.post(API_URL, headers=headers, data=payload) as response:
            print("API Response: ", response.status)
            if response.headers['Content-Type'] == 'application/json':
                return await response.json()
            elif response.headers['Content-Type'] == 'application/octet-stream':
                return await response.read()
            elif response.headers['Content-Type'] == 'text/plain':
                return await response.text()
            else:
                raise RuntimeError("Error Fetching API")


def ping(name):
    url = f'https://huggingface.co/api/telemetry/spaces/radames/edit-video-by-editing-text/{name}'
    print("ping: ", url)

    async def req():
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                print("pong: ", response.status)
    asyncio.create_task(req())


def find_words_to_cut(transcription, words):
    """
    Find timestamps of words to cut in the transcription
    """
    timestamps_to_cut = []
    for word in words:
        word_lower = word.lower()
        word_positions = [i for i, x in enumerate(transcription.split()) if x.lower() == word_lower]
        for pos in word_positions:
            timestamps_to_cut.append([transcription[pos:].find(' '), pos + 1])

    return timestamps_to_cut


# ---- Gradio Layout -----
video_in = gr.Video(label="Video file", elem_id="video-container")
text_in = gr.Textbox(label="Transcription", lines=10, interactive=True)
words_to_cut_input = gr.Textbox(label="Words to Cut (comma-separated)", lines=1, default="")
video_out = gr.Video(label="Video Out")
diff_out = gr.HighlightedText(label="Cuts Diffs", combine_adjacent=True)
examples = gr.Dataset(components=[video_in], samples=VIDEOS, type="index")

css = """
#cut_btn, #reset_btn { align-self:stretch; }
#\\31 3 { max-width: 540px; }
.output-markdown {max-width: 65ch !important;}
#video-container{
    max-width: 40rem;
}
"""
with gr.Blocks(css=css) as demo:
    transcription_var = gr.State()
    timestamps_var = gr.State()
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            # Edit Video By Editing Text
            This project is a quick proof of concept of a simple video editor where the edits
            are made by editing the audio transcription.
            Using the [Huggingface Automatic Speech Recognition Pipeline](https://huggingface.co/tasks/automatic-speech-recognition)
            with a fine-tuned model, you can predict the text transcription and character or word-based timestamps.
            """)

    with gr.Row():
        examples.render()

        def load_example(id):
            video = SAMPLES[id]['video']
            transcription = SAMPLES[id]['transcription'].lower()
            timestamps = SAMPLES[id]['timestamps']

            return (video, transcription, transcription, timestamps)

        examples.click(
            load_example,
            inputs=[examples],
            outputs=[video_in, text_in, transcription_var, timestamps_var],
            queue=False)

    with gr.Row():
        with gr.Column():
            video_in.render()
            transcribe_btn = gr.Button("Transcribe Audio")
            transcribe_btn.click(speech_to_text, [video_in], [
                text_in, transcription_var, timestamps_var])

    with gr.Row():
        gr.Markdown("""
        ### Now edit as text
        After running the video transcription, you can make cuts to the text below (only cuts, not additions!)
        You can also specify words to cut in the "Words to Cut" input box.
        """)

    with gr.Row():
        with gr.Column():
            text_in.render()
            words_to_cut_input.render()
            with gr.Row():
                cut_btn = gr.Button("Cut to video", elem_id="cut_btn")
                # send audio path and hidden variables
                cut_btn.click(cut_timestamps_to_video, [
                    video_in, transcription_var, text_in, timestamps_var, words_to_cut_input], [diff_out, video_out])

                reset_transcription = gr.Button(
                    "Reset to last transcription", elem_id="reset_btn")
                reset_transcription.click(
                    lambda x: x, transcription_var, text_in)
        with gr.Column():
            video_out.render()
            diff_out.render()

    with gr.Row():
        gr.Markdown("""
        #### Video Credits

        1. [Cooking](https://vimeo.com/573792389)
        2. [Shia LaBeouf "Just Do It"](https://www.youtube.com/watch?v=n2lTxIk_Dr0)
        3. [Mark Zuckerberg & Yuval Noah Harari in Conversation](https://www.youtube.com/watch?v=Boj9eD0Wug8)
        """)

demo.launch(debug=True)