from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query, status, Depends
from sqlalchemy.ext.asyncio import AsyncSession
from app.auth import get_user_from_token
from app.database import get_transactional_session
from app.crud import list as crud_list, group as crud_group
from app.core.websocket import websocket_manager, WebSocketEvent
from app.models import User
import asyncio
import json
import logging

router = APIRouter()
logger = logging.getLogger(__name__)

@router.websocket("/ws/household/{household_id}")
async def household_websocket_endpoint(
    websocket: WebSocket,
    household_id: int,
    token: str = Query(...),
    db: AsyncSession = Depends(get_transactional_session),
):
    """
    Main WebSocket endpoint for household-wide real-time features
    Handles: chores, expenses, general household activity, member presence
    """
    user = await get_user_from_token(token, db)
    if not user:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    try:
        # Verify user is member of this household
        household = await crud_group.get_group_with_user_check(db, household_id, user.id)
        if not household:
            await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)
            return
    except Exception as e:
        logger.error(f"Household permission check failed: {e}")
        await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)
        return

    # Connect to WebSocket manager
    connected = await websocket_manager.connect(websocket, user.id, user.name or f"User {user.id}")
    if not connected:
        return

    # Subscribe to household room
    household_room = f"household:{household_id}"
    websocket_manager.join_room(user.id, household_room)

    # Notify other members that user came online
    member_activity_event = WebSocketEvent(
        event="group:member_activity",
        payload={
            "user_id": user.id,
            "last_seen": user.created_at.isoformat(),  # Use created_at as placeholder
            "is_active": True,
            "action": "connected"
        },
        timestamp=websocket_manager.connections[user.id].connected_at.isoformat(),
        user_id=user.id
    )
    await websocket_manager.send_to_room(household_room, member_activity_event, exclude_user=user.id)

    try:
        # Send welcome message
        welcome_event = WebSocketEvent(
            event="connection:established",
            payload={
                "message": f"Connected to household {household.name}",
                "household_id": household_id,
                "rooms": [household_room]
            },
            timestamp=websocket_manager.connections[user.id].connected_at.isoformat()
        )
        await websocket_manager.send_to_user(user.id, welcome_event)

        # Keep connection alive and handle incoming messages
        while True:
            try:
                # Wait for incoming messages
                message = await websocket.receive_text()
                data = json.loads(message)
                
                # Handle different message types
                await handle_websocket_message(user.id, household_id, data, db)
                
            except WebSocketDisconnect:
                break
            except json.JSONDecodeError:
                logger.warning(f"Invalid JSON from user {user.id}")
            except Exception as e:
                logger.error(f"Error handling WebSocket message from user {user.id}: {e}")

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"WebSocket error for user {user.id} in household {household_id}: {e}")
    finally:
        # Notify other members that user went offline
        offline_event = WebSocketEvent(
            event="group:member_activity",
            payload={
                "user_id": user.id,
                "last_seen": websocket_manager.connections[user.id].last_activity.isoformat() if user.id in websocket_manager.connections else None,
                "is_active": False,
                "action": "disconnected"
            },
            timestamp=websocket_manager.connections[user.id].last_activity.isoformat() if user.id in websocket_manager.connections else None,
            user_id=user.id
        )
        await websocket_manager.send_to_room(household_room, offline_event, exclude_user=user.id)
        
        # Disconnect from WebSocket manager
        await websocket_manager.disconnect(user.id)

