#!/usr/bin/env python
# -*- coding: utf-8 -*-

from tempfile import NamedTemporaryFile

import streamlit as st

from conette import CoNeTTEModel, conette


@st.cache_resource
def load_conette(*args, **kwargs) -> CoNeTTEModel:
    return conette(*args, **kwargs)


def main() -> None:
    st.header("CoNeTTE model test")
    audios = st.file_uploader(
        "Upload an audio file",
        type=["wav", "flac", "mp3", "ogg", "avi"],
        accept_multiple_files=True,
    )

    if audios is not None and len(audios) > 0:
        model = load_conette(model_kwds=dict(device="cpu"))

        for audio in audios:
            with NamedTemporaryFile() as temp:
                temp.write(audio.getvalue())
                fpath = temp.name
                outputs = model(fpath)
                cand = outputs["cands"][0]

                st.write(f"Output for {audio.name}:")
                st.write(" - ", cand)


if __name__ == "__main__":
    main()