koc-loop/lib/collection-service.ts

import {
  collectXhsMetricsFromMcp,
  type CollectionMcpConfig,
} from "./mcp-collection-client";
import { uid } from "./mvp-db";
import { isCollectionScheduleDue } from "./collection-schedule";

export { isCollectionScheduleDue } from "./collection-schedule";

type DistributionForCollection = {
  id: string;
  task_id: string;
  publish_url: string | null;
  ocr_status: string;
};

type ScheduledTask = {
  id: string;
  collection_start_date: string;
  collection_days: string;
};

type CollectionSource = "automatic" | "catchup" | "manual";

function utcDay(value: string) {
  const match = value.match(/^(\d{4})-(\d{2})-(\d{2})$/);
  if (!match) return null;
  return Date.UTC(Number(match[1]), Number(match[2]) - 1, Number(match[3]));
}

function dateForScheduleDay(startDate: string, scheduleDay: number) {
  const start = utcDay(startDate);
  if (start === null) return null;
  return new Date(start + (scheduleDay - 1) * 86_400_000)
    .toISOString()
    .slice(0, 10);
}

function parseDays(value: string) {
  try {
    const parsed = JSON.parse(value) as unknown;
    if (!Array.isArray(parsed)) return [];
    return [...new Set(parsed.map(Number))]
      .filter((day) => Number.isInteger(day) && day >= 1 && day <= 7)
      .sort((a, b) => a - b);
  } catch {
    return [];
  }
}

export function shanghaiDateFromTimestamp(timestamp: number) {
  return new Date(timestamp + 8 * 60 * 60 * 1000)
    .toISOString()
    .slice(0, 10);
}

function dueSchedules(
  startDate: string,
  days: number[],
  timestamp: number,
) {
  return days
    .map((scheduleDay) => ({
      scheduleDay,
      scheduledDate: dateForScheduleDay(startDate, scheduleDay),
    }))
    .filter(
      (
        value,
      ): value is { scheduleDay: number; scheduledDate: string } =>
        Boolean(
          value.scheduledDate &&
            isCollectionScheduleDue(value.scheduledDate, timestamp),
        ),
    );
}

export async function createCollectionRunTasks(
  db: D1Database,
  taskId: string,
  startDate: string,
  days: number[],
) {
  const normalizedDays = [...new Set(days)]
    .filter((day) => Number.isInteger(day) && day >= 1 && day <= 7)
    .sort((a, b) => a - b);
  const distributions = await db
    .prepare(
      `SELECT id FROM distributions
       WHERE task_id = ?
         AND publish_url IS NOT NULL
         AND publish_url != ''`,
    )
    .bind(taskId)
    .all<{ id: string }>();

  await db
    .prepare(
      `DELETE FROM collection_runs
       WHERE task_id = ?
         AND status = 'pending'`,
    )
    .bind(taskId)
    .run();

  const statements: D1PreparedStatement[] = [];
  for (const distribution of distributions.results) {
    for (const scheduleDay of normalizedDays) {
      const scheduledDate = dateForScheduleDay(startDate, scheduleDay);
      if (!scheduledDate) continue;
      statements.push(
        db
          .prepare(
            `INSERT OR IGNORE INTO collection_runs
              (id, task_id, distribution_id, scheduled_date, schedule_day,
               scheduled_at, status, status_description)
             VALUES (?, ?, ?, ?, ?, ?, 'pending', ?)`,
          )
          .bind(
            uid("run"),
            taskId,
            distribution.id,
            scheduledDate,
            scheduleDay,
            `${scheduledDate}T09:00:00+08:00`,
            `等待第${scheduleDay}天 09:00自动采集`,
          ),
      );
    }
  }
  for (let index = 0; index < statements.length; index += 100) {
    await db.batch(statements.slice(index, index + 100));
  }
  return {
    distributions: distributions.results.length,
    days: normalizedDays.length,
    tasks: statements.length,
  };
}

