from fastapi import APIRouter, Depends, HTTPException, status, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session
from sqlalchemy import and_, or_
from typing import List, Optional
from datetime import datetime
import json
from ...shared.config.database import get_db
from ...shared.config.logging_config import get_logger
from ...security.middleware.auth import get_current_user, get_current_user_optional
from ...auth.models.user import User
from ...ai.models.chat import Chat, ChatMessage, ChatStatus
from ...auth.models.role import Role

logger = get_logger(__name__)
router = APIRouter(prefix='/chat', tags=['chat'])

class ConnectionManager:

    def __init__(self):
        self.active_connections: dict[int, List[WebSocket]] = {}
        self.staff_connections: dict[int, WebSocket] = {}
        self.visitor_connections: dict[int, WebSocket] = {}

    async def connect_chat(self, websocket: WebSocket, chat_id: int, user_type: str):
        await websocket.accept()
        if chat_id not in self.active_connections:
            self.active_connections[chat_id] = []
        self.active_connections[chat_id].append(websocket)
        if user_type == 'staff':
            pass
        elif user_type == 'visitor':
            self.visitor_connections[chat_id] = websocket

    def disconnect_chat(self, websocket: WebSocket, chat_id: int):
        if chat_id in self.active_connections:
            if websocket in self.active_connections[chat_id]:
                self.active_connections[chat_id].remove(websocket)
            if not self.active_connections[chat_id]:
                del self.active_connections[chat_id]
        if chat_id in self.visitor_connections and self.visitor_connections[chat_id] == websocket:
            del self.visitor_connections[chat_id]

    async def send_personal_message(self, message: dict, websocket: WebSocket):
        try:
            await websocket.send_json(message)
        except Exception as e:
            logger.error(f'Error sending message: {str(e)}', exc_info=True)

    async def broadcast_to_chat(self, message: dict, chat_id: int):
        if chat_id in self.active_connections:
            disconnected = []
            for connection in self.active_connections[chat_id]:
                try:
                    await connection.send_json(message)
                except Exception as e:
                    logger.error(f'Error broadcasting to connection: {str(e)}', exc_info=True, extra={'chat_id': chat_id})
                    disconnected.append(connection)
            for conn in disconnected:
                self.active_connections[chat_id].remove(conn)

    async def notify_staff_new_chat(self, chat_data: dict):
        disconnected = []
        for user_id, websocket in self.staff_connections.items():
            try:
                await websocket.send_json({'type': 'new_chat', 'data': chat_data})
            except Exception as e:
                logger.error(f'Error notifying staff {user_id}: {str(e)}', exc_info=True, extra={'staff_id': user_id})
                disconnected.append(user_id)
        for user_id in disconnected:
            del self.staff_connections[user_id]

    async def notify_staff_new_message(self, chat_id: int, message_data: dict, chat: Chat):
        if message_data.get('sender_type') == 'visitor':
            notification_data = {'type': 'new_message_notification', 'data': {'chat_id': chat_id, 'chat': {'id': chat.id, 'visitor_name': chat.visitor_name, 'visitor_email': chat.visitor_email, 'status': chat.status.value, 'created_at': chat.created_at.isoformat()}, 'message': {'id': message_data.get('id'), 'message': message_data.get('message'), 'sender_type': message_data.get('sender_type'), 'created_at': message_data.get('created_at')}}}
            disconnected = []
            for user_id, websocket in self.staff_connections.items():
                try:
                    await websocket.send_json(notification_data)
                except Exception as e:
                    logger.error(f'Error notifying staff {user_id}: {str(e)}', exc_info=True, extra={'staff_id': user_id, 'chat_id': chat_id})
                    disconnected.append(user_id)
            for user_id in disconnected:
                del self.staff_connections[user_id]

    def connect_staff(self, user_id: int, websocket: WebSocket):
        self.staff_connections[user_id] = websocket

    def disconnect_staff(self, user_id: int):
        if user_id in self.staff_connections:
            del self.staff_connections[user_id]
manager = ConnectionManager()

