import asyncio
import numpy as np
import logging
import warnings
import os
import threading

logging.getLogger("nemo_logger").setLevel(logging.ERROR)
logging.getLogger("nemo").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from asr import ASRPipeline

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_DIR = "/app/models"
pipeline = ASRPipeline(MODEL_DIR)

@app.api_route("/", methods=["GET", "HEAD"])
def root():
    return {"status": "CaptionAI server running"}

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    print("=" * 50)
    print("Client connected")
    print("=" * 50)

    loop = asyncio.get_event_loop()
    caption_queue = asyncio.Queue()

    async def send_results():
        while True:
            text = await caption_queue.get()
            if text is None:
                break
            try:
                await websocket.send_text(text)
            except:
                break

    sender = asyncio.create_task(send_results())

    try:
        while True:
            data = await websocket.receive_bytes()
            audio_chunk = np.frombuffer(data, dtype=np.float32)

            def process(chunk=audio_chunk):
                result = pipeline.transcribe_chunk(chunk)
                if result:
                    asyncio.run_coroutine_threadsafe(
                        caption_queue.put(result),
                        loop
                    )

            threading.Thread(target=process, daemon=True).start()

    except WebSocketDisconnect:
        pipeline.stop()
        await caption_queue.put(None)
        await sender
        print("=" * 50)
        print("Client disconnected")
        print("=" * 50)
    except Exception as e:
        print(f"Error: {e}")
        pipeline.stop()
        await caption_queue.put(None)
        await websocket.close()