from datetime import datetime
import wave
import json

# -------------- 公共类 ------------ #        
class SentenceSegmentation():
    def __init__(self,):
        self.is_first_sentence = True
        self.cache = ""
        
    def __sentenceSegmentation(self, llm_frame):
        results = []
        if llm_frame['is_end']:
            if self.cache:
                results.append(self.cache)
            self.cache = ""
            return results
        for char in llm_frame['msg']:
            self.cache += char
            if self.is_first_sentence and char in ',.?!，。？！':
                results.append(self.cache)
                self.cache = ""
                self.is_first_sentence = False
            elif char in '。？！':
                results.append(self.cache)
                self.cache = ""
        return results
    
    def split(self, llm_chunk):
        return self.__sentenceSegmentation(llm_chunk)

class Recorder:
    def __init__(self, user_id, input_sr, output_sr):
        self.input_wav_path = 'storage/wav/'+ datetime.now().strftime('%Y%m%d%H%M%S') + 'U' + user_id + 'i.wav'
        self.output_wav_path = 'storage/wav/'+ datetime.now().strftime('%Y%m%d%H%M%S') + 'U' + user_id + 'o.wav'
        self.out_put_text_path = 'storage/record/'+ datetime.now().strftime('%Y%m%d%H%M%S') + 'U' + user_id + 'o.txt'
        self.input_sr = input_sr
        self.output_sr = output_sr
        self.user_audio = b''
        self.tts_audio = b''
        self.input_text = ""
        self.output_text = ""
        
    def write(self):
        record = {"input_wav":self.input_wav_path,"input_text":self.input_text,"input_sr":self.input_sr,"output_wav":self.output_wav_path,"output_text":self.output_text,"output_sr":self.output_sr}
        with wave.open(self.input_wav_path, 'wb') as wav_file:
            wav_file.setparams((1, 2, self.input_sr, 0, 'NONE', 'not compressed'))
            wav_file.writeframes(self.user_audio)
        with wave.open(self.output_wav_path, 'wb') as wav_file:
            wav_file.setparams((1, 2, self.output_sr, 0, 'NONE', 'not compressed'))
            wav_file.writeframes(self.tts_audio)
        with open(self.out_put_text_path, 'w', encoding='utf-8') as file:
            file.write(json.dumps(record, ensure_ascii=False))
# ---------------------------------- #