// Nithic on-device AI-text-detection engine (offscreen document).
//
// Runs the same 2-model ensemble as the cloud, but entirely in the browser so
// page text never leaves the device:
//   chatgpt-detector-roberta (precision) + e5-small-lora (recall), averaged.
//
// transformers.js + onnxruntime-web need a Document for WebGPU/WASM, which MV3
// service workers don't provide — hence this offscreen document. The background
// service worker creates it on demand and relays infer requests here.
//
// Model weights are NOT bundled (~153MB); they download once from our own infra
// (api.nithic.ai/v1/models) and cache in the browser's Cache API.

// IMPORTANT: transformers.js is loaded LAZILY via dynamic import() inside
// ensureTransformers(), NOT statically at the top of this module. A static
// import of the 470KB bundle must fully evaluate before the onMessage listener
// (at the bottom) registers — so the background would fire a message before the
// listener exists ("Could not establish connection. Receiving end does not
// exist"), and any import error would kill the whole module silently. Loading
// it lazily makes the listener register in microseconds and turns import
// failures into catchable, reportable errors.
let pipeline = null;
let env = null;
let transformersPromise = null;

async function ensureTransformers() {
  if (pipeline) return;
  if (!transformersPromise) {
    transformersPromise = (async () => {
      // transformers.min.js is the self-contained build (inlines onnxruntime).
      // The .web.min build externalizes onnxruntime-common/web as bare module
      // specifiers, which the browser can't resolve without a bundler.
      const mod = await import('./lib/transformers/transformers.min.js');
      pipeline = mod.pipeline;
      env = mod.env;
      env.allowLocalModels  = false;
      env.allowRemoteModels = true;
      env.useBrowserCache   = true;
      env.remoteHost        = 'https://api.nithic.ai';
      env.remotePathTemplate = 'v1/models/{model}';
      env.backends = env.backends || {};
      env.backends.onnx = env.backends.onnx || {};
      env.backends.onnx.wasm = env.backends.onnx.wasm || {};
      env.backends.onnx.wasm.wasmPaths = chrome.runtime.getURL('lib/transformers/');
      env.backends.onnx.wasm.proxy = false;
      env.backends.onnx.wasm.numThreads = 1; // offscreen docs aren't cross-origin isolated
      logEvent('transformers.js loaded');
    })().catch(e => {
      transformersPromise = null;
      logEvent('transformers.js load FAILED: ' + (e && e.message || e));
      throw e;
    });
  }
  return transformersPromise;
}

const REMOTE_HOST = 'https://api.nithic.ai';
const REMOTE_PATH_TPL = 'v1/models/{model}';

const MODELS = [
  { id: 'chatgpt-detector-roberta-ONNX',            aiLabels: ['chatgpt', 'fake', 'label_1'] },
  { id: 'e5-small-lora-ai-generated-detector-ONNX', aiLabels: ['label_1', 'ai', '1', 'generated'] }
];
// The single weight file we mirror per model (q8). Used for the update check.
const WEIGHT_FILE = 'onnx/model_quantized.onnx';

let detectorsPromise = null;
let device = 'wasm';        // default: WASM is reliable for int8/q8 (WebGPU isn't)
let loadState = 'idle';     // idle | loading | ready | error
let lastError = null;
const fileBytes = {};       // file → { loaded, total }
const eventLog = [];        // recent diagnostic events (capped)

function logEvent(msg) {
  // Timestamps are injected by the reader (Date.* is fine in a document).
  const line = `${new Date().toISOString().slice(11, 19)}  ${msg}`;
  eventLog.push(line);
  if (eventLog.length > 60) eventLog.shift();
  console.log('[Nithic engine]', msg);
}

// Overall download % weighted by bytes (so the big 120MB file dominates, the
// way a user expects — not a flat average across many tiny tokenizer files).
function aggregateProgress() {
  let loaded = 0, total = 0;
  for (const f of Object.values(fileBytes)) {
    if (f.total) { loaded += f.loaded; total += f.total; }
  }
  if (total === 0) return 0;
  return Math.round((loaded / total) * 100);
}

