background.js

importScripts('shared.js');
importScripts('dataService.js'); // 复用存储管理（getStorageUsage / pruneOldData）做后台自动剪枝

const FLUSH_COUNT = 100;
// 存储自动剪枝：独立 alarm 周期触发（不能用"每 N 次 flush 计一次"——worker 回收后模块级计数器归零）
const PRUNE_ALARM = 'storagePrune';
const PRUNE_CHECK_MINUTES = 30;
const PRUNE_KEEP_DAYS = 30;
const PRUNE_THRESHOLD = 0.8;

let pendingData = {};
let pendingRawData = {};
let pendingTotal = 0;
let pendingDownloads = {};
let pendingDownloadItems = [];
let pendingDownloadCount = 0;
let excludedDomains = new Set();
let requestCounter = 0;

// ── flush 串行化 ────────────────────────────────────────────
// 并发 flush 会各自读到同一份 storage 基线再各自写回，后写覆盖前写 → 前一批字节凭空消失。
// 所有 flush 排进同一条 Promise 链，保证"读-改-写"原子。
let flushChain = Promise.resolve();
let browseFlushQueued = false;
let downloadFlushQueued = false;

/**
 * 把 flush 任务排进串行链；链上任意任务抛错都不会阻塞后续任务。
 * @param {() => Promise<void>} task 任务
 * @returns {Promise<void>}
 */
function enqueueFlush(task) {
  flushChain = flushChain
    .then(task)
    .catch(e => console.warn('[FlowWatch] flush 任务失败，已回填待下次重试', e));
  return flushChain;
}

/**
 * 从 URL 提取主机名与根域名（非 http 协议返回 null）
 * @param {string} url 请求 URL
 * @returns {{hostname: string, root: string}|null}
 */
function getDomainInfo(url) {
  try {
    const u = new URL(url);
    if (!u.protocol.startsWith('http')) return null;
    const hostname = u.hostname;
    const root = getRootDomain(hostname);
    return { hostname, root };
  } catch {
    return null;
  }
}

/**
 * 把一批计数按 key 累加进目标对象（flush 失败回填专用，不能用展开覆盖）
 * @param {Object<string, number>} target 目标累加对象
 * @param {Object<string, number>} src 待回填批次
 */
function mergeCounterInto(target, src) {
  for (const [key, value] of Object.entries(src)) {
    target[key] = (target[key] || 0) + value;
  }
}

/**
 * 把当前内存待写队列全量写一份到 chrome.storage.session 作崩溃恢复快照；
 * 队列已空则删除快照。flush 成功后也走这里——直接 remove 会把 flush 期间
 * 新采集到的数据一并抹掉，而那时它尚无落库副本。
 */
async function writePendingSnapshot() {
  try {
    const empty = pendingTotal === 0 && pendingDownloadCount === 0 &&
      pendingDownloadItems.length === 0;
    if (empty) {
      await chrome.storage.session.remove([
        STORAGE_PENDING, STORAGE_PENDING_RAW, STORAGE_PENDING_DL, STORAGE_PENDING_DL_RAW
      ]);
      return;
    }
    await chrome.storage.session.set({
      [STORAGE_PENDING]: pendingData,
      [STORAGE_PENDING_RAW]: pendingRawData,
      [STORAGE_PENDING_DL]: pendingDownloads,
      [STORAGE_PENDING_DL_RAW]: pendingDownloadItems
    });
  } catch (e) {
    // 快照失败不影响采集，下个周期重试
  }
}

/**
 * 将积累的浏览流量 flush 到 chrome.storage.local（失败回填 pending 下次重试）。
 * 已排队时直接返回同一条链，避免并发 flush 各自覆盖对方写入。
 * @returns {Promise<void>}
 */
