// ImposterShield service worker (v0.3).
// Two detection tiers, kept strictly separate so we never claim proof we
// don't have:
//   PROVEN  - byte-level markers (C2PA AI stamps, tool names in metadata)
//             and AI-only image hosts. We can point at the evidence.
//   LIKELY  - OUR OWN detector, running locally in the browser: a SwinV2
//             AI-image classifier (ONNX via transformers.js). No third-party
//             API, no per-image cost, nothing leaves the machine. Shown as
//             "likely," never as fact.
//
// Toolbar light: red = a lot of AI, yellow = some, green = none found,
// gray = still checking. Green means "nothing found," not "guaranteed real."

var COLORS = {
  red: '#E5484D',
  yellow: '#FFB020',
  green: '#3DDC97',
  gray: '#5D667A'
};

var AI_MARKERS = [
  { s: 'trainedalgorithmicmedia', why: 'The file itself is stamped as AI-generated (its Content Credentials say so).' },
  { s: 'midjourney', why: 'Made with Midjourney. The file says so in its own metadata.' },
  { s: 'stable diffusion', why: 'Made with Stable Diffusion. The file says so in its own metadata.' },
  { s: 'stablediffusion', why: 'Made with Stable Diffusion. The file says so in its own metadata.' },
  { s: 'dall-e', why: 'Made with DALL-E. The file says so in its own metadata.' },
  { s: 'dall·e', why: 'Made with DALL-E. The file says so in its own metadata.' },
  { s: 'adobe firefly', why: 'Made with Adobe Firefly. The file says so in its own metadata.' },
  { s: 'ideogram', why: 'Made with Ideogram. The file says so in its own metadata.' },
  { s: 'leonardo.ai', why: 'Made with Leonardo AI. The file says so in its own metadata.' },
  { s: 'black forest labs', why: 'Made with FLUX (Black Forest Labs). The file says so in its own metadata.' },
  { s: 'openai', why: 'Carries an OpenAI generation record in its metadata.' }
];

var AI_HOSTS = [
  { s: 'oaiusercontent.com', why: 'Served from OpenAI’s image host.' },
  { s: 'mj-gallery.com', why: 'Served from Midjourney’s gallery host.' },
  { s: 'cdn.midjourney.com', why: 'Served from Midjourney’s image host.' },
  { s: 'replicate.delivery', why: 'Served from an AI model host (Replicate).' },
  { s: 'ideogram.ai', why: 'Served from Ideogram’s image host.' },
  { s: 'cdn.leonardo.ai', why: 'Served from Leonardo AI’s image host.' }
];

var MAX_BYTES = 10 * 1024 * 1024;
var FETCH_TIMEOUT_MS = 6000;

var CLASSIFY_MIN_DIM = 120;
var CLASSIFY_MAX_FRESH_PER_SCAN = 8;
var CLASSIFY_THRESHOLD = 0.75;
var classifyCache = new Map(); // url -> {score} | {error}
var CLASSIFY_CACHE_MAX = 500;

// ---- Local detector, hosted in an offscreen document ----
// MV3 service workers cannot use dynamic import() (the ONNX runtime needs
// it), so the detector lives in offscreen.html and we message it. Everything
// it loads comes from us: runtime vendored in the extension, weights from
// models.impostershield.com. Model: ViT fine-tune, Apache-2.0, REAL/FAKE.

var modelState = { state: 'idle', pct: 0 }; // idle | loading | ready | failed

function publishModelState() {
  chrome.storage.session.set({ modelState: modelState });
}