@router.post('/create', status_code=status.HTTP_201_CREATED)
async def create_chat(visitor_name: Optional[str]=None, visitor_email: Optional[str]=None, visitor_phone: Optional[str]=None, current_user: Optional[User]=Depends(get_current_user_optional), db: Session=Depends(get_db)):
    try:
        if current_user:
            chat = Chat(visitor_id=current_user.id, visitor_name=current_user.full_name, visitor_email=current_user.email, status=ChatStatus.pending)
        else:
            if not visitor_name:
                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Visitor name is required')
            if not visitor_email:
                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Visitor email is required')
            chat = Chat(visitor_name=visitor_name, visitor_email=visitor_email, status=ChatStatus.pending)
        db.add(chat)
        db.commit()
        db.refresh(chat)
        chat_data = {'id': chat.id, 'visitor_name': chat.visitor_name, 'visitor_email': chat.visitor_email, 'status': chat.status.value, 'created_at': chat.created_at.isoformat()}
        await manager.notify_staff_new_chat(chat_data)
        return {'success': True, 'data': {'id': chat.id, 'visitor_name': chat.visitor_name, 'visitor_email': chat.visitor_email, 'status': chat.status.value, 'created_at': chat.created_at.isoformat()}}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))

@router.post('/{chat_id}/accept')
async def accept_chat(chat_id: int, current_user: User=Depends(get_current_user), db: Session=Depends(get_db)):
    try:
        if current_user.role.name not in ['staff', 'admin']:
            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail='Only staff members can accept chats')
        chat = db.query(Chat).filter(Chat.id == chat_id).first()
        if not chat:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='Chat not found')
        if chat.status != ChatStatus.pending:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Chat is not pending')
        chat.staff_id = current_user.id
        chat.status = ChatStatus.active
        db.commit()
        db.refresh(chat)
        await manager.broadcast_to_chat({'type': 'chat_accepted', 'data': {'staff_name': current_user.full_name, 'staff_id': current_user.id}}, chat_id)
        return {'success': True, 'data': {'id': chat.id, 'staff_id': chat.staff_id, 'staff_name': current_user.full_name, 'status': chat.status.value}}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))

@router.get('/list')
async def list_chats(status_filter: Optional[str]=None, current_user: User=Depends(get_current_user), db: Session=Depends(get_db)):
    if current_user.role.name in ['staff', 'admin']:
        query = db.query(Chat)
        if status_filter:
            try:
                status_enum = ChatStatus(status_filter)
                query = query.filter(Chat.status == status_enum)
            except ValueError:
                pass
        chats = query.order_by(Chat.created_at.desc()).all()
    else:
        chats = db.query(Chat).filter(Chat.visitor_id == current_user.id).order_by(Chat.created_at.desc()).all()
    return {'success': True, 'data': [{'id': chat.id, 'visitor_id': chat.visitor_id, 'visitor_name': chat.visitor_name, 'visitor_email': chat.visitor_email, 'staff_id': chat.staff_id, 'staff_name': chat.staff.full_name if chat.staff else None, 'status': chat.status.value, 'created_at': chat.created_at.isoformat(), 'updated_at': chat.updated_at.isoformat(), 'message_count': len(chat.messages)} for chat in chats]}

@router.get('/{chat_id}')
async def get_chat(chat_id: int, current_user: Optional[User]=Depends(get_current_user_optional), db: Session=Depends(get_db)):
    chat = db.query(Chat).filter(Chat.id == chat_id).first()
    if not chat:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='Chat not found')
    if current_user:
        if current_user.role.name not in ['staff', 'admin']:
            if chat.visitor_id != current_user.id:
                raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="You don't have permission to view this chat")
    return {'success': True, 'data': {'id': chat.id, 'visitor_id': chat.visitor_id, 'visitor_name': chat.visitor_name, 'visitor_email': chat.visitor_email, 'staff_id': chat.staff_id, 'staff_name': chat.staff.full_name if chat.staff else None, 'status': chat.status.value, 'created_at': chat.created_at.isoformat(), 'updated_at': chat.updated_at.isoformat()}}

@router.get('/{chat_id}/messages')
async def get_messages(chat_id: int, current_user: Optional[User]=Depends(get_current_user_optional), db: Session=Depends(get_db)):
    chat = db.query(Chat).filter(Chat.id == chat_id).first()
    if not chat:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='Chat not found')
    if current_user:
        if current_user.role.name not in ['staff', 'admin']:
            if chat.visitor_id != current_user.id:
                raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="You don't have permission to view this chat")
    else:
        pass
    messages = db.query(ChatMessage).filter(ChatMessage.chat_id == chat_id).order_by(ChatMessage.created_at.asc()).all()
    return {'success': True, 'data': [{'id': msg.id, 'chat_id': msg.chat_id, 'sender_id': msg.sender_id, 'sender_type': msg.sender_type, 'sender_name': msg.sender.full_name if msg.sender else None, 'message': msg.message, 'is_read': msg.is_read, 'created_at': msg.created_at.isoformat()} for msg in messages]}

