from fastapi import WebSocket, APIRouter, WebSocketDisconnect
from backend.domains.llms.llm_chat_srvc import construir_agente_llm
from src.Logger.logger_db import LoggerDB, logger
from entrypoint.init_db import db_credencial
import json

LoggerDB(db_credencial, "logger_llm_ws", created_by="sistema")

router = APIRouter()
agente = construir_agente_llm()

@router.websocket("/ws/chat")
async def chat_ws(websocket: WebSocket):
    await websocket.accept()
    try:
        data = await websocket.receive_text()
        parsed = json.loads(data)
        prompt = parsed.get("prompt")
        if not prompt:
            await websocket.send_text("⚠️ Prompt vacío.")
            await websocket.close()
            return

        # ✅ Solución: hacer await antes de iterar
        respuesta_gen = await agente.interactuar_en_bucle(prompt=prompt, stream=True)
        async for token in respuesta_gen:
            await websocket.send_text(token)

        await websocket.close()

    except WebSocketDisconnect:
        logger.info("🔌 WebSocket desconectado por el cliente.")
    except Exception as e:
        logger.exception("❌ Error en WebSocket:")
        await websocket.close()