var offscreenReady = null;
function ensureOffscreen() {
  // chrome.offscreen is missing on older Chrome. Reading .hasDocument() off undefined throws a
  // SYNCHRONOUS TypeError, which aborts this whole script - so the extension dies silently and
  // never even reports its own install. Fail as a rejected promise instead of a hard throw.
  if (!chrome.offscreen || !chrome.offscreen.hasDocument) {
    return Promise.reject(new Error('This browser is too old to run the detector (needs Chrome 116+).'));
  }
  if (!offscreenReady) {
    offscreenReady = chrome.offscreen.hasDocument().then(function (has) {
      if (has) return;
      return chrome.offscreen.createDocument({
        url: 'offscreen.html',
        reasons: ['WORKERS'],
        justification: 'Runs the on-device AI image detector (WebAssembly/WebGPU), which MV3 service workers cannot host.'
      });
    }).catch(function (e) {
      // Racing createDocument calls throw "already exists" - that is fine.
      if (String(e).indexOf('already exists') === -1) { offscreenReady = null; throw e; }
    });
  }
  return offscreenReady;
}

function classifyImage(url) {
  if (classifyCache.has(url)) return Promise.resolve(classifyCache.get(url));
  if (url.indexOf('http') !== 0) return Promise.resolve({ skipped: true });
  return ensureOffscreen().then(function () {
    return chrome.runtime.sendMessage({ target: 'offscreen', type: 'IS_CLASSIFY', url: url });
  }).then(function (out) {
    out = out || { error: 'noresponse' };
    if (typeof out.score === 'number' || out.error === 'labels') {
      if (classifyCache.size >= CLASSIFY_CACHE_MAX) classifyCache.clear();
      classifyCache.set(url, out);
    }
    return out;
  }).catch(function () {
    return { error: 'model' };
  });
}

// ---- Marker tier ----

function latin1(buf) {
  try { return new TextDecoder('latin1').decode(buf).toLowerCase(); }
  catch (e) { return ''; }
}

function checkUrl(url) {
  var u = url.toLowerCase();
  for (var i = 0; i < AI_HOSTS.length; i++) {
    if (u.indexOf(AI_HOSTS[i].s) !== -1) return AI_HOSTS[i].why;
  }
  return null;
}

function checkBytes(text) {
  for (var i = 0; i < AI_MARKERS.length; i++) {
    if (text.indexOf(AI_MARKERS[i].s) !== -1) return AI_MARKERS[i].why;
  }
  return null;
}

function checkSigned(text) {
  return text.indexOf('c2pa') !== -1 || text.indexOf('jumb') !== -1;
}

function fetchImage(url) {
  var ctrl = new AbortController();
  var timer = setTimeout(function () { ctrl.abort(); }, FETCH_TIMEOUT_MS);
  return fetch(url, { signal: ctrl.signal, credentials: 'omit' })
    .then(function (r) {
      var len = parseInt(r.headers.get('content-length') || '0', 10);
      if (len > MAX_BYTES) throw new Error('too big');
      return r.arrayBuffer();
    })
    .then(function (buf) {
      clearTimeout(timer);
      if (buf.byteLength > MAX_BYTES) throw new Error('too big');
      return buf;
    })
    .catch(function (e) {
      clearTimeout(timer);
      throw e;
    });
}

function analyzeImage(img) {
  var urlWhy = checkUrl(img.url);
  if (urlWhy) {
    return Promise.resolve({ url: img.url, w: img.w, h: img.h, ai: true, why: urlWhy, signed: false });
  }
  return fetchImage(img.url).then(function (buf) {
    var text = latin1(buf);
    var why = checkBytes(text);
    if (why) return { url: img.url, w: img.w, h: img.h, ai: true, why: why, signed: false };
    if (checkSigned(text)) return { url: img.url, w: img.w, h: img.h, ai: false, why: null, signed: true };
    return { url: img.url, w: img.w, h: img.h, ai: false, why: null, signed: false };
  }).catch(function () {
    return { url: img.url, w: img.w, h: img.h, ai: false, why: null, signed: false, unreadable: true };
  });
}

// ---- Verdict + light ----

function verdictFor(provenCount, likelyCount, pageDeclared) {
  var score = provenCount * 2 + likelyCount;
  if (score >= 4 || provenCount >= 3) return 'red';
  if (score >= 1 || pageDeclared) return 'yellow';
  return 'green';
}