@router.post('/{chat_id}/message')
async def send_message(chat_id: int, message: str, current_user: Optional[User]=Depends(get_current_user_optional), db: Session=Depends(get_db)):
    try:
        chat = db.query(Chat).filter(Chat.id == chat_id).first()
        if not chat:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='Chat not found')
        if chat.status == ChatStatus.closed:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Chat is closed')
        sender_type = 'visitor'
        sender_id = None
        if current_user:
            if current_user.role.name in ['staff', 'admin']:
                sender_type = 'staff'
                sender_id = current_user.id
            else:
                sender_type = 'visitor'
                sender_id = current_user.id
                if chat.visitor_id != current_user.id:
                    raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="You don't have permission to send messages in this chat")
        else:
            sender_type = 'visitor'
            sender_id = None
        chat_message = ChatMessage(chat_id=chat_id, sender_id=sender_id, sender_type=sender_type, message=message)
        db.add(chat_message)
        db.commit()
        db.refresh(chat_message)
        message_data = {'type': 'new_message', 'data': {'id': chat_message.id, 'chat_id': chat_message.chat_id, 'sender_id': chat_message.sender_id, 'sender_type': chat_message.sender_type, 'sender_name': chat_message.sender.full_name if chat_message.sender else None, 'message': chat_message.message, 'is_read': chat_message.is_read, 'created_at': chat_message.created_at.isoformat()}}
        await manager.broadcast_to_chat(message_data, chat_id)
        if chat_message.sender_type == 'visitor':
            await manager.notify_staff_new_message(chat_id, message_data['data'], chat)
        return {'success': True, 'data': {'id': chat_message.id, 'chat_id': chat_message.chat_id, 'sender_type': chat_message.sender_type, 'message': chat_message.message, 'created_at': chat_message.created_at.isoformat()}}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))

@router.post('/{chat_id}/close')
async def close_chat(chat_id: int, current_user: Optional[User]=Depends(get_current_user_optional), db: Session=Depends(get_db)):
    try:
        chat = db.query(Chat).filter(Chat.id == chat_id).first()
        if not chat:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='Chat not found')
        if current_user:
            if current_user.role.name not in ['staff', 'admin']:
                if chat.visitor_id != current_user.id:
                    raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="You don't have permission to close this chat")
        chat.status = ChatStatus.closed
        chat.closed_at = datetime.utcnow()
        db.commit()
        await manager.broadcast_to_chat({'type': 'chat_closed', 'data': {'chat_id': chat_id}}, chat_id)
        return {'success': True, 'data': {'id': chat.id, 'status': chat.status.value}}
    except HTTPException:
        raise
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))

