const express = require('express');
const WebSocket = require('ws');
const os = require('os');
const app = express();

app.use(express.json());

const OLLAMA_URL = 'http://127.0.0.1:11434';

function getCpuMetrics() {
  const cpus = os.cpus();
  const totalLoad = cpus.reduce((sum, cpu) => sum + (1 - cpu.times.idle / Object.values(cpu.times).reduce((a, b) => a + b, 0)), 0) / cpus.length;
  return { id: process.env.WORKER_ID || 'cpu1', cpuLoad: Math.round(totalLoad * 100) };
}

let broker = null;
let taskQueue = [];
let currentTask = null;

function connectBroker() {
  const brokerHost = process.env.BROKER_HOST || 'wss://choruyt-broker-exocore.hf.space';
  broker = new WebSocket(brokerHost);
  
  broker.on('open', () => {
    console.log(`✅ Connected to Broker as ${getCpuMetrics().id}`);
    broker.send(JSON.stringify({ type: 'register', role: 'cpu_worker', metrics: getCpuMetrics() }));
  });
  
  broker.on('message', async (data) => {
    const msg = JSON.parse(data.toString());
    if (msg.type === 'cpu_task') {
      if (!currentTask) {
        currentTask = msg;
        const result = await executeInferenceTask(msg.task);
        if (broker.readyState === WebSocket.OPEN) {
          broker.send(JSON.stringify({ type: 'cpu_result', taskId: msg.task.id, result, from: getCpuMetrics().id }));
        }
        currentTask = null;
        processTaskQueue();
      } else {
        taskQueue.push(msg);
      }
    }
  });
  
  broker.on('close', () => setTimeout(connectBroker, 3000));
}

// ITO NA YUNG LLM CALL SA LOOB NG WORKER!
async function executeInferenceTask(task) {
  if (task.type !== 'inference') return { error: "Unknown task type" };
  
  try {
    const response = await fetch(`${OLLAMA_URL}/api/generate`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ model: task.model, prompt: task.prompt, stream: false })
    });
    
    const data = await response.json();
    return { success: true, response: data.response, from: getCpuMetrics().id };
  } catch (e) {
    return { success: false, error: `Ollama Worker Error: ${e.message}` };
  }
}

function processTaskQueue() {
  if (taskQueue.length > 0 && !currentTask && broker?.readyState === WebSocket.OPEN) {
    broker.send(JSON.stringify(taskQueue.shift()));
  }
}

app.get('/', (req, res) => res.json({ status: 'OLLAMA Worker Active', workerId: getCpuMetrics().id }));

setInterval(() => {
  if (broker?.readyState === WebSocket.OPEN) {
    broker.send(JSON.stringify({ type: 'heartbeat', metrics: getCpuMetrics() }));
  }
}, 30000);

connectBroker();
app.listen(7860, () => console.log(`🚀 Ollama Worker ready on 7860`));