import os
from os.path import basename
from zipfile import ZipFile, ZIP_DEFLATED
from shutil import rmtree

import sentry_sdk
import sphn
import gradio as gr
from inference import inference_file

if os.environ["SENTRY_DSN"]:
    sentry_sdk.init(
        dsn=os.environ["SENTRY_DSN"],
        send_default_pii=True,
    )
    print("Sentry SDK is activated")


def extract_chunks(file, min_sec, max_sec):
    archive_name = "tmp.zip"
    n_files = 0
    duration_secs = 0

    with ZipFile(
        archive_name, "w", compression=ZIP_DEFLATED, allowZip64=True, compresslevel=9
    ) as zip_file:
        results = inference_file(file)

        filenames = [it["filename"] for it in results]
        durations = sphn.durations(filenames)

        for idx, result in enumerate(results):
            duration = durations[idx]

            print(result, duration)

            if duration <= min_sec or duration >= max_sec:
                print("Skipping...")
                continue

            arc_name = basename(result["filename"])
            zip_file.write(result["filename"], arc_name)

            duration_secs += duration
            n_files += 1

    # Remove files
    rmtree("chunks")

    mins = round(duration_secs / 60, 4)

    gr.Success(
        f"VAD model identified {n_files} files in interval [{min_sec}:{max_sec}], total duration = {mins} min."
    )

    return archive_name


demo = gr.Interface(
    title="MarbleNet",
    fn=extract_chunks,
    inputs=[
        gr.File(label="WAV file to process", file_count="single", file_types=[".wav"]),
        gr.Number(label="Minimum seconds", value=0.1, minimum=0.01, maximum=59.99),
        gr.Number(label="Maximum seconds", value=30, minimum=0.02, maximum=60),
    ],
    outputs=[gr.File(label="ZIP file with voice chunks")],
    submit_btn="Inference",
)
demo.launch()