const WebSocket = require('ws');
const http = require('http');

const SERVER_URL = process.env.SERVER_URL || 'wss://choruyt-ai-exocore.hf.space/cpu-pool';
const WORKER_ID = process.env.WORKER_ID || 'cpu' + (Math.floor(Math.random() * 1000));
const CPU_CORES = 2; // HF Spaces = 2 vCPU

console.log(`🚀 CPU Worker ${WORKER_ID} starting...`);

let ws;
let heartbeatInterval;

function connectToPool() {
    ws = new WebSocket(SERVER_URL);
    
    ws.on('open', () => {
        console.log(`✅ ${WORKER_ID} connected to CPU POOL`);
        ws.send(JSON.stringify({
            type: 'register',
            workerId: WORKER_ID,
            cpuCores: CPU_CORES,
            status: 'ready'
        }));
        
        // Heartbeat
        heartbeatInterval = setInterval(() => {
            if (ws.readyState === WebSocket.OPEN) {
                ws.send(JSON.stringify({
                    type: 'heartbeat',
                    workerId: WORKER_ID,
                    load: process.uptime() % 10 // Fake load
                }));
            }
        }, 5000);
    });

    ws.on('message', async (data) => {
        try {
            const task = JSON.parse(data.toString());
            if (task.type === 'cpu-task') {
                const result = await heavyCPUWork(task.data);
                ws.send(JSON.stringify({
                    type: 'task-complete',
                    taskId: task.taskId,
                    workerId: WORKER_ID,
                    result: result
                }));
            }
        } catch (e) {
            console.error('Task error:', e);
        }
    });

    ws.on('close', () => {
        console.log('🔌 Disconnected, reconnecting...');
        clearInterval(heartbeatInterval);
        setTimeout(connectToPool, 3000);
    });
}

// 🔥 HEAVY CPU WORK (for Qwen3 token processing)
function heavyCPUWork(data) {
    return new Promise((resolve) => {
        let result = 0;
        const start = Date.now();
        
        // Simulate Qwen3 token generation (CPU intensive)
        for (let i = 0; i < data.iterations || 10000000; i++) {
            result += Math.sqrt(i) * Math.sin(i / 1000);
        }
        
        resolve({
            tokensGenerated: data.iterations || 512,
            computationTime: Date.now() - start,
            result: result.toFixed(2)
        });
    });
}

// Auto-reconnect
connectToPool();

// Health check server
const server = http.createServer((req, res) => {
    res.writeHead(200);
    res.end(`CPU Worker ${WORKER_ID} OK | Cores: ${CPU_CORES}`);
});

server.listen(process.env.PORT || 7860, '0.0.0.0', () => {
    console.log(`✅ Worker ${WORKER_ID} listening on ${process.env.PORT || 7860}`);
});