import llm_asr
import sounddevice as sd
import numpy as np
import tempfile
import threading
from pydub import AudioSegment
from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS

app = Flask(__name__, static_folder='../asr-front/dist', static_url_path='/')
CORS(app)

sensevoice_path = "/usr/share/spacemit-asr/sensevoice"
sensevoiceModel = llm_asr.llm_asr(sensevoice_path)

# 录音相关变量
recording_thread = None
is_recording = False
frames = []
fs = 48000
channels = 2
dtype = 'float32'

# 处理音频并识别文字
def process_audio(mp3_path):
    if not mp3_path:
        return "No audio detected."
    text = sensevoiceModel.audio_to_text(mp3_path)
    return text

# 录音函数
def record_audio():
    global is_recording, frames
    with sd.InputStream(samplerate=fs, channels=channels, dtype=dtype) as stream:
        while is_recording:
            data, _ = stream.read(1024)
            frames.append(data)

# 后端状态检测
@app.route('/health', methods=['GET'])
def health():
    app.logger.debug(request.headers)
    return 'OK'

# 开始录音
@app.route('/start', methods=['GET'])
def start_recording():
    app.logger.debug('开始录音')
    global is_recording, frames, recording_thread
    if not is_recording:
        is_recording = True
        frames = []
        recording_thread = threading.Thread(target=record_audio)
        recording_thread.start()
        return jsonify({"status": "1"})
    else:
        return jsonify({"status": "0"})

# 停止录音
@app.route('/stop', methods=['GET'])
def stop_recording():
    app.logger.debug('停止录音')
    global is_recording, frames, recording_thread
    if is_recording:
        is_recording = False
        if recording_thread:
            recording_thread.join()
        recording_thread = None
        if frames:
            recording = np.concatenate(frames, axis=0)
            mono_recording = np.mean(recording, axis=1)
            audio_data = (mono_recording * 32767).astype(np.int16)
            audio_segment = AudioSegment(
                audio_data.tobytes(),
                frame_rate=fs,
                sample_width=2,
                channels=1
            )
            temp_mp3_file = tempfile.NamedTemporaryFile(suffix='.mp3', delete=False)
            mp3_path = temp_mp3_file.name
            audio_segment.export(mp3_path, format='mp3', bitrate='256k')
            temp_mp3_file.close()
            transcription = process_audio(mp3_path)
            return jsonify({"status": "1", "transcription": transcription})
        else:
            return jsonify({"status": "1", "transcription": "没有检测到语音帧！"})
    else:
        return jsonify({"status": "0"})

@app.route('/')
def index():
    return send_from_directory(app.static_folder, 'index.html')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860,  threaded=True)