@router.websocket("/ws/list/{list_id}")
async def list_websocket_endpoint(
    websocket: WebSocket,
    list_id: int,
    token: str = Query(...),
    db: AsyncSession = Depends(get_transactional_session),
):
    """
    WebSocket endpoint for list-specific real-time features
    Handles: item claiming, real-time collaboration, purchase confirmations
    """
    user = await get_user_from_token(token, db)
    if not user:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    try:
        # Verify user has access to this list
        await crud_list.check_list_permission(db, list_id, user.id)
    except Exception as e:
        logger.error(f"List permission check failed: {e}")
        await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)
        return

    # Connect to WebSocket manager (reuse existing connection if available)
    if user.id not in websocket_manager.connections:
        connected = await websocket_manager.connect(websocket, user.id, user.name or f"User {user.id}")
        if not connected:
            return
    else:
        # Use existing connection but accept this websocket for this specific list
        await websocket.accept()

    # Subscribe to list-specific room
    list_room = f"list:{list_id}"
    websocket_manager.join_room(user.id, list_room)

    try:
        # Send list connection confirmation
        list_welcome = WebSocketEvent(
            event="list:connected",
            payload={
                "message": f"Connected to list {list_id}",
                "list_id": list_id,
                "active_users": websocket_manager.get_room_members(list_room)
            },
            timestamp=websocket_manager.connections[user.id].last_activity.isoformat()
        )
        await websocket_manager.send_to_user(user.id, list_welcome)

        # Keep connection alive for list-specific events
        while True:
            try:
                # For list-specific endpoint, we primarily listen for events
                # Most interactions will be handled through the main household WebSocket
                await asyncio.sleep(30)  # Periodic check
                
                # Send periodic ping to keep list connection active
                ping_event = WebSocketEvent(
                    event="list:ping",
                    payload={"list_id": list_id},
                    timestamp=websocket_manager.connections[user_id].last_activity.isoformat() if user.id in websocket_manager.connections else None
                )
                await websocket_manager.send_to_user(user.id, ping_event)
                
            except WebSocketDisconnect:
                break
            except Exception as e:
                logger.error(f"List WebSocket error for user {user.id}: {e}")
                break

    except WebSocketDisconnect:
        pass
    finally:
        # Leave list room but don't disconnect main connection
        websocket_manager.leave_room(user.id, list_room)

async def handle_websocket_message(user_id: int, household_id: int, data: dict, db: AsyncSession):
    """
    Handle incoming WebSocket messages from clients
    Supports editing indicators, presence updates, etc.
    """
    message_type = data.get("type")
    payload = data.get("payload", {})

    if message_type == "editing_started":
        # User started editing an entity
        entity_type = payload.get("entity_type")
        entity_id = payload.get("entity_id")
        field = payload.get("field")
        
        if entity_type and entity_id:
            event = WebSocketEvent(
                event=f"{entity_type}:editing_started",
                payload={
                    f"{entity_type}_id": entity_id,
                    "user_id": user_id,
                    "user_name": websocket_manager.connections[user_id].user_name,
                    "field": field
                },
                timestamp=websocket_manager.connections[user_id].last_activity.isoformat(),
                user_id=user_id
            )
            
            # Send to appropriate room
            if entity_type == "expense":
                await websocket_manager.broadcast_to_expense(entity_id, event, exclude_user=user_id)
            elif entity_type == "chore":
                await websocket_manager.broadcast_to_household(household_id, event, exclude_user=user_id)
            elif entity_type == "list":
                await websocket_manager.broadcast_to_list(entity_id, event, exclude_user=user_id)

    elif message_type == "editing_stopped":
        # User stopped editing an entity
        entity_type = payload.get("entity_type")
        entity_id = payload.get("entity_id")
        
        if entity_type and entity_id:
            event = WebSocketEvent(
                event=f"{entity_type}:editing_stopped",
                payload={f"{entity_type}_id": entity_id},
                timestamp=websocket_manager.connections[user_id].last_activity.isoformat(),
                user_id=user_id
            )
            
            # Send to appropriate room
            if entity_type == "expense":
                await websocket_manager.broadcast_to_expense(entity_id, event, exclude_user=user_id)
            elif entity_type == "chore":
                await websocket_manager.broadcast_to_household(household_id, event, exclude_user=user_id)
            elif entity_type == "list":
                await websocket_manager.broadcast_to_list(entity_id, event, exclude_user=user_id)

    elif message_type == "presence_update":
        # User presence/activity update
        action = payload.get("action", "active")
        
        event = WebSocketEvent(
            event="group:member_activity",
            payload={
                "user_id": user_id,
                "last_seen": websocket_manager.connections[user_id].last_activity.isoformat(),
                "is_active": True,
                "action": action
            },
            timestamp=websocket_manager.connections[user_id].last_activity.isoformat(),
            user_id=user_id
        )
        
        await websocket_manager.broadcast_to_household(household_id, event, exclude_user=user_id)

    else:
        logger.warning(f"Unknown WebSocket message type: {message_type} from user {user_id}")

# Utility endpoint to get current WebSocket status
@router.get("/ws/status")
async def websocket_status():
    """Get current WebSocket connection status"""
    return {
        "connected_users": websocket_manager.get_connected_users(),
        "total_connections": len(websocket_manager.connections),
        "total_rooms": len(websocket_manager.rooms),
        "rooms": {room: len(users) for room, users in websocket_manager.rooms.items()}
    } 