from datetime import datetime

import numpy as np
import scipy.io.wavfile as wavfile
import whisper
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

# Whisper 모델 로드 (tiny 모델로 실시간성 유지)
model = whisper.load_model("tiny")

# 클라이언트 관리
clients = {}


@app.websocket("/audio-stream")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    client_id = str(id(websocket))
    clients[client_id] = websocket
    print(f"Client {client_id} connected")

    try:
        while True:
            # 오디오 청크 수신
            print("클라이언트 대기중......")
            audio_chunk = await websocket.receive_bytes()

            # 오디오 데이터를 새로운 버퍼에 저장 (기존 데이터 누적 방지)
            audio_buffer = bytearray(audio_chunk)  # 🔥 새로운 데이터로 덮어쓰기

            # 수신 크기 확인
            print(f"Received data size: {len(audio_chunk)} bytes")

            # 오디오 바이너리 데이터 => 숫자배열(numpy)로 해석
            audio_np = np.frombuffer(audio_buffer, dtype=np.int16).copy()

            # WAV 파일로 저장 (덮어쓰기)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"recorded_audio_{timestamp}.wav"
            wavfile.write(output_file, 16000, audio_np)

            # STT 처리
            print("STT 처리중.....")
            stt_result = model.transcribe(output_file, language="ko")
            transcription = stt_result["text"]
            print("STT 처리완료!")

            # 빈 문자열이 아닌 경우만 전송
            if transcription.strip():
                print("클라이언트에 데이터 전송")
                await websocket.send_text(transcription)
            print("=================== END =======================")
    except WebSocketDisconnect:
        print(f"Client {client_id} disconnected")
        del clients[client_id]
    except Exception as e:
        print(f"Error: {e}")
        await websocket.send_text(f"Error: {str(e)}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)