let socket = null;
let sessionId = null;
let workerId = null;
let token = null;
let heartbeatInterval = null;
let deviceKeys = null;

function log(msg) {
  const el = document.getElementById('log');
  const line = document.createElement('div');
  line.textContent = new Date().toLocaleTimeString() + ' ' + msg;
  el.appendChild(line);
  el.scrollTop = el.scrollHeight;
  console.log(msg);
}

function parseSessionFromUrl() {
  const params = new URLSearchParams(window.location.search);
  sessionId = params.get('session_id');
  token = params.get('token');
  return { sessionId, token };
}

async function connectToSpace() {
  parseSessionFromUrl();
  if (!sessionId || !token) {
    log('No session_id or token in URL');
    return;
  }

  workerId = 'wk_' + Math.random().toString(36).slice(2, 10);
  log('Connecting as ' + workerId + ' to session ' + sessionId);

  const wsUrl = (location.protocol === 'https:' ? 'wss://' : 'ws://') + location.host + '/ws/worker/' + sessionId + '/' + workerId;
  socket = new WebSocket(wsUrl);

  socket.onopen = async () => {
    log('WebSocket connected');
    updateStatus('connecting');
    deviceKeys = await ReceiptCrypto.loadDeviceKeypair();
    const runtime = PhoneRuntime.selectBestRuntime();
    sendWorkerHello(runtime);
  };

  socket.onmessage = (event) => {
    const msg = JSON.parse(event.data);
    handleSocketMessage(msg);
  };

  socket.onclose = () => {
    log('WebSocket closed');
    updateStatus('offline');
    stopHeartbeat();
  };

  socket.onerror = (err) => {
    log('WebSocket error');
    updateStatus('offline');
  };

  document.getElementById('btnConnect').disabled = true;
  document.getElementById('btnDisconnect').disabled = false;
}

function sendWorkerHello(runtimeType) {
  if (!socket) return;
  socket.send(JSON.stringify({
    op: 'worker_hello',
    worker_id: workerId,
    runtime_type: runtimeType,
    device_public_key: deviceKeys ? deviceKeys.publicKey : null,
  }));
}

function detectDeviceInfo() {
  const battery = navigator.getBattery ? navigator.getBattery() : Promise.resolve({});
  battery.then(b => {
    document.getElementById('batteryValue').textContent = (b.level ? Math.round(b.level * 100) + '%' : 'unknown');
  }).catch(() => {});
  document.getElementById('networkValue').textContent = navigator.connection ? navigator.connection.effectiveType : 'unknown';
  document.getElementById('thermalValue').textContent = 'nominal'; // iOS Web API limitation
  document.getElementById('runtimeValue').textContent = PhoneRuntime.selectBestRuntime();
}

function advertiseCapabilities() {
  if (!socket) return;
  const caps = [
    { capability_name: 'iphone.text.embedding.private', runtime_type: PhoneRuntime.selectBestRuntime(), model_id: 'mock-embedding' },
    { capability_name: 'iphone.image.classify.local', runtime_type: PhoneRuntime.selectBestRuntime(), model_id: 'mock-image' },
    { capability_name: 'iphone.privacy.redact.local', runtime_type: PhoneRuntime.selectBestRuntime(), model_id: 'mock-redact' },
  ];
  socket.send(JSON.stringify({ op: 'capabilities', worker_id: workerId, capabilities: caps }));
  renderCapabilities(caps);
  log('Capabilities advertised');
}

function handleSocketMessage(msg) {
  switch (msg.op) {
    case 'worker_welcome':
      updateStatus('online');
      detectDeviceInfo();
      advertiseCapabilities();
      startHeartbeat();
      break;
    case 'job_offer':
      handleJobOffer(msg);
      break;
    case 'session_update':
      log('Session update: ' + JSON.stringify(msg));
      break;
    case 'error':
      log('Server error: ' + msg.error);
      break;
  }
}

