# app/api/v1/endpoints/ws_chat.py
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from app.api import deps
from app.services import message_service
from typing import Dict
from sqlalchemy.orm import Session
from app.core.logging_config import setup_logging

router = APIRouter()
# Setup logging
logger = setup_logging()

# Dictionary to store active connections
active_connections: Dict[int, WebSocket] = {}

@router.websocket("/ws/chat/{friend_id}")
async def websocket_chat(
    websocket: WebSocket,
    friend_id: int,
    db: Session = Depends(deps.get_db),
    current_user = Depends(deps.get_current_ws_user)
):
    """
    Supports real-time chat between the currently logged-in user and the user with friend_id.
    This example primarily functions as an echo for messages.
    """
    logger.info(f"WebSocket connection attempt from user {current_user.id} to friend {friend_id}")
    await websocket.accept()
    # Store the current user's connection
    active_connections[current_user.id] = websocket
    
    try:
        while True:
            data = await websocket.receive_text()
            
            # Save the message to the database
            message = message_service.create_message(
                db=db,
                sender_id=current_user.id,
                receiver_id=friend_id,
                content=data
            )
            
            # 저장된 메시지를 로그로 출력하거나, 필요한 후처리 진행
            logger.info(f"Message saved with id {message.id}")
            
            # Send the message back to the sender
            await websocket.send_text(f"You: {data}")
            
            # If the receiver is connected, send the message to them
            if friend_id in active_connections:
                friend_ws = active_connections[friend_id]
                await friend_ws.send_text(f"{current_user.first_name}: {data}")
                
    except WebSocketDisconnect:
        # Remove from active connections on disconnect
        if current_user.id in active_connections:
            del active_connections[current_user.id]
        logger.info(f"Client #{current_user.id} disconnected")