massive-cloud/server/db.ts

import { eq, desc, sql } from "drizzle-orm";
import { drizzle } from "drizzle-orm/mysql2";
import type { ResultSetHeader } from "mysql2";
import {
  users, InsertUser,
  vehicleSeries, InsertVehicleSeries, VehicleSeries,
  vehicleModels, InsertVehicleModel, VehicleModel,
  modelCadFiles, InsertModelCadFile, ModelCadFile,
  taskCadFiles, InsertTaskCadFile, TaskCadFile,
  inspectionTasks, InsertInspectionTask, InspectionTask,
  taskPointcloudFiles, InsertTaskPointcloudFile, TaskPointcloudFile,
  taskAnalysisResults, InsertTaskAnalysisResult, TaskAnalysisResult,
} from "../drizzle/schema";
import { ENV } from './_core/env';

let _db: ReturnType<typeof drizzle> | null = null;

const sleep = (ms: number) => new Promise(r => setTimeout(r, ms));

/** drizzle mysql2 insert resolves to mysql2 ResultSetHeader with insertId (not result[0].insertId). */
function getMysqlInsertId(result: unknown): number {
  const r = result as ResultSetHeader & { 0?: ResultSetHeader };
  if (typeof r.insertId === "number" && r.insertId > 0) return r.insertId;
  if (r[0] && typeof r[0].insertId === "number") return r[0].insertId;
  throw new Error("insertId missing from insert result");
}

/**
 * Clears the cached pool so the next getDb() creates a new connection.
 * Call after ECONNREFUSED / PROTOCOL_CONNECTION_LOST so retries can recover.
 */
export function invalidateDb() {
  _db = null;
}

/**
 * Returns the drizzle DB instance, creating it if needed.
 * Retries the initial TCP probe while MySQL is still starting (common in Docker).
 */
export async function getDb() {
  if (_db) return _db;
  if (!process.env.DATABASE_URL) return null;

  const maxAttempts = 40;
  const delayMs = 500;
  // One drizzle instance / pool; retry execute only (avoids leaking a pool per attempt).
  const instance = drizzle(process.env.DATABASE_URL);

  for (let attempt = 1; attempt <= maxAttempts; attempt++) {
    try {
      await instance.execute(sql`SELECT 1`);
      _db = instance;
      if (attempt > 1) {
        console.log(`[Database] Connected after ${attempt} attempt(s).`);
      }
      return _db;
    } catch (error: any) {
      const msg = error?.message ?? String(error);
      const retriable =
        error?.code === "ECONNREFUSED" ||
        error?.code === "ETIMEDOUT" ||
        error?.code === "ENOTFOUND" ||
        msg.includes("connect");

      if (attempt < maxAttempts && retriable) {
        if (attempt === 1 || attempt % 5 === 0) {
          console.warn(`[Database] Waiting for MySQL... (${attempt}/${maxAttempts}) ${msg}`);
        }
        await sleep(delayMs);
        continue;
      }
      console.error("[Database] Giving up:", msg);
      _db = null;
      return null;
    }
  }
  return null;
}

// ─── 用户 ─────────────────────────────────────────────────────────────────────

export async function upsertUser(user: InsertUser): Promise<void> {
  if (!user.openId) throw new Error("User openId is required for upsert");

  const db = await getDb();
  if (!db) {
    console.warn("[Database] Cannot upsert user: database not available");
    return;
  }

  try {
    const values: InsertUser = { openId: user.openId };
    const updateSet: Record<string, unknown> = {};

    const textFields = ["name", "email", "loginMethod"] as const;
    type TextField = (typeof textFields)[number];

    const assignNullable = (field: TextField) => {
      const value = user[field];
      if (value === undefined) return;
      const normalized = value ?? null;
      values[field] = normalized;
      updateSet[field] = normalized;
    };
    textFields.forEach(assignNullable);

    if (user.lastSignedIn !== undefined) {
      values.lastSignedIn = user.lastSignedIn;
      updateSet.lastSignedIn = user.lastSignedIn;
    }
    if (user.role !== undefined) {
      values.role = user.role;
      updateSet.role = user.role;
    } else if (user.openId === ENV.ownerOpenId) {
      values.role = 'admin';
      updateSet.role = 'admin';
    }

    if (!values.lastSignedIn) values.lastSignedIn = new Date();
    if (Object.keys(updateSet).length === 0) updateSet.lastSignedIn = new Date();

    await db.insert(users).values(values).onDuplicateKeyUpdate({ set: updateSet });
  } catch (error) {
    console.error("[Database] Failed to upsert user:", error);
    throw error;
  }
}

