import json
import os

from appPublic.dictObject import DictObject
from appPublic.zmq_reqrep import ZmqReplier
from appPublic.jsonConfig import getConfig

from time import time
from faster_whisper import WhisperModel

class ZmqASR:
    def __init__(self):
        config = getConfig()
        self.model_dir = config.model_dir or "large-v3"
        self.device = config.device or "cpu"
        self.zmq_url = config.zmq_url
        self.compute_type=config.compute_type or "float16"
        self.model = WhisperModel(self.model_dir, 
                                        device=self.device, 
                                        compute_type=self.compute_type)
        self.replier = ZmqReplier(self.zmq_url, self.generate)

    def run(self):
        print(f'running {self.zmq_url}')
        self.replier._run()
        print('ended ...')

    def generate(self, d):
        msg= d.decode('utf-8')
        data = DictObject(**json.loads(msg))
        t1 = time()
        segs, info = self.model.transcribe(data.audio_file)
        t2 = time()
        d = DictObject(**{
            'segments':[],
            'language':info.language
        })
        d.content = ''
        for s in segs:
            seg = {
                "start":s.start,
                "end":s.end,
                "text":s.text
            }
            d.segments.append(seg)
            d.content += s.text
        d.time_cost = t2 - t1
        print(f'{d}')
        return json.dumps(d)

if __name__ == '__main__':
    workdir = os.getcwd()
    config = getConfig(workdir)
    print(f'{config=}')
    asr = ZmqASR()
    print('here')
    asr.run()