export async function collectDistributionMetrics(
  db: D1Database,
  distributionId: string,
  scheduledDate: string,
  scheduleDay: number | null,
  source: CollectionSource,
  mcpConfig: CollectionMcpConfig,
) {
  const current = await db
    .prepare("SELECT * FROM distributions WHERE id = ?")
    .bind(distributionId)
    .first<DistributionForCollection>();
  if (!current) throw new Error("分发记录不存在");
  if (!current.publish_url) throw new Error("笔记尚未回填发布链接");

  const scheduledAt = `${scheduledDate}T09:00:00+08:00`;
  const runId = uid("run");
  await db
    .prepare(
      `INSERT OR IGNORE INTO collection_runs
        (id, task_id, distribution_id, scheduled_date, schedule_day,
         scheduled_at, status, status_description)
       VALUES (?, ?, ?, ?, ?, ?, 'pending', ?)`,
    )
    .bind(
      runId,
      current.task_id,
      distributionId,
      scheduledDate,
      scheduleDay,
      scheduledAt,
      source === "automatic"
        ? "等待自动采集"
        : source === "catchup"
          ? "等待自动追采"
          : "等待手动采集",
    )
    .run();

  const run = await db
    .prepare(
      `SELECT id, status FROM collection_runs
       WHERE distribution_id = ? AND scheduled_date = ?`,
    )
    .bind(distributionId, scheduledDate)
    .first<{ id: string; status: string }>();
  if (!run) throw new Error("采集任务创建失败");
  if (run.status === "success") return { skipped: true };

  const collectingDescription =
    source === "automatic"
      ? `第${scheduleDay ?? "—"}天自动采集中`
      : source === "catchup"
        ? `第${scheduleDay ?? "—"}天自动追采中`
        : "正在手动采集";
  await db.batch([
    db
      .prepare(
        `UPDATE collection_runs
         SET schedule_day = COALESCE(schedule_day, ?),
             scheduled_at = ?,
             status = 'collecting',
             status_description = ?,
             started_at = CURRENT_TIMESTAMP
         WHERE id = ?`,
      )
      .bind(scheduleDay, scheduledAt, collectingDescription, run.id),
    db
      .prepare(
        `UPDATE distributions
         SET collection_status = 'collecting',
             collection_status_description = ?,
             updated_at = CURRENT_TIMESTAMP
         WHERE id = ?`,
      )
      .bind(collectingDescription, distributionId),
  ]);

  try {
    const { likes, comments, collects } =
      await collectXhsMetricsFromMcp(current.publish_url, mcpConfig);
    const dayWeight = scheduleDay ?? 1;
    const successDescription =
      source === "automatic"
        ? `成功 · 第${dayWeight}天 09:00自动采集`
        : source === "catchup"
          ? `成功 · 第${dayWeight}天自动追采`
          : "成功 · 手动采集";

    await db.batch([
      db
        .prepare(
          `UPDATE collection_runs
           SET status = 'success',
               likes = ?,
               comments = ?,
               collects = ?,
               status_description = ?,
               completed_at = CURRENT_TIMESTAMP
           WHERE id = ?`,
        )
        .bind(
          likes,
          comments,
          collects,
          successDescription,
          run.id,
        ),
      db
        .prepare(
          `UPDATE distributions
           SET latest_likes = ?,
               latest_comments = ?,
               latest_collects = ?,
               collection_status = 'success',
               collection_status_description = ?,
               collection_updated_at = CURRENT_TIMESTAMP,
               last_collection_day = ?,
               status = CASE
                 WHEN exposure IS NOT NULL AND views IS NOT NULL THEN 'complete'
                 ELSE 'collecting'
               END,
               updated_at = CURRENT_TIMESTAMP
           WHERE id = ?`,
        )
        .bind(
          likes,
          comments,
          collects,
          successDescription,
          scheduleDay,
          distributionId,
        ),
    ]);
    return { skipped: false, likes, comments, collects };
  } catch (error) {
    const message =
      error instanceof Error ? error.message : "公开数据采集失败";
    await db.batch([
      db
        .prepare(
          `UPDATE collection_runs
           SET status = 'failed',
               status_description = ?,
               completed_at = CURRENT_TIMESTAMP
           WHERE id = ?`,
        )
        .bind(message, run.id),
      db
        .prepare(
          `UPDATE distributions
           SET collection_status = 'failed',
               collection_status_description = ?,
               collection_updated_at = CURRENT_TIMESTAMP,
               updated_at = CURRENT_TIMESTAMP
           WHERE id = ?`,
        )
        .bind(message, distributionId),
    ]);
    throw error;
  }
}