export async function getUserByOpenId(openId: string) {
  const db = await getDb();
  if (!db) return undefined;
  const result = await db.select().from(users).where(eq(users.openId, openId)).limit(1);
  return result[0];
}

// ─── 车系 ─────────────────────────────────────────────────────────────────────

export async function createVehicleSeries(data: InsertVehicleSeries): Promise<VehicleSeries> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(vehicleSeries).values(data);
  const [row] = await db.select().from(vehicleSeries).where(eq(vehicleSeries.id, getMysqlInsertId(result)));
  return row;
}

export async function listVehicleSeries(): Promise<VehicleSeries[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(vehicleSeries).orderBy(desc(vehicleSeries.createdAt));
}

export async function updateVehicleSeries(id: number, data: Partial<InsertVehicleSeries>): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.update(vehicleSeries).set(data).where(eq(vehicleSeries.id, id));
}

export async function deleteVehicleSeries(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(vehicleSeries).where(eq(vehicleSeries.id, id));
}

// ─── 车型 ─────────────────────────────────────────────────────────────────────

export async function createVehicleModel(data: InsertVehicleModel): Promise<VehicleModel> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(vehicleModels).values(data);
  const [row] = await db.select().from(vehicleModels).where(eq(vehicleModels.id, getMysqlInsertId(result)));
  return row;
}

export async function listVehicleModelsBySeries(seriesId: number): Promise<VehicleModel[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(vehicleModels)
    .where(eq(vehicleModels.seriesId, seriesId))
    .orderBy(desc(vehicleModels.createdAt));
}

export async function getVehicleModelById(id: number): Promise<VehicleModel | undefined> {
  const db = await getDb();
  if (!db) return undefined;
  const [row] = await db.select().from(vehicleModels).where(eq(vehicleModels.id, id));
  return row;
}

export async function updateVehicleModel(id: number, data: Partial<InsertVehicleModel>): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.update(vehicleModels).set(data).where(eq(vehicleModels.id, id));
}

export async function deleteVehicleModel(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(vehicleModels).where(eq(vehicleModels.id, id));
}

// ─── 车型 CAD 文件 ─────────────────────────────────────────────────────────────

export async function createModelCadFile(data: InsertModelCadFile): Promise<ModelCadFile> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(modelCadFiles).values(data);
  const [row] = await db.select().from(modelCadFiles).where(eq(modelCadFiles.id, getMysqlInsertId(result)));
  return row;
}

export async function listModelCadFiles(modelId: number): Promise<ModelCadFile[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(modelCadFiles)
    .where(eq(modelCadFiles.modelId, modelId))
    .orderBy(desc(modelCadFiles.createdAt));
}

export async function getModelCadFileById(id: number): Promise<ModelCadFile | undefined> {
  const db = await getDb();
  if (!db) return undefined;
  const [row] = await db.select().from(modelCadFiles).where(eq(modelCadFiles.id, id));
  return row;
}

export async function deleteModelCadFile(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(modelCadFiles).where(eq(modelCadFiles.id, id));
}

// ─── 任务 CAD 文件 ─────────────────────────────────────────────────────────────

export async function createTaskCadFile(data: InsertTaskCadFile): Promise<TaskCadFile> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(taskCadFiles).values(data);
  const [row] = await db.select().from(taskCadFiles).where(eq(taskCadFiles.id, getMysqlInsertId(result)));
  return row;
}

export async function listTaskCadFiles(taskId: number): Promise<TaskCadFile[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(taskCadFiles)
    .where(eq(taskCadFiles.taskId, taskId))
    .orderBy(taskCadFiles.createdAt);
}

export async function getTaskCadFileById(id: number): Promise<TaskCadFile | undefined> {
  const db = await getDb();
  if (!db) return undefined;
  const [row] = await db.select().from(taskCadFiles).where(eq(taskCadFiles.id, id));
  return row;
}

export async function deleteTaskCadFile(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(taskCadFiles).where(eq(taskCadFiles.id, id));
}

// ─── 检测任务 ─────────────────────────────────────────────────────────────────

export async function createInspectionTask(data: InsertInspectionTask): Promise<InspectionTask> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(inspectionTasks).values(data);
  const [row] = await db.select().from(inspectionTasks).where(eq(inspectionTasks.id, getMysqlInsertId(result)));
  return row;
}

export async function listInspectionTasksByModel(modelId: number): Promise<InspectionTask[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(inspectionTasks)
    .where(eq(inspectionTasks.modelId, modelId))
    .orderBy(desc(inspectionTasks.createdAt));
}

