"""
@File    :   stt.py
@Time    :   2024/07/10 09:57:46
@Author  :   RayLam
@Contact :   1027196450@qq.com
"""

from typing import Optional, Union, List, Dict
import argparse
import time

import numpy as np
import torch
import pyaudio
import pynput
import pyperclip
# from faster_whisper import WhisperModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline


parser = argparse.ArgumentParser("stt")
parser.add_argument(
    "-r",
    "--record_time",
    default=5,
    type=int,
    help="record time(second)",
)
parser.add_argument(
    "-l",
    "--record_src_language",
    default="zh",
    help="record src language",
)

parser.add_argument("-i", "--is_translate", action="store_true", help="is_translate")
parser.add_argument(
    "-s",
    "--translate_src_lang",
    default="zho_Hans",
    help="translate zho_Hans, zho_Hant,eng_Latn, jpn_Jpan,yue_Hant,kor_Hang",
)
parser.add_argument(
    "-t",
    "--translate_tgt_lang",
    default="eng_Latn",
    help="translate zho_Hans, zho_Hant,eng_Latn, jpn_Jpan,yue_Hant,kor_Hang",  # zho_Hant繁体
)
parser.add_argument(
    "-m",
    "--translate_model",
    default="facebook/nllb-200-distilled-600M",
    help="translate model",
)
args = parser.parse_args()

signal = 0
INT16_MAX_ABS_VALUE = 32768.0
CHUNK = 512
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000


model_id = "openai/whisper-large-v3-turbo"
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
# stt_model = WhisperModel(model_size, device=device, compute_type="float16")
stt_model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True
)
stt_model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=stt_model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    chunk_length_s=30,
    batch_size=16,  # batch size for inference - set based on your device
    torch_dtype=torch_dtype,
    device=device,
)

generate_kwargs={"language": "chinese",    
                 "logprob_threshold": -1.0,
                 "no_speech_threshold": 0.6,
                 "return_timestamps": False,}

# language:
#   - en
#   - zh
#   - de
#   - es
#   - ru
#   - ko
#   - fr
#   - ja
#   - pt
#   - tr
#   - pl
#   - ca
#   - nl
#   - ar
#   - sv
#   - it
#   - id
#   - hi
#   - fi
#   - vi
#   - he
#   - uk
#   - el
#   - ms
#   - cs
#   - ro
#   - da
#   - hu
#   - ta
#   - 'no'
#   - th
#   - ur
#   - hr
#   - bg
#   - lt
#   - la
#   - mi
#   - ml
#   - cy
#   - sk
#   - te
#   - fa
#   - lv
#   - bn
#   - sr
#   - az
#   - sl
#   - kn
#   - et
#   - mk
#   - br
#   - eu
#   - is
#   - hy
#   - ne
#   - mn
#   - bs
#   - kk
#   - sq
#   - sw
#   - gl
#   - mr
#   - pa
#   - si
#   - km
#   - sn
#   - yo
#   - so
#   - af
#   - oc
#   - ka
#   - be
#   - tg
#   - sd
#   - gu
#   - am
#   - yi
#   - lo
#   - uz
#   - fo
#   - ht
#   - ps
#   - tk
#   - nn
#   - mt
#   - sa
#   - lb
#   - my
#   - bo
#   - tl
#   - mg
#   - as
#   - tt
#   - haw
#   - ln
#   - ha
#   - ba
#   - jw
#   - su

def on_press(key: pynput.keyboard.Key) -> Optional[bool]:
    if key == pynput.keyboard.Key.esc:
        global signal
        signal = 1
        return False
    elif hasattr(key, "char"):
        if key.char == "f":
            return False
    # print("Key {0} pressed".format(key))


def kb_listener() -> None:
    listener = pynput.keyboard.Listener(on_press=on_press)
    listener.start()
    listener.join()


def speech_to_text() -> str:
    # 创建一个PyAudio对象
    p = pyaudio.PyAudio()

    # 打开麦克风流
    stream = p.open(
        format=FORMAT,
        channels=CHANNELS,
        rate=RATE,
        input=True,
        output=True,
        frames_per_buffer=CHUNK,
    )

    audio_collector = []
    counter = 0
    unit = RATE // CHUNK
    text = ""
    start = time.time()
    print("start record...")
    while True:
        audio_data = stream.read(CHUNK)
        audio_collector.append(audio_data)
        counter += 1
        if counter % (int(unit * args.record_time) + 1) == 0:
            audio_collector = np.frombuffer(b"".join(audio_collector), dtype=np.int16)
            audio_collector = audio_collector.astype(np.float32) / INT16_MAX_ABS_VALUE
            print("record time: ", time.time() - start)
            # segments, info = stt_model.transcribe(
            #     audio_collector,
            #     language=args.record_src_language,
            #     beam_size=5,
            #     log_prob_threshold=-0.5,
            # )
            segments = pipe(audio_collector, generate_kwargs=generate_kwargs)

            text = ""
            # print("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text))
            if segments['text'] != "":
                text += segments['text']
            # initial
            audio_collector = []
            counter = 0
            break

    # 关闭流和PyAudio对象
    stream.stop_stream()
    stream.close()
    p.terminate()

    return text


def main() -> None:
    global signal

    print("record_lang: ", args.record_src_language)

    if args.is_translate:
        tokenizer = AutoTokenizer.from_pretrained(args.translate_model)
        translate_model = AutoModelForSeq2SeqLM.from_pretrained(args.translate_model)

        translator = pipeline(
            "translation",
            model=translate_model,
            tokenizer=tokenizer,
            src_lang=args.translate_src_lang,
            tgt_lang=args.translate_tgt_lang,
            max_length=512,
            device=0,
        )
        print(
            "translate model is ready, src_lang: ",
            args.translate_src_lang,
            "tgt_lang: ",
            args.translate_tgt_lang,
        )

    while True:
        kb_listener()
        if signal == 1:
            break
        text = speech_to_text()
        print("record is done")
        translated_text = ""
        if args.is_translate and text != "":
            for tl in translator(text):
                translated_text += tl["translation_text"]
            print("sequence is translated")
            text = translated_text
        pyperclip.copy(text)


if __name__ == "__main__":
    main()