async function handleJobOffer(job) {
  log('Job offered: ' + job.job_type + ' (' + job.job_id + ')');

  // Auto-accept for v1
  socket.send(JSON.stringify({ op: 'job_accept', job_id: job.job_id, worker_id: workerId }));

  const startTime = Date.now();
  let output = {};

  try {
    switch (job.job_type) {
      case 'text_embedding':
        output = await PhoneRuntime.runTextEmbedding(job.payload.text || '');
        break;
      case 'image_classification':
        output = await PhoneRuntime.runImageClassification(job.payload.image || '');
        break;
      case 'privacy_redaction':
        output = await PhoneRuntime.runPrivacyRedaction(job.payload.text || '');
        break;
      default:
        throw new Error('Unsupported job type: ' + job.job_type);
    }
  } catch (err) {
    log('Job failed: ' + err.message);
    socket.send(JSON.stringify({ op: 'job_result', job_id: job.job_id, worker_id: workerId, output: { error: err.message }, latency_ms: Date.now() - startTime }));
    return;
  }

  const latencyMs = Date.now() - startTime;
  const inputHash = await ReceiptCrypto.makeInputHash(job.payload);
  const outputHash = await ReceiptCrypto.makeOutputHash(output);
  const deviceSignature = await ReceiptCrypto.signReceipt({
    job_id: job.job_id,
    worker_id: workerId,
    input_hash: inputHash,
    output_hash: outputHash,
    latency_ms: latencyMs,
  });

  socket.send(JSON.stringify({
    op: 'job_result',
    job_id: job.job_id,
    worker_id: workerId,
    output: output,
    latency_ms: latencyMs,
    input_hash: inputHash,
    output_hash: outputHash,
    device_signature: deviceSignature,
  }));

  addReceiptCard(job.job_id, job.job_type, latencyMs, outputHash);
  log('Job completed: ' + job.job_id + ' in ' + latencyMs + 'ms');
}

function startHeartbeat() {
  heartbeatInterval = setInterval(() => {
    if (!socket || socket.readyState !== WebSocket.OPEN) return;
    socket.send(JSON.stringify({
      op: 'heartbeat',
      worker_id: workerId,
      battery_level: 0.8, // mock for v1
      thermal_state: 'nominal',
      jobs_completed: document.querySelectorAll('.receipt-card').length,
    }));
  }, 5000);
}

function stopHeartbeat() {
  if (heartbeatInterval) clearInterval(heartbeatInterval);
  heartbeatInterval = null;
}

function disconnectWorker() {
  stopHeartbeat();
  if (socket) {
    socket.send(JSON.stringify({ op: 'disconnect', worker_id: workerId, reason: 'user' }));
    socket.close();
    socket = null;
  }
  updateStatus('offline');
  document.getElementById('btnConnect').disabled = false;
  document.getElementById('btnDisconnect').disabled = true;
}

function updateStatus(status) {
  const badge = document.getElementById('statusBadge');
  badge.className = 'status-badge status-' + status;
  badge.textContent = status.charAt(0).toUpperCase() + status.slice(1);
}

function renderCapabilities(caps) {
  const el = document.getElementById('capabilitiesList');
  el.innerHTML = caps.map(c =>
    '<div class="job-card"><div class="job-type">' + c.capability_name + '</div><div style="font-size:12px;color:#64748b">' + c.runtime_type + ' / ' + c.model_id + '</div></div>'
  ).join('');
}

function addReceiptCard(jobId, jobType, latencyMs, outputHash) {
  const el = document.getElementById('receiptLog');
  const card = document.createElement('div');
  card.className = 'receipt-card';
  card.innerHTML = '<div style="font-weight:700;margin-bottom:4px;">' + jobType + '</div><div>Latency: ' + latencyMs + 'ms</div><div style="color:#059669">' + outputHash.slice(0, 24) + '...</div>';
  el.prepend(card);
}

// Auto-connect if URL has session params
if (window.location.search.includes('session_id')) {
  window.addEventListener('load', () => {
    setTimeout(connectToSpace, 500);
  });
}