export async function getInspectionTaskById(id: number): Promise<InspectionTask | undefined> {
  const db = await getDb();
  if (!db) return undefined;
  const [row] = await db.select().from(inspectionTasks).where(eq(inspectionTasks.id, id));
  return row;
}

export async function updateInspectionTask(id: number, data: Partial<InsertInspectionTask>): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.update(inspectionTasks).set(data).where(eq(inspectionTasks.id, id));
}

export async function deleteInspectionTask(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(inspectionTasks).where(eq(inspectionTasks.id, id));
}

// ─── 任务点云文件 ──────────────────────────────────────────────────────────────

export async function createTaskPointcloudFile(data: InsertTaskPointcloudFile): Promise<TaskPointcloudFile> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(taskPointcloudFiles).values(data);
  const [row] = await db.select().from(taskPointcloudFiles).where(eq(taskPointcloudFiles.id, getMysqlInsertId(result)));
  return row;
}

/** 把配准算法产出的点云注册为任务点云（source='registered'）。
 *  filePath 必须是绝对路径，文件必须已存在于磁盘。
 */
export async function createRegisteredPointcloudFile(
  taskId: number,
  name: string,
  filePath: string,
  sourceResultId: number,
): Promise<TaskPointcloudFile> {
  const fs = await import('node:fs');
  const fileSize = fs.existsSync(filePath) ? fs.statSync(filePath).size : 0;
  return createTaskPointcloudFile({
    taskId,
    name,
    filePath,
    fileSize,
    source: 'registered',
    sourceResultId,
  } as InsertTaskPointcloudFile);
}

export async function listTaskPointcloudFiles(taskId: number): Promise<TaskPointcloudFile[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(taskPointcloudFiles)
    .where(eq(taskPointcloudFiles.taskId, taskId))
    .orderBy(taskPointcloudFiles.createdAt);
}

export async function deleteTaskPointcloudFile(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(taskPointcloudFiles).where(eq(taskPointcloudFiles.id, id));
}

// ─── 任务分析结果 ──────────────────────────────────────────────────────────────

export async function createTaskAnalysisResult(data: InsertTaskAnalysisResult): Promise<TaskAnalysisResult> {
  const db = await getDb();
  if (!db) throw new Error("Database not available");
  const result = await db.insert(taskAnalysisResults).values(data);
  const [row] = await db.select().from(taskAnalysisResults).where(eq(taskAnalysisResults.id, getMysqlInsertId(result)));
  return row;
}

export async function listTaskAnalysisResults(taskId: number): Promise<TaskAnalysisResult[]> {
  const db = await getDb();
  if (!db) return [];
  return db.select().from(taskAnalysisResults)
    .where(eq(taskAnalysisResults.taskId, taskId))
    .orderBy(desc(taskAnalysisResults.createdAt));
}

export async function getTaskAnalysisResultById(id: number): Promise<TaskAnalysisResult | undefined> {
  const db = await getDb();
  if (!db) return undefined;
  const [row] = await db.select().from(taskAnalysisResults).where(eq(taskAnalysisResults.id, id));
  return row;
}

export async function updateTaskAnalysisResult(id: number, data: Partial<InsertTaskAnalysisResult>): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.update(taskAnalysisResults).set(data).where(eq(taskAnalysisResults.id, id));
}

/**
 * 非破坏性局部更新 task_analysis_results.resultData。
 * 与 updateTaskAnalysisResult(resultData) 不同——后者会把整个 JSON 列覆写，
 * 适合「最终结果」；本函数把 patch 合并到现有 JSON 上，适合
 * 子进程频繁上报进度的场景（例如算法插件 progress / stage）。
 *
 * 实现策略：先读取一次再 set 写回；并发更新时按「最后写入者获胜」处理，
 * 对进度上报来说足够安全（覆盖丢失最多 1~2 个心跳）。
 */
export async function patchTaskAnalysisResultData(
  id: number,
  patch: Record<string, unknown>,
): Promise<void> {
  const db = await getDb();
  if (!db) return;
  const existing = await getTaskAnalysisResultById(id);
  if (!existing) return;
  const prev = (existing.resultData as Record<string, unknown> | null) ?? {};
  const merged: Record<string, unknown> = { ...prev };
  for (const [k, v] of Object.entries(patch)) {
    if (v === undefined) continue;
    merged[k] = v;
  }
  await db
    .update(taskAnalysisResults)
    .set({ resultData: merged })
    .where(eq(taskAnalysisResults.id, id));
}

export async function deleteTaskAnalysisResult(id: number): Promise<void> {
  const db = await getDb();
  if (!db) return;
  await db.delete(taskAnalysisResults).where(eq(taskAnalysisResults.id, id));
}