function bytesSummary() {
  let loaded = 0, total = 0;
  for (const f of Object.values(fileBytes)) { loaded += f.loaded || 0; total += f.total || 0; }
  return { loaded, total };
}

function postStatus(extra = {}) {
  chrome.runtime.sendMessage({
    target: 'nithic-bg', type: 'offscreen-status',
    ...statusSnapshot(), ...extra
  }).catch(() => {});
}

function statusSnapshot() {
  const b = bytesSummary();
  return {
    state: loadState, device, progress: aggregateProgress(),
    bytesLoaded: b.loaded, bytesTotal: b.total,
    models: MODELS.map(m => m.id),
    error: lastError,
    files: Object.entries(fileBytes).map(([file, v]) => ({
      file, loaded: v.loaded || 0, total: v.total || 0
    }))
  };
}

function progress_callback(p) {
  if (!p || !p.file) return;
  if (p.status === 'progress' || typeof p.loaded === 'number') {
    fileBytes[p.file] = { loaded: p.loaded || 0, total: p.total || fileBytes[p.file]?.total || 0 };
    postStatus();
  } else if (p.status === 'done' && fileBytes[p.file]) {
    fileBytes[p.file].loaded = fileBytes[p.file].total;
  } else if (p.status === 'initiate') {
    logEvent(`fetch ${p.file}`);
  }
}

async function getDetectors() {
  if (detectorsPromise) return detectorsPromise;
  loadState = 'loading';
  lastError = null;
  logEvent(`loading ensemble on ${device}`);
  postStatus();

  detectorsPromise = (async () => {
    await ensureTransformers();   // lazy-load the library (errors bubble up here)
    const built = [];
    for (const m of MODELS) {
      logEvent(`pipeline ${m.id} …`);
      const pipe = await pipeline('text-classification', m.id, {
        dtype: 'q8', device, progress_callback
      });
      built.push({ pipe, aiLabels: m.aiLabels });
      logEvent(`loaded ${m.id}`);
    }
    loadState = 'ready';
    logEvent(`ensemble READY (${built.length} models, ${device})`);
    postStatus();
    return built;
  })().catch(err => {
    loadState = 'error';
    lastError = String(err && err.message || err);
    detectorsPromise = null;     // allow retry
    logEvent(`ERROR: ${lastError}`);
    postStatus();
    throw err;
  });

  return detectorsPromise;
}

// Probe the Cache API for the weights — definitive "already downloaded" proof.
async function isCached() {
  try {
    const names = await caches.keys();
    for (const n of names) {
      const c = await caches.open(n);
      const keys = await c.keys();
      if (keys.some(r => r.url.includes('/v1/models/') && r.url.includes('model_quantized.onnx'))) return true;
    }
  } catch (_) {}
  return false;
}

// HEAD each weight file on the server and compare its size with what's cached.
// A size mismatch (or nothing cached) means an update is available.
async function checkUpdates() {
  const out = [];
  for (const m of MODELS) {
    const url = `${REMOTE_HOST}/${REMOTE_PATH_TPL.replace('{model}', m.id)}/${WEIGHT_FILE}`;
    let serverBytes = null, cachedBytes = null;
    try {
      const head = await fetch(url, { method: 'HEAD', cache: 'no-store' });
      serverBytes = Number(head.headers.get('content-length')) || null;
    } catch (e) { /* offline */ }
    try {
      const names = await caches.keys();
      for (const n of names) {
        const c = await caches.open(n);
        const match = await c.match(url);
        if (match) { cachedBytes = Number(match.headers.get('content-length')) || null; break; }
      }
    } catch (_) {}
    out.push({ model: m.id, serverBytes, cachedBytes, updateAvailable: !cachedBytes || (serverBytes && serverBytes !== cachedBytes) });
  }
  return out;
}

