const WebSocket = require('ws');
const express = require('express');
const http = require('http');
const cors = require('cors');
const dotenv = require('dotenv');
const { v4: uuidv4 } = require('uuid');

dotenv.config();

const app = express();
app.use(cors());
app.use(express.json());

const server = http.createServer(app);
const wss = new WebSocket.Server({ server });

// In-memory storage for connections
const connections = new Map();
const userConnections = new Map();
const typingStatus = new Map();

// Helper to get user key
const getUserKey = (userType, userId) => `${userType}:${userId}`;

// Helper to broadcast to specific user
const broadcastToUser = (userType, userId, message) => {
    const userKey = getUserKey(userType, userId);
    const userSockets = userConnections.get(userKey) || [];
    
    userSockets.forEach(ws => {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify(message));
        }
    });
};

// Helper to broadcast typing status
const broadcastTypingStatus = (fromUserType, fromUserId, toUserType, toUserId, isTyping) => {
    broadcastToUser(toUserType, toUserId, {
        type: 'typing',
        fromUser: { type: fromUserType, id: fromUserId },
        isTyping: isTyping
    });
};

wss.on('connection', (ws, req) => {
    const connectionId = uuidv4();
    connections.set(connectionId, ws);
    
    console.log(`New connection: ${connectionId}`);
    
    ws.on('message', (message) => {
        try {
            const data = JSON.parse(message);
            console.log('Received:', data.type);
            
            switch (data.type) {
                case 'authenticate':
                    handleAuthentication(connectionId, data, ws);
                    break;
                    
                case 'message':
                    handleMessage(data);
                    break;
                    
                case 'typing':
                    handleTyping(data);
                    break;
                    
                case 'read_receipt':
                    handleReadReceipt(data);
                    break;
                    
                case 'online_status':
                    handleOnlineStatus(data);
                    break;
                    
                case 'ping':
                    ws.send(JSON.stringify({ type: 'pong' }));
                    break;
            }
        } catch (error) {
            console.error('Error processing message:', error);
        }
    });
    
    ws.on('close', () => {
        console.log(`Connection closed: ${connectionId}`);
        connections.delete(connectionId);
        
        // Remove from user connections
        for (const [userKey, sockets] of userConnections.entries()) {
            const index = sockets.indexOf(ws);
            if (index > -1) {
                sockets.splice(index, 1);
                if (sockets.length === 0) {
                    userConnections.delete(userKey);
                    
                    // Broadcast offline status
                    const [userType, userId] = userKey.split(':');
                    broadcastToUser('admin', userId, {
                        type: 'online_status',
                        user: { type: userType, id: parseInt(userId) },
                        isOnline: false
                    });
                }
                break;
            }
        }
    });
    
    ws.on('error', (error) => {
        console.error('WebSocket error:', error);
    });
});

function handleAuthentication(connectionId, data, ws) {
    const { userType, userId } = data;
    const userKey = getUserKey(userType, userId);
    
    // Store connection
    if (!userConnections.has(userKey)) {
        userConnections.set(userKey, []);
    }
    userConnections.get(userKey).push(ws);
    
    // Send connection confirmation
    ws.send(JSON.stringify({
        type: 'authenticated',
        connectionId: connectionId,
        timestamp: new Date().toISOString()
    }));
    
    // Broadcast online status (if admin, notify users)
    if (userType === 'admin') {
        // Notify all users connected to this admin
        // This would require fetching all users for this admin
    } else {
        // Notify admin that user is online
        broadcastToUser('admin', data.adminId, {
            type: 'online_status',
            user: { type: userType, id: userId },
            isOnline: true
        });
    }
}

function handleMessage(data) {
    const { from, to, message, messageId, timestamp } = data;
    
    // Store message in database (via HTTP API)
    // For now, just broadcast
    
    // Broadcast to recipient
    broadcastToUser(to.type, to.id, {
        type: 'message',
        from: from,
        message: message,
        messageId: messageId,
        timestamp: timestamp || new Date().toISOString()
    });
    
    // Send delivery confirmation to sender
    broadcastToUser(from.type, from.id, {
        type: 'message_delivered',
        messageId: messageId,
        timestamp: new Date().toISOString()
    });
}

function handleTyping(data) {
    const { from, to, isTyping } = data;
    
    // Update typing status
    const typingKey = `${from.type}:${from.id}:${to.type}:${to.id}`;
    typingStatus.set(typingKey, {
        isTyping: isTyping,
        lastUpdate: Date.now()
    });
    
    // Broadcast typing status
    broadcastTypingStatus(from.type, from.id, to.type, to.id, isTyping);
}

function handleReadReceipt(data) {
    const { reader, messageId, chatPartner } = data;
    
    // Broadcast read receipt to sender
    broadcastToUser(chatPartner.type, chatPartner.id, {
        type: 'read_receipt',
        messageId: messageId,
        reader: reader,
        timestamp: new Date().toISOString()
    });
}

function handleOnlineStatus(data) {
    const { user, isOnline } = data;
    
    // Broadcast online status to relevant parties
    if (user.type === 'penghuni') {
        // Notify admin
        // Need admin ID - would come from session or database
    }
}

// Clean up old typing status every minute
setInterval(() => {
    const now = Date.now();
    for (const [key, status] of typingStatus.entries()) {
        if (now - status.lastUpdate > 5000) { // 5 seconds
            const [fromType, fromId, toType, toId] = key.split(':');
            if (status.isTyping) {
                broadcastTypingStatus(fromType, fromId, toType, toId, false);
            }
            typingStatus.delete(key);
        }
    }
}, 60000);

const PORT = process.env.WS_PORT || 3001;
server.listen(PORT, () => {
    console.log(`WebSocket server running on port ${PORT}`);
});

module.exports = { wss, server };