function flushData() {
  if (browseFlushQueued || pendingTotal === 0) return flushChain;
  browseFlushQueued = true;
  return enqueueFlush(async () => {
    browseFlushQueued = false;
    if (pendingTotal === 0) return;
    // 取走当前批次并立即清空，让采集侧继续往新的空对象累加
    const data = pendingData;
    const rawData = pendingRawData;
    const total = pendingTotal;
    pendingData = {};
    pendingRawData = {};
    pendingTotal = 0;
    const todayKey = getDateKey(new Date());
    const aggKey = DAILY_PREFIX + todayKey;
    const rawKey = RAW_PREFIX + todayKey;
    try {
      const result = await chrome.storage.local.get([aggKey, rawKey, STORAGE_DATES]);
      const storedAgg = result[aggKey] || {};
      for (const [domain, bytes] of Object.entries(data)) {
        storedAgg[domain] = (storedAgg[domain] || 0) + bytes;
      }
      const storedRaw = result[rawKey] || {};
      for (const [hostname, bytes] of Object.entries(rawData)) {
        storedRaw[hostname] = (storedRaw[hostname] || 0) + bytes;
      }
      const toSet = { [aggKey]: storedAgg, [rawKey]: storedRaw };
      let dates = result[STORAGE_DATES] || [];
      if (!dates.includes(todayKey)) {
        dates = [...dates, todayKey].sort();
        toSet[STORAGE_DATES] = dates;
      }
      await chrome.storage.local.set(toSet);
      // 落库成功：快照改写为"当前剩余的待写队列"，而非无脑删除
      await writePendingSnapshot();
    } catch (e) {
      // 回填必须累加：flush 期间采集侧可能已往 pending 里加了新数据，
      // 用展开覆盖会把那部分新数据连同本批一起丢掉。
      mergeCounterInto(pendingData, data);
      mergeCounterInto(pendingRawData, rawData);
      pendingTotal += total;
      await writePendingSnapshot();
    }
  });
}

/**
 * 将积累的下载流量 flush 到 chrome.storage.local（失败回填 pending 下次重试）
 * @returns {Promise<void>}
 */
function flushDownloads() {
  if (downloadFlushQueued || pendingDownloadCount === 0) return flushChain;
  downloadFlushQueued = true;
  return enqueueFlush(async () => {
    downloadFlushQueued = false;
    if (pendingDownloadCount === 0) return;
    const data = pendingDownloads;
    const items = pendingDownloadItems;
    const count = pendingDownloadCount;
    pendingDownloads = {};
    pendingDownloadItems = [];
    pendingDownloadCount = 0;
    const todayKey = getDateKey(new Date());
    const aggKey = DL_DAILY_PREFIX + todayKey;
    const rawKey = DL_RAW_PREFIX + todayKey;
    try {
      const result = await chrome.storage.local.get([aggKey, rawKey, STORAGE_DATES]);
      const storedAgg = result[aggKey] || {};
      for (const [domain, bytes] of Object.entries(data)) {
        storedAgg[domain] = (storedAgg[domain] || 0) + bytes;
      }
      const storedRaw = result[rawKey] || [];
      storedRaw.push(...items);
      const toSet = { [aggKey]: storedAgg, [rawKey]: storedRaw };
      let dates = result[STORAGE_DATES] || [];
      if (!dates.includes(todayKey)) {
        dates = [...dates, todayKey].sort();
        toSet[STORAGE_DATES] = dates;
      }
      await chrome.storage.local.set(toSet);
      await writePendingSnapshot();
    } catch (e) {
      // 同 flushData：累加回填，旧批次排在时间序前面
      mergeCounterInto(pendingDownloads, data);
      pendingDownloadItems = [...items, ...pendingDownloadItems];
      pendingDownloadCount += count;
      await writePendingSnapshot();
    }
  });
}

