from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from contextlib import asynccontextmanager
from pydantic import BaseModel
from core.backend import AIBackend
import uvicorn, json, os
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from services.stt import StreamingSTT
from services.tts import text_to_speech_stream
from fastapi.staticfiles import StaticFiles


chatbot_obj = AIBackend()

@asynccontextmanager
async def lifespan(app: FastAPI):
    await chatbot_obj.async_setup()
    yield
    if chatbot_obj.conn:
        await chatbot_obj.conn.close()

app = FastAPI(lifespan=lifespan)

class UserRequest(BaseModel):
    user_id: str
    user_query: str

@app.post("/chat")
async def chat(request: UserRequest):
    stream = await chatbot_obj.main(
        user_id=request.user_id,
        user_query=request.user_query,
    )
    return StreamingResponse(stream, media_type="text/event-stream")

@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # receive frontend message
            data = await websocket.receive_text()
            payload = json.loads(data)
            
            user_id = payload["user_id"]
            user_query = payload["user_query"]
            
            # stream AI response
            stream = await chatbot_obj.main(
                user_id=user_id,
                user_query=user_query
            )
            
            async for chunk in stream:
                await websocket.send_text(chunk)
            
            # notify frontend response finished
            await websocket.send_text("[[END]]")
    except WebSocketDisconnect:
        print("Client disconnected")

@app.websocket("/ws/voice")
async def voice_ws(websocket: WebSocket):
    await websocket.accept()
    stt = StreamingSTT()
    try:
        while True:
            message = await websocket.receive()
            # 🎤 AUDIO INPUT
            if "bytes" in message:
                audio_chunk = message["bytes"]
                stt.add_audio(audio_chunk)
                text = stt.transcribe_if_ready()
                if not text:
                    continue
                await websocket.send_text(f"[STT]: {text}")
                # 🤖 LLM STREAM
                stream = chatbot_obj.main(
                    user_id="voice_user",
                    user_query=text
                )
                full_response = ""
                async for token in stream:
                    full_response += token
                    await websocket.send_text(f"[LLM]: {token}")
                # 🔊 TTS STREAM
                async for audio_chunk in text_to_speech_stream(full_response):
                    await websocket.send_bytes(audio_chunk)
                await websocket.send_text("[END]")
    except WebSocketDisconnect:
        print("Voice client disconnected")

if __name__ == "__main__":
    uvicorn.run("app:app", host="127.0.0.1", port=8679, reload=True)