@router.websocket('/ws/{chat_id}')
async def websocket_chat(websocket: WebSocket, chat_id: int, user_type: str=None, token: Optional[str]=None):
    query_params = dict(websocket.query_params)
    user_type = query_params.get('user_type', 'visitor')
    
    # Security: Read token from cookies instead of query parameters
    # Cookies are sent automatically with WebSocket connections
    token = None
    if websocket.cookies:
        token = websocket.cookies.get('accessToken')
    
    # Fallback to query parameter for backward compatibility (deprecated)
    if not token:
        token = query_params.get('token')
    
    current_user = None
    if user_type == 'staff' and token:
        try:
            from ...security.middleware.auth import verify_token
            from ...shared.config.database import get_db
            payload = verify_token(token)
            user_id = payload.get('userId')
            db_gen = get_db()
            db = next(db_gen)
            try:
                current_user = db.query(User).filter(User.id == user_id).first()
                if not current_user or current_user.role.name not in ['staff', 'admin']:
                    await websocket.close(code=1008, reason='Unauthorized')
                    return
            finally:
                db.close()
        except Exception as e:
            await websocket.close(code=1008, reason='Invalid token')
            return
    await manager.connect_chat(websocket, chat_id, user_type)
    if user_type == 'staff' and current_user:
        manager.connect_staff(current_user.id, websocket)
    try:
        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)
            if message_data.get('type') == 'message':
                from ...shared.config.database import get_db
                db_gen = get_db()
                db = next(db_gen)
                try:
                    chat = db.query(Chat).filter(Chat.id == chat_id).first()
                    if not chat:
                        continue
                    sender_id = current_user.id if current_user else None
                    sender_type = 'staff' if user_type == 'staff' else 'visitor'
                    chat_message = ChatMessage(chat_id=chat_id, sender_id=sender_id, sender_type=sender_type, message=message_data.get('message', ''))
                    db.add(chat_message)
                    db.commit()
                    db.refresh(chat_message)
                except Exception:
                    db.rollback()
                    raise
                finally:
                    db.close()
                message_data = {'type': 'new_message', 'data': {'id': chat_message.id, 'chat_id': chat_message.chat_id, 'sender_id': chat_message.sender_id, 'sender_type': chat_message.sender_type, 'sender_name': chat_message.sender.full_name if chat_message.sender else None, 'message': chat_message.message, 'is_read': chat_message.is_read, 'created_at': chat_message.created_at.isoformat()}}
                await manager.broadcast_to_chat(message_data, chat_id)
                if chat_message.sender_type == 'visitor':
                    await manager.notify_staff_new_message(chat_id, message_data['data'], chat)
    except WebSocketDisconnect:
        manager.disconnect_chat(websocket, chat_id)
        if user_type == 'staff' and current_user:
            manager.disconnect_staff(current_user.id)

@router.websocket('/ws/staff/notifications')
async def websocket_staff_notifications(websocket: WebSocket):
    current_user = None
    try:
        await websocket.accept()
        
        # Security: Read token from cookies instead of query parameters
        # Cookies are sent automatically with WebSocket connections
        token = None
        if websocket.cookies:
            token = websocket.cookies.get('accessToken')
        
        # Fallback to query parameter for backward compatibility (deprecated)
        if not token:
            query_params = dict(websocket.query_params)
            token = query_params.get('token')
        
        if not token:
            await websocket.close(code=1008, reason='Token required')
            return
        try:
            from ...security.middleware.auth import verify_token
            from ...shared.config.database import get_db
            payload = verify_token(token)
            user_id = payload.get('userId')
            if not user_id:
                await websocket.close(code=1008, reason='Invalid token payload')
                return
            db_gen = get_db()
            db = next(db_gen)
            try:
                current_user = db.query(User).filter(User.id == user_id).first()
                if not current_user:
                    await websocket.close(code=1008, reason='User not found')
                    return
                role = db.query(Role).filter(Role.id == current_user.role_id).first()
                if not role or role.name not in ['staff', 'admin', 'housekeeping']:
                    await websocket.close(code=1008, reason='Unauthorized role')
                    return
            finally:
                db.close()
        except Exception as e:
            logger.error(f'WebSocket token verification error: {str(e)}', exc_info=True)
            await websocket.close(code=1008, reason=f'Token verification failed: {str(e)}')
            return
        manager.connect_staff(current_user.id, websocket)
        try:
            await websocket.send_json({'type': 'connected', 'data': {'message': 'WebSocket connected'}})
        except Exception as e:
            logger.error(f'Error sending initial message: {str(e)}', exc_info=True, extra={'user_id': current_user.id})
        while True:
            try:
                data = await websocket.receive_text()
                try:
                    message_data = json.loads(data)
                    if message_data.get('type') == 'ping':
                        await websocket.send_json({'type': 'pong', 'data': 'pong'})
                except json.JSONDecodeError:
                    await websocket.send_json({'type': 'pong', 'data': 'pong'})
            except WebSocketDisconnect:
                logger.info('WebSocket disconnected normally', extra={'user_id': current_user.id})
                break
            except Exception as e:
                logger.error(f'WebSocket receive error: {str(e)}', exc_info=True, extra={'user_id': current_user.id})
                break
    except WebSocketDisconnect:
        logger.info('WebSocket disconnected')
    except Exception as e:
        logger.error(f'WebSocket error: {str(e)}', exc_info=True)
    finally:
        if current_user:
            try:
                manager.disconnect_staff(current_user.id)
            except:
                pass