// ── Service Worker 冷启动恢复 ────────────────────────────────
// MV3 service worker 空闲约 30s 即被回收、下次事件唤醒时内存状态全丢，
// 而 chrome.runtime.onStartup 只在**浏览器启动**时触发——唤醒不走它。
// 因此恢复必须挂在每次 worker 求值时（本 IIFE），不能只在 onStartup 里做。
// await 会延长 worker 生命周期，保证恢复完成后再处理引发唤醒的事件；
// 失败(如 alarms 上下文无 session 权限)只降级为不恢复，绝不影响后续 flush。
let restoreReady = (async () => {
  try {
    const ex = await chrome.storage.local.get(EXCLUDED_KEY);
    excludedDomains = new Set(ex[EXCLUDED_KEY] || []);
  } catch (e) {
    excludedDomains = new Set();
  }
  // 分组规则必须加载，采集层才能用「排除组内任一成员 = 整组排除」的口径判断。
  // 否则排除了 hdslb.com，bilibili.com 的 CDN 流量仍会照常入库并占配额。
  try {
    await initDomainGroups();
  } catch (e) {
    // 加载失败时 resolveGroup 退化为恒等，「整组排除」降级为「精确匹配排除」
  }
  try {
    const result = await chrome.storage.session.get([
      STORAGE_PENDING, STORAGE_PENDING_RAW, STORAGE_PENDING_DL, STORAGE_PENDING_DL_RAW
    ]);
    if (result[STORAGE_PENDING]) pendingData = result[STORAGE_PENDING];
    if (result[STORAGE_PENDING_RAW]) pendingRawData = result[STORAGE_PENDING_RAW];
    if (result[STORAGE_PENDING_DL]) pendingDownloads = result[STORAGE_PENDING_DL];
    if (result[STORAGE_PENDING_DL_RAW]) pendingDownloadItems = result[STORAGE_PENDING_DL_RAW];
    // 计数语义与运行时一致：pendingTotal 是"待 flush 的请求条数"，不是字节和。
    // 旧写法把字节和赋给计数，会让恢复后第一个请求就误触发 flush（100 字节即达标）。
    pendingTotal = Object.keys(pendingData).length;
    pendingDownloadCount = pendingDownloadItems.length;
  } catch (e) {
    // 无已存快照：保持空状态
  }
})();

chrome.storage.onChanged.addListener((changes, area) => {
  if (area !== 'local') return;
  if (changes[EXCLUDED_KEY]) {
    excludedDomains = new Set(changes[EXCLUDED_KEY].newValue || []);
  }
  // 分组规则变更后重载，让后续请求按新规则判断整组排除
  if (changes[DOMAIN_GROUP_KEY]) initDomainGroups().catch(() => {});
});

// dashboard「重置数据」后清空内存待写队列：只删 storage 不够，
// 内存里已积累的数据会在下一次 alarm 被写回，导致刚重置就有数据冒出来。
chrome.runtime.onMessage.addListener((msg) => {
  if (!msg || msg.type !== MSG_RESET_PENDING) return;
  pendingData = {};
  pendingRawData = {};
  pendingTotal = 0;
  pendingDownloads = {};
  pendingDownloadItems = [];
  pendingDownloadCount = 0;
  requestCounter = 0;
  return false; // 不需要响应，避免打开异步响应通道
});

chrome.alarms.create('trafficFlush', { periodInMinutes: 1 });
chrome.alarms.onAlarm.addListener((alarm) => {
  if (alarm.name === 'trafficFlush') {
    if (pendingTotal > 0) flushData();
    if (pendingDownloadCount > 0) flushDownloads();
    return;
  }
  // 存储压力检查走独立 alarm：用"每 N 次 flush 计一次"的模块级计数器不可靠，
  // worker 空闲回收后计数器归零，低流量下永远攒不满 N 次 → 配额写满后 flush 永久失败、数据静默停记。
  if (alarm.name === PRUNE_ALARM) checkStoragePressure();
});

/**
 * 确保剪枝 alarm 已注册。只在 alarm 不存在时创建，避免每次 worker 求值都
 * 重新排期、把下一次触发时间不断往后推（低流量下 alarm 间隔可能远大于周期）。
 */
async function ensurePruneAlarm() {
  try {
    const existing = await chrome.alarms.get(PRUNE_ALARM);
    if (existing) return;
    chrome.alarms.create(PRUNE_ALARM, { periodInMinutes: PRUNE_CHECK_MINUTES });
  } catch (e) {
    // 查询失败时退回一次性创建（重复 create 同名 alarm 本身是幂等的）
    chrome.alarms.create(PRUNE_ALARM, { periodInMinutes: PRUNE_CHECK_MINUTES });
  }
}