// Drop cached weights + loaded pipelines so the next load re-downloads.
async function resetEngine({ clearCache } = {}) {
  detectorsPromise = null;
  loadState = 'idle';
  lastError = null;
  for (const k of Object.keys(fileBytes)) delete fileBytes[k];
  if (clearCache) {
    try {
      const names = await caches.keys();
      for (const n of names) {
        const c = await caches.open(n);
        const keys = await c.keys();
        for (const r of keys) if (r.url.includes('/v1/models/')) await c.delete(r);
      }
      logEvent('cleared model cache');
    } catch (e) { logEvent('cache clear failed: ' + e.message); }
  }
  postStatus();
}

function aiProb(output, aiLabels) {
  const arr = Array.isArray(output) ? output : [output];
  const ai = arr.find(i => aiLabels.includes(String(i.label).toLowerCase()));
  if (ai) return ai.score;
  const human = arr.find(i => /^(human|real|label_0|0)$/i.test(i.label));
  if (human) return 1 - human.score;
  return arr[0] ? arr[0].score : 0;
}

async function scoreText(detectors, text) {
  let sum = 0;
  for (const d of detectors) {
    const out = await d.pipe(text, { top_k: 2 });
    sum += aiProb(out, d.aiLabels);
  }
  return sum / detectors.length;
}

function verdictFor(score) {
  if (score >= 0.65) return { verdict: 'ai_generated', confidence: score };
  if (score >= 0.45) return { verdict: 'mixed',        confidence: score };
  if (score >= 0.30) return { verdict: 'uncertain',    confidence: 1 - score };
  return { verdict: 'human', confidence: 1 - score };
}

chrome.runtime.onMessage.addListener((msg, sender, sendResponse) => {
  if (!msg || msg.target !== 'nithic-offscreen') return false;

  if (msg.type === 'status') {
    isCached().then(cached => sendResponse({ ...statusSnapshot(), cached, log: eventLog.slice(-40) }));
    return true;
  }

  if (msg.type === 'setDevice') {
    // Allow the diagnostics panel to force wasm vs webgpu (experimental).
    if (msg.device && msg.device !== device) {
      device = msg.device;
      logEvent(`device set to ${device}`);
      resetEngine();
    }
    sendResponse({ ok: true, device });
    return false;
  }

  if (msg.type === 'checkUpdates') {
    checkUpdates().then(r => sendResponse({ ok: true, results: r }))
                  .catch(e => sendResponse({ ok: false, error: String(e && e.message || e) }));
    return true;
  }

  if (msg.type === 'reload') {
    resetEngine({ clearCache: msg.clearCache }).then(() => getDetectors())
      .then(() => sendResponse({ ok: true, state: loadState }))
      .catch(e => sendResponse({ ok: false, error: String(e && e.message || e) }));
    return true;
  }

  if (msg.type === 'warmup') {
    // Kick off loading but DON'T await it — model load can take minutes and we
    // must not hold the caller (or the background service worker) open that
    // long. Progress is reported via postStatus + the 'status' query.
    getDetectors().catch(() => {});
    sendResponse({ ok: true, state: loadState, device, started: true });
    return false;
  }

  if (msg.type === 'infer') {
    const chunks = Array.isArray(msg.chunks) ? msg.chunks : [];
    getDetectors().then(async (detectors) => {
      const out = [];
      for (const c of chunks) {
        if (!c || !c.text || c.text.length < 50) { out.push({ id: c && c.id, verdict: 'uncertain', skipped: true }); continue; }
        try {
          const score = await scoreText(detectors, c.text.slice(0, 2000));
          const v = verdictFor(score);
          out.push({ id: c.id, cloudVerdict: v.verdict, cloudConfidence: Math.round(v.confidence * 100) / 100,
                     cloudDetails: { model_score: Math.round(score * 100) / 100, engine: 'on-device', device } });
        } catch (e) {
          out.push({ id: c.id, verdict: 'uncertain', error: String(e && e.message || e) });
        }
      }
      sendResponse({ chunks: out });
    }).catch(e => sendResponse({ chunks: [], error: String(e && e.message || e) }));
    return true;
  }

  return false;
});

logEvent('offscreen booted');
postStatus({ booted: true });