function drawDot(color, size) {
  var c = new OffscreenCanvas(size, size);
  var ctx = c.getContext('2d');
  ctx.clearRect(0, 0, size, size);
  ctx.fillStyle = color;
  ctx.beginPath();
  ctx.arc(size / 2, size / 2, size * 0.42, 0, Math.PI * 2);
  ctx.fill();
  ctx.lineWidth = Math.max(1, size * 0.08);
  ctx.strokeStyle = 'rgba(10,12,16,.55)';
  ctx.stroke();
  return ctx.getImageData(0, 0, size, size);
}

function setLight(tabId, verdict, count) {
  var color = COLORS[verdict] || COLORS.gray;
  try {
    chrome.action.setIcon({
      tabId: tabId,
      imageData: { 16: drawDot(color, 16), 32: drawDot(color, 32) }
    });
    chrome.action.setBadgeBackgroundColor({ tabId: tabId, color: color });
    chrome.action.setBadgeText({ tabId: tabId, text: count > 0 ? String(count) : '' });
  } catch (e) { /* tab may be gone */ }
}

function saveResults(tabId, results) {
  var obj = {};
  obj['tab:' + tabId] = results;
  chrome.storage.session.set(obj);
}

function pool(items, limit, fn) {
  var i = 0;
  var out = new Array(items.length);
  function next() {
    if (i >= items.length) return Promise.resolve();
    var idx = i++;
    return fn(items[idx]).then(function (r) { out[idx] = r; return next(); });
  }
  var starters = [];
  for (var k = 0; k < Math.min(limit, items.length); k++) starters.push(next());
  return Promise.all(starters).then(function () { return out; });
}

function handleScan(msg, sender) {
  var tabId = sender.tab.id;
  var gen = ((msg.signals && msg.signals.generator) || '').toLowerCase();
  var pageDeclared = /midjourney|stable diffusion|dall-e|firefly|ideogram/.test(gen);

  pool(msg.images || [], 4, analyzeImage).then(function (results) {
    var proven = results.filter(function (r) { return r.ai; });
    var signed = results.filter(function (r) { return r.signed; });
    var clean = results.filter(function (r) { return !r.ai && !r.signed && !r.unreadable; });

    var candidates = [];
    var fresh = 0;
    clean.forEach(function (r) {
      if (r.w < CLASSIFY_MIN_DIM || r.h < CLASSIFY_MIN_DIM) return;
      if (classifyCache.has(r.url)) {
        candidates.push(r);
      } else if (fresh < CLASSIFY_MAX_FRESH_PER_SCAN) {
        candidates.push(r);
        fresh++;
      }
    });

    // Classify one at a time: local inference saturates the GPU/CPU anyway.
    pool(candidates, 1, function (r) {
      return classifyImage(r.url).then(function (v) { return { r: r, v: v }; });
    }).then(function (scored) {
      var likely = [];
      var topScore = null;
      var classifiedCount = 0;
      scored.forEach(function (s) {
        if (s.v && typeof s.v.score === 'number') {
          classifiedCount++;
          if (topScore === null || s.v.score > topScore) topScore = s.v.score;
          if (s.v.score >= CLASSIFY_THRESHOLD) {
            likely.push({
              url: s.r.url,
              score: s.v.score,
              why: 'Our detector rates this ' + Math.round(s.v.score * 100) + '% likely AI-made. Its metadata carries no proof either way.'
            });
          }
        }
      });

      var verdict = verdictFor(proven.length, likely.length, pageDeclared);
      var summary = {
        verdict: verdict,
        checked: results.length,
        aiCount: proven.length,
        likelyCount: likely.length,
        signedCount: signed.length,
        pageDeclared: pageDeclared,
        detector: modelState.state,
        detectorPct: modelState.pct,
        classified: classifiedCount,
        topScore: topScore,
        findings: proven.map(function (r) { return { url: r.url, why: r.why, kind: 'proven', score: null }; })
          .concat(likely.map(function (l) { return { url: l.url, why: l.why, kind: 'likely', score: l.score }; })),
        url: sender.tab.url || '',
        at: Date.now()
      };

      setLight(tabId, verdict, proven.length + likely.length);
      saveResults(tabId, summary);

      if (proven.length > 0 || likely.length > 0) {
        chrome.tabs.sendMessage(tabId, {
          type: 'IS_MARK',
          proven: proven.map(function (r) { return r.url; }),
          likely: likely.map(function (l) { return l.url; }),
          color: verdict === 'red' ? COLORS.red : COLORS.yellow
        }).catch(function () {});
      }
    });
  });
}

