import logging
import os
import warnings

import librosa
import transformers
import torch
from dotenv import load_dotenv
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor


warnings.filterwarnings("ignore")
transformers.logging.set_verbosity_error()
load_dotenv()
MODEL = os.getenv("MODEL", "m3hrdadfi/wav2vec2-large-xlsr-persian")
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG")


def initLogger():
    logger = logging.getLogger("speech2text-fa")
    level = getattr(logging, LOG_LEVEL.upper())
    
    fmt = "%(asctime)s | %(levelname)-8s | %(message)s"
    datefmt = "%Y-%m-%d %H:%M:%S"
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    formatter = logging.Formatter(fmt, datefmt)
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    logger.setLevel(level)
    return logger


def mp3_to_text(mp3_file_path):
    # Load the MP3 file and resample to 16kHz
    audio, sample_rate = librosa.load(mp3_file_path, sr=16000)
    logger.info("Resampling is Done!")

    # Load tokenizer and model from Hugging Face
    tokenizer = Wav2Vec2Processor.from_pretrained(MODEL)
    model = Wav2Vec2ForCTC.from_pretrained(MODEL)
    logger.info("Loading model is Done!")

    # Preprocess the audio
    input_values = tokenizer(audio, sampling_rate=16000, return_tensors="pt", padding="longest").input_values
    logits = model(input_values).logits
    logger.info("Processing the audio is Done!")

    # Decode the predicted IDs
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = tokenizer.batch_decode(predicted_ids)
    logger.info("Decoding the prodicted IDs is Done!")

    return transcription[0]


if __name__ == "__main__":
    logger = initLogger()
    text = mp3_to_text("samples/sample1.wav")
    print()
    print(text)