export async function runScheduledCollections(
  db: D1Database,
  scheduledTimestamp: number,
  mcpConfig: CollectionMcpConfig,
) {
  return runDueScheduledCollections(
    db,
    scheduledTimestamp,
    mcpConfig,
    "automatic",
  );
}

export async function runDueScheduledCollections(
  db: D1Database,
  timestamp: number,
  mcpConfig: CollectionMcpConfig,
  source: Extract<CollectionSource, "automatic" | "catchup"> = "catchup",
  onlyTaskId?: string,
) {
  const currentDate = shanghaiDateFromTimestamp(timestamp);
  const tasks = await db
    .prepare(
      `SELECT id, collection_start_date, collection_days
       FROM tasks
       WHERE collection_start_date IS NOT NULL
         AND collection_start_date != ''
         AND collection_days IS NOT NULL
         AND collection_days != '[]'
         AND (? IS NULL OR id = ?)`,
    )
    .bind(onlyTaskId ?? null, onlyTaskId ?? null)
    .all<ScheduledTask>();
  let dueTasks = 0;
  let attempted = 0;
  let succeeded = 0;
  let failed = 0;

  for (const task of tasks.results) {
    const schedules = dueSchedules(
      task.collection_start_date,
      parseDays(task.collection_days),
      timestamp,
    );
    for (const { scheduleDay, scheduledDate } of schedules) {
      dueTasks += 1;
      const distributions = await db
        .prepare(
          `SELECT d.id FROM distributions d
           WHERE d.task_id = ?
             AND d.publish_url IS NOT NULL
             AND d.publish_url != ''
             AND NOT EXISTS (
               SELECT 1 FROM collection_runs r
               WHERE r.distribution_id = d.id
                 AND r.scheduled_date = ?
                 AND (
                   r.status = 'success'
                   OR (
                     r.status = 'collecting'
                     AND r.started_at >= datetime('now', '-30 minutes')
                   )
                   OR (
                     r.status = 'failed'
                     AND r.completed_at >= datetime('now', '-15 minutes')
                   )
                 )
             )`,
        )
        .bind(task.id, scheduledDate)
        .all<{ id: string }>();

      for (const distribution of distributions.results) {
        attempted += 1;
        try {
          await collectDistributionMetrics(
            db,
            distribution.id,
            scheduledDate,
            scheduleDay,
            source,
            mcpConfig,
          );
          succeeded += 1;
        } catch {
          failed += 1;
        }
      }
    }
  }

  return { currentDate, dueTasks, attempted, succeeded, failed };
}

export async function retryFailedCollections(
  db: D1Database,
  taskId: string,
  mcpConfig: CollectionMcpConfig,
) {
  const failedDistributions = await db
    .prepare(
      `SELECT
         d.id,
         COALESCE(
           (
             SELECT r.scheduled_date
             FROM collection_runs r
             WHERE r.distribution_id = d.id
               AND r.status = 'failed'
             ORDER BY COALESCE(r.completed_at, r.created_at) DESC
             LIMIT 1
           ),
           ?
         ) AS scheduled_date,
         (
           SELECT r.schedule_day
           FROM collection_runs r
           WHERE r.distribution_id = d.id
             AND r.status = 'failed'
           ORDER BY COALESCE(r.completed_at, r.created_at) DESC
           LIMIT 1
         ) AS schedule_day
       FROM distributions d
       WHERE d.task_id = ?
         AND d.publish_url IS NOT NULL
         AND d.publish_url != ''
         AND d.collection_status = 'failed'
       ORDER BY COALESCE(d.collection_updated_at, d.updated_at)
       LIMIT 50`,
    )
    .bind(shanghaiDateFromTimestamp(Date.now()), taskId)
    .all<{
      id: string;
      scheduled_date: string;
      schedule_day: number | null;
    }>();

  let succeeded = 0;
  let failed = 0;
  for (const distribution of failedDistributions.results) {
    try {
      await collectDistributionMetrics(
        db,
        distribution.id,
        distribution.scheduled_date,
        distribution.schedule_day,
        "manual",
        mcpConfig,
      );
      succeeded += 1;
    } catch {
      failed += 1;
    }
  }
  return {
    attempted: failedDistributions.results.length,
    succeeded,
    failed,
  };
}