chrome.runtime.onMessage.addListener(function (msg, sender, sendResponse) {
  if (msg && msg.target === 'offscreen') return; // not for us

  if (msg && msg.target === 'sw' && msg.type === 'IS_MODEL_STATE') {
    modelState = { state: msg.state, pct: msg.pct, error: msg.error };
    publishModelState();
    if (msg.state === 'ready') {
      // Model just came online: re-check whatever the user is looking at.
      chrome.tabs.query({ active: true }, function (tabs) {
        tabs.forEach(function (t) {
          chrome.tabs.sendMessage(t.id, { type: 'IS_RESCAN' }).catch(function () {});
        });
      });
    }
    return;
  }

  if (msg && msg.type === 'IS_SCAN' && sender.tab && sender.tab.id != null) {
    handleScan(msg, sender);
    return;
  }
  if (msg && msg.type === 'IS_GET') {
    chrome.storage.session.get(['tab:' + msg.tabId, 'modelState'], function (data) {
      var res = data['tab:' + msg.tabId] || null;
      var ms = data.modelState || modelState;
      if (res) { res.detector = ms.state; res.detectorPct = ms.pct; res.detectorError = ms.error; }
      sendResponse(res ? res : (ms.state !== 'idle' ? { onlyModel: true, detector: ms.state, detectorPct: ms.pct, detectorError: ms.error } : null));
    });
    return true;
  }
  if (msg && msg.type === 'IS_RESCAN') {
    chrome.tabs.sendMessage(msg.tabId, { type: 'IS_RESCAN' }).catch(function () {});
    return;
  }
  if (msg && msg.type === 'IS_WARM') {
    ensureOffscreen().then(function () {
      chrome.runtime.sendMessage({ target: 'offscreen', type: 'IS_WARM' }).catch(function () {});
    }).catch(function () {});
    return;
  }
});

chrome.tabs.onUpdated.addListener(function (tabId, info) {
  if (info.status === 'loading') {
    setLight(tabId, 'gray', 0);
    chrome.storage.session.remove('tab:' + tabId);
  }
});

// Anonymous daily heartbeat: random install id + version, nothing else.
function dailyPing() {
  chrome.storage.local.get(['installId', 'lastPingDay'], function (d) {
    var id = d.installId;
    if (!id) {
      id = crypto.randomUUID();
      chrome.storage.local.set({ installId: id });
    }
    var today = new Date().toISOString().slice(0, 10);
    if (d.lastPingDay === today) return;
    fetch('https://impostershield.com/api/ping', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ id: id, version: chrome.runtime.getManifest().version })
    }).then(function () {
      chrome.storage.local.set({ lastPingDay: today });
    }).catch(function () {});
  });
}

// ORDER MATTERS. dailyPing() goes first so an install always reports itself even if the detector
// cannot start on this machine. When these were the other way round, anything that broke during
// offscreen setup also erased our only evidence that the person installed at all - which reads as
// "nobody installed it" instead of "it crashed for them".
dailyPing();

// Start the offscreen host (which begins the one-time model download) as
// soon as the extension wakes up.
try {
  ensureOffscreen().catch(function () {});
} catch (e) { /* never let detector startup take down the service worker */ }
