"""
ToolMate AI Plugin - analyze audio file

analyze audio file

[TOOL_CALL]
"""

if not config.isLite:

    from toolmate import config, showErrors, getGroqClient
    from toolmate import print1, print2, print3
    from toolmate.utils.sttLanguages import googleSpeeckToTextLanguages
    import os, whisper, io, shutil, subprocess
    import speech_recognition as sr
    from pydub import AudioSegment

    # Function method
    def transcribe_audio_whisper(function_args):
        from toolmate import config
        def check_file_format(file_path):
            # List of allowed file extensions
            allowed_extensions = ('.mp3', '.mp4', '.mpeg', '.mpga', '.m4a', '.wav', '.webm')
            # Getting the file extension
            _, file_extension = os.path.splitext(file_path)
            # Checking if the file extension is in the list of allowed extensions
            return True if file_extension.lower() in allowed_extensions else False

        audio_file = function_args.get("audio_filepath") # required
        language = function_args.get("language") # required

        if audio_file and os.path.isfile(audio_file):

            if not shutil.which("ffmpeg"):
                print1("Install 'ffmpeg' first!")
                print1("Read https://github.com/openai/whisper/tree/main#setup")
                return ""
            # https://github.com/openai/whisper/tree/main#python-usage
            # platform: llamacpp or ollama
            if language.lower() in ("english", "non-english"):
                model = whisper.load_model(config.voiceTypingWhisperEnglishModel if language.lower() == "english" else "large")
                result = model.transcribe(audio_file)
            else:
                # non-English
                model = whisper.load_model("large")

                # load audio and pad/trim it to fit 30 seconds
                audio = whisper.load_audio(audio_file)
                audio = whisper.pad_or_trim(audio)

                # make log-Mel spectrogram and move to the same device as the model
                mel = whisper.log_mel_spectrogram(audio).to(model.device)

                # detect the spoken language
                _, probs = model.detect_language(mel)
                print(f"Detected language: {max(probs, key=probs.get)}")

                # decode the audio
                options = whisper.DecodingOptions()
                result = whisper.decode(model, mel, options)

            config.toolTextOutput = result['text']
            print2("```transcript")
            print1(config.toolTextOutput)
            print2("```")
            return ""

        return "[INVALID]"

    # Function Signature
    functionSignature = {
        "examples": [
            "Transcribe audio",
        ],
        "name": "transcribe_audio_whisper",
        "description": f'''Transcribe audio into text with Whisper''',
        "parameters": {
            "type": "object",
            "properties": {
                "audio_filepath": {
                    "type": "string",
                    "description": "Return the audio file path that I specified in my requests. Return an empty string '' if it is not specified.",
                },
                "language": {
                    "type": "string",
                    "description": "Audio language",
                    "enum": ["English", "non-English"],
                },
            },
            "required": ["audio_filepath", "language"],
        },
    }

    # Integrate the signature and method into LetMeDoIt AI
    config.addFunctionCall(signature=functionSignature, method=transcribe_audio_whisper)