from pathlib import Path

import numpy as np
import streamlit as st
from fastcore.xtras import globtastic

from annotator.utils import (
    annotate,
    df_from_result,
    find_word_timestamp,
    generate_srt,
    get_audio,
    get_v_from_url,
    working_directory,
    write_srt,
)

st.set_page_config(layout="wide")


SRT_PATH = Path("srt")
if not SRT_PATH.exists():
    SRT_PATH.mkdir()

AUDIO_PATH = Path("./audio")
if not AUDIO_PATH.exists():
    AUDIO_PATH.mkdir()


def make_sidebar():
    with st.sidebar:
        st.markdown("## yt-video-annotator")
        st.write("Link to the GitHub repo")


@st.cache_resource
def caption_from_url(url):
    audio_src = get_audio(url)
    v = get_v_from_url(url)
    audio_src = globtastic(AUDIO_PATH, file_glob="*.mp3", file_re=v)[0]
    result = annotate(audio_src)
    df = df_from_result(result)
    return audio_src, df


def main():
    url, name = None, None
    make_sidebar()
    place = "https://www.youtube.com/watch?v=C0DPdy98e4c&ab_channel=SimonYapp"
    col1, col2 = st.columns([1.2, 1])
    with col1:
        url = st.text_input("Enter URL for the YT video", place)
        st.video(url)

    with col2:
        default_opt = "Search for words"
        opt = st.radio(
            "What do you wish to do?",
            [default_opt, "Generate subtitles for the entire video"],
        )
        if opt == default_opt:
            st.markdown("### Search for words in the video")
            words = st.text_input("Enter words separated by a comma")
            words = words.split(",")

            if st.button("Get Timestamps"):
                audio_src, df = caption_from_url(url)
                times = find_word_timestamp(df, *words)
                times = np.asarray(times).reshape(len(words), -1)
                # st.write(times)
                for i, word in enumerate(words):
                    st.write(f"{word} is said on {times[i].flatten()} timestamp")

        else:
            if st.button("Generate SRT"):
                audio_src, df = caption_from_url(url)
                name = Path(audio_src).stem
                s = generate_srt(df)
                with working_directory(SRT_PATH):
                    write_srt(s, name)

        if name is not None:
            with working_directory(SRT_PATH):
                key = get_v_from_url(url)
                srt = globtastic(".", file_glob="*.srt", file_re=key)[0]
                with open(srt) as f:
                    st.download_button("Download SRT", f, file_name=f"{name}.srt")

    # subprocess.run(['rm', '-rf', 'audio'])
    # subprocess.run(['rm', '-rf', 'srt'])


if __name__ == "__main__":
    main()