import { readFile } from "node:fs/promises";
import mysql from "mysql2/promise";

const allowedTables = [
  "partners",
  "tasks",
  "contents",
  "accounts",
  "claims",
  "delegation_bundles",
  "distributions",
  "collection_runs",
  "users",
  "auth_sessions",
  "mcp_export_tokens",
];

function databaseUrl() {
  if (process.env.DATABASE_URL) return process.env.DATABASE_URL;
  const host = process.env.MYSQL_HOST ?? "127.0.0.1";
  const port = process.env.MYSQL_PORT ?? "3306";
  const user = encodeURIComponent(process.env.MYSQL_USER ?? "koc");
  const password = encodeURIComponent(process.env.MYSQL_PASSWORD ?? "");
  const database = process.env.MYSQL_DATABASE ?? "koc_loop";
  return `mysql://${user}:${password}@${host}:${port}/${database}`;
}

function normalizeImportedValue(value) {
  if (
    typeof value === "string" &&
    /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,3})?Z$/.test(value)
  ) {
    return value.replace("T", " ").replace(/Z$/, "");
  }
  return value;
}

const inputPath = process.argv[2];
if (!inputPath) {
  throw new Error("用法：npm run db:import-json -- /path/to/d1-export.json");
}
const parsed = JSON.parse(await readFile(inputPath, "utf8"));
const tables = parsed.tables && typeof parsed.tables === "object"
  ? parsed.tables
  : parsed;
const pool = mysql.createPool({
  uri: databaseUrl(),
  connectionLimit: 2,
  charset: "utf8mb4",
  timezone: "Z",
  dateStrings: true,
});

try {
  for (const table of allowedTables) {
    const rows = Array.isArray(tables[table]) ? tables[table] : [];
    if (rows.length === 0) continue;
    const [columnRows] = await pool.query(
      `SELECT column_name, extra FROM information_schema.columns
       WHERE table_schema = DATABASE() AND table_name = ? ORDER BY ordinal_position`,
      [table],
    );
    const allowedColumns = new Set(
      columnRows
        .filter((column) => !String(column.extra ?? "").includes("GENERATED"))
        .map((column) => String(column.column_name)),
    );
    const columns = [...new Set(rows.flatMap((row) => Object.keys(row)))]
      .filter((column) => allowedColumns.has(column));
    if (columns.length === 0) continue;
    const quotedColumns = columns.map((column) => `\`${column}\``).join(", ");
    const updates = columns
      .filter((column) => column !== "id" && column !== "token_hash")
      .map((column) => `\`${column}\` = VALUES(\`${column}\`)`)
      .join(", ");
    const connection = await pool.getConnection();
    try {
      await connection.beginTransaction();
      for (let offset = 0; offset < rows.length; offset += 200) {
        const chunk = rows.slice(offset, offset + 200);
        const placeholders = chunk
          .map(() => `(${columns.map(() => "?").join(", ")})`)
          .join(", ");
        const sql = `INSERT INTO \`${table}\` (${quotedColumns}) VALUES ${placeholders}${
          updates ? ` ON DUPLICATE KEY UPDATE ${updates}` : ""
        }`;
        const values = chunk.flatMap((row) =>
          columns.map((column) => normalizeImportedValue(row[column] ?? null)),
        );
        await connection.query(sql, values);
      }
      await connection.commit();
      console.info(`imported ${table}: ${rows.length}`);
    } catch (error) {
      await connection.rollback();
      throw error;
    } finally {
      connection.release();
    }
  }
} finally {
  await pool.end();
}