import time
import json
import queue
import requests
import threading
import datetime as dt
import numpy as np
from funasr import AutoModel
from .shared import audio_q, stop_flag, asr_flag


class ASRRunner:
    def __init__(self, cfg):
        self.chunk_size = [0, 10, 5]        # [past_ms, cur_ms, subsample]，10*60ms=600ms窗口
        self.encoder_chunk_look_back = cfg['encoder_chunk_look_back']
        self.decoder_chunk_look_back = cfg['decoder_chunk_look_back']
        self.chunk_stride = self.chunk_size[1] * 960  # 600ms
        self.chunk_stride = max(int(self.chunk_stride), 1600)  # 不低于约100ms分辨率
        self.max_dur = cfg['max_duration']
        self.url = cfg['url']

        self.model = AutoModel(
            model=cfg['asr_model'],
            disable_update=cfg['disable_update'],  
            device=cfg['device'],
            )
        self.punc = AutoModel(model=cfg['punc_model'], disable_update=cfg['disable_update'])
        self.cache = {}
        self.start_time = None
        self.buf = np.empty((0,), dtype=np.int16)

        self.empty_txt_count = 0
        self.max_empty_txt_count = cfg['max_empty_txt_count']
        self.text_buffer = ""
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

    def run(self):
        while not stop_flag.is_set():
            if not asr_flag.is_set():
                time.sleep(0.5)
                continue
            else:
                if self.start_time is None:
                    self.start_time = dt.datetime.now()

            try:
                chunk = audio_q.get(timeout=0.5)            # int16 mono
            except queue.Empty:
                time.sleep(0.1)
                continue
            
            # TODO: get audio frame from auido_q and concat them into fixed length (chunk_stride)
            if chunk is None:
                continue
            if chunk.dtype != np.int16:
                chunk = chunk.astype(np.int16, copy=False)
            self.buf = np.concatenate((self.buf, chunk), dtype=np.int16)
            if self.buf.shape[0] < self.chunk_stride:
                continue

            res = None
            while self.buf.shape[0] >= self.chunk_stride:
                speech_chunk = self.buf[: self.chunk_stride]
                self.buf = self.buf[self.chunk_stride :]
                speech_chunk = np.ascontiguousarray(
                    speech_chunk.astype(np.float32, copy=False)
                )
                res = self.model.generate(
                    input=speech_chunk,
                    cache=self.cache,  # 关键：启用流式缓存
                    is_final=False,
                    chunk_size=self.chunk_size,
                    encoder_chunk_look_back=self.encoder_chunk_look_back,
                    decoder_chunk_look_back=self.decoder_chunk_look_back,
                    fs=16000,  # 明确采样率（SoundInputStream 已以 16k 推送）
                )

            if res is None:
                continue
            
            print(res)
            cur_time = dt.datetime.now()
            if (cur_time - self.start_time).total_seconds() > self.max_dur:
                self.switch()
            
            for r in res:
                if r['text'] == '':
                    self.empty_txt_count += 1
                else:
                    self.empty_txt_count = 0
                    self.text_buffer += r['text']

                if self.empty_txt_count > self.max_empty_txt_count:
                    self.switch()
                    continue

    def do_task(self):
        if len(self.text_buffer) == 0:
            return
        res = self.punc.generate(input=self.text_buffer)
        data = ""
        for r in res:
            data += r['text']
        print(data)
        # try:
        #     data_json= json.dumps({'text': data})
        #     response = requests.post(self.url, data = data_json)
        #     # if response.result != "OK":
        #     #     print("request Failed")
        # except Exception as e:
        #     print("ASRRunner send request failed: ", e)

    def switch(self):
        asr_flag.clear()
        self.do_task()
        self.cache = {}
        self.text_buffer = ""
        self.buf = np.empty((0,), dtype=np.int16)
        self.start_time = None
        self.empty_txt_count = 0
        