/** 存储占用超阈值时清理最旧数据；失败只记日志，不影响采集 */
async function checkStoragePressure() {
  try {
    const { pct } = await getStorageUsage();
    if (pct <= PRUNE_THRESHOLD) return;
    const pruned = await pruneOldData(PRUNE_KEEP_DAYS);
    if (pruned > 0) console.log(`存储占用 ${(pct * 100).toFixed(0)}%，已自动清理 ${pruned} 天前的旧数据`);
  } catch (e) {
    // 剪枝失败不阻塞采集
  }
}

/**
 * 把当前 pending 队列写一份到 chrome.storage.session 作崩溃恢复快照。
 * 只在快照不存在时写：快照存在说明它保存的批次已包含当前内存数据，
 * 每次唤醒都覆盖会把"尚未落库的内存态"退回成"已落库的旧快照"，造成重算。
 */
async function savePendingSnapshot() {
  const existing = await chrome.storage.session.get([
    STORAGE_PENDING, STORAGE_PENDING_RAW, STORAGE_PENDING_DL, STORAGE_PENDING_DL_RAW
  ]).catch(() => null);
  if (!existing) return;
  if (existing[STORAGE_PENDING] || existing[STORAGE_PENDING_RAW] ||
      existing[STORAGE_PENDING_DL] || existing[STORAGE_PENDING_DL_RAW]) return;
  await writePendingSnapshot();
}
restoreReady.then(savePendingSnapshot);
restoreReady.then(ensurePruneAlarm);

chrome.webRequest.onCompleted.addListener(
  async (details) => {
    await restoreReady; // 先确保 pending 已从崩溃恢复，否则会用空对象覆盖已恢复的数据
    let bytes = details.responseSize;
    if (!bytes || bytes <= 0) {
      if (details.responseHeaders) {
        const cl = details.responseHeaders.find(
          h => h.name.toLowerCase() === 'content-length'
        );
        if (cl) bytes = parseInt(cl.value, 10);
      }
      if (!bytes || bytes <= 0) return;
    }
    const info = getDomainInfo(details.url);
    if (!info) return;
    // 分组语义：排除组内任一成员即整组排除（排除 bilibili.com 后 hdslb.com 也不再采集）
    if (isDomainExcluded(excludedDomains, info.root)) return;
    pendingData[info.root] = (pendingData[info.root] || 0) + bytes;
    pendingRawData[info.hostname] = (pendingRawData[info.hostname] || 0) + bytes;
    pendingTotal++;
    requestCounter++;
    if (requestCounter >= 10) {
      requestCounter = 0;
      writePendingSnapshot(); // 崩溃恢复快照失败不影响采集，下个周期会重试
    }
    if (pendingTotal >= FLUSH_COUNT) {
      flushData();
    }
  },
  { urls: ['<all_urls>'] },
  ['responseHeaders', 'extraHeaders']
);

chrome.downloads.onChanged.addListener(async (delta) => {
  if (!delta.state || delta.state.current !== 'complete') return;
  await restoreReady; // 同 webRequest：等恢复完成再累加，避免覆盖崩溃恢复出的 pending
  try {
    const items = await chrome.downloads.search({ id: delta.id });
    if (!items || items.length === 0) return;
    const item = items[0];
    const fileSize = (delta.fileSize && delta.fileSize.current) || item.fileSize;
    if (!fileSize || fileSize <= 0) return;
    if (!item.url) return;
    const info = getDomainInfo(item.url);
    if (!info) return;
    // 分组语义同上：排除组内任一成员即整组排除
    if (isDomainExcluded(excludedDomains, info.root)) return;
    pendingDownloads[info.root] = (pendingDownloads[info.root] || 0) + fileSize;
    pendingDownloadItems.push({
      url: item.url,
      hostname: info.hostname,
      root: info.root,
      fileName: item.filename || '',
      fileSize: fileSize,
      mime: item.mime || '',
      time: new Date().toISOString()
    });
    pendingDownloadCount++;
    if (pendingDownloadCount >= FLUSH_COUNT) {
      flushDownloads();
    }
  } catch (e) {
    // Silently handle errors
  }
});