karate_halloween/lib/storage.ts

import fs from "fs";
import path from "path";
import { validateAndSanitizeImage } from "@/lib/security";

export interface GenerationRecord {
  id: string;
  timestamp: string; // ISO 8601 string
  dateStr: string;   // YYYY-MM-DD for fast grouping
  styleId: string;
  styleName: string;
  provider: "gemini" | "openrouter" | "demo_mock" | "unknown";
  model: string;
  durationMs: number;
  status: "success" | "refused" | "error";
  imageUrl?: string;
  imageFilename?: string;
  errorType?: "safety_refusal" | "provider_error" | "rate_limit" | "server_error";
  errorMessage?: string;
  refusalReason?: string;
  clientId?: string;
}

export interface StatsSummary {
  period: string;
  fromDate?: string;
  toDate?: string;
  totalAttempts: number;
  successCount: number;
  refusedCount: number;
  errorCount: number;
  successRate: number;
  avgDurationMs: number;
  dailyTrends: {
    date: string;
    total: number;
    success: number;
    refused: number;
    error: number;
  }[];
  styleBreakdown: {
    styleId: string;
    styleName: string;
    count: number;
    percentage: number;
  }[];
  providerBreakdown: {
    provider: string;
    model: string;
    count: number;
    percentage: number;
  }[];
  images: GenerationRecord[];
  errorLogs: GenerationRecord[];
}

const DATA_DIR = path.join(process.cwd(), "data");
const IMAGES_DIR = path.join(DATA_DIR, "images");
const RECORDS_FILE = path.join(DATA_DIR, "records.json");
const BACKUP_FILE = path.join(DATA_DIR, "records.json.bak");

function ensureDirectories() {
  if (!fs.existsSync(DATA_DIR)) {
    fs.mkdirSync(DATA_DIR, { recursive: true });
  }
  if (!fs.existsSync(IMAGES_DIR)) {
    fs.mkdirSync(IMAGES_DIR, { recursive: true });
  }
  if (!fs.existsSync(RECORDS_FILE)) {
    fs.writeFileSync(RECORDS_FILE, JSON.stringify([], null, 2), "utf8");
  }
}

/**
 * Mutex / Write Queue to prevent concurrent write collisions
 */
let writeQueuePromise: Promise<void> = Promise.resolve();

function queueWrite(task: () => Promise<void> | void): Promise<void> {
  const next = writeQueuePromise.then(() => task()).catch((err) => {
    console.error("[Storage Mutex] Write task error:", err);
  });
  writeQueuePromise = next;
  return next;
}

/**
 * Load all records from disk safely with automatic backup recovery
 */
export function getAllRecords(): GenerationRecord[] {
  ensureDirectories();
  try {
    const raw = fs.readFileSync(RECORDS_FILE, "utf8");
    if (!raw.trim()) {
      return recoverFromBackup();
    }
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) {
      console.warn("[Storage] records.json is not an array, attempting backup recovery...");
      return recoverFromBackup();
    }
    return parsed;
  } catch (err) {
    console.error("[Storage] JSON parse error in records.json, attempting backup recovery:", err);
    return recoverFromBackup();
  }
}

/**
 * Recover from backup file if primary file was corrupted
 */
function recoverFromBackup(): GenerationRecord[] {
  try {
    if (fs.existsSync(BACKUP_FILE)) {
      const backupRaw = fs.readFileSync(BACKUP_FILE, "utf8");
      const parsed = JSON.parse(backupRaw);
      if (Array.isArray(parsed)) {
        console.warn("[Storage] Successfully recovered records from records.json.bak!");
        // Re-write valid records back to records.json
        fs.writeFileSync(RECORDS_FILE, backupRaw, "utf8");
        return parsed;
      }
    }
  } catch (backupErr) {
    console.error("[Storage] Backup recovery failed:", backupErr);
  }
  return [];
}

/**
 * Persist records to disk atomically with fsync and backup rotation
 */
async function saveAllRecordsAtomic(records: GenerationRecord[]): Promise<void> {
  await queueWrite(() => {
    ensureDirectories();

    const serialized = JSON.stringify(records, null, 2);
    // Safety check: Never write invalid or empty string if we had records
    if (!serialized || serialized.length < 2) {
      console.error("[Storage] Refusing to write invalid empty JSON payload.");
      return;
    }

    // 1. Maintain backup copy of current good records file before replacing
    if (fs.existsSync(RECORDS_FILE)) {
      try {
        const currentContent = fs.readFileSync(RECORDS_FILE, "utf8");
        if (currentContent.trim().length > 2) {
          fs.writeFileSync(BACKUP_FILE, currentContent, "utf8");
        }
      } catch (backupErr) {
        console.warn("[Storage] Failed to update backup copy:", backupErr);
      }
    }

    // 2. Write to unique temporary file
    const tempFile = path.join(
      DATA_DIR,
      `records.${Date.now()}_${Math.random().toString(36).substring(2, 8)}.tmp`
    );

    const fd = fs.openSync(tempFile, "w");
    fs.writeSync(fd, serialized, 0, "utf8");
    // Ensure data is physically flushed to disk platter/SSD before rename
    fs.fsyncSync(fd);
    fs.closeSync(fd);

    // 3. Atomic rename replaces the destination file instantaneously
    fs.renameSync(tempFile, RECORDS_FILE);
  });
}

/**
 * Save an image file from base64 or URL buffer safely with binary validation
 */
export async function persistImageFile(
  id: string,
  imageInput: string
): Promise<{ imageUrl: string; filename: string }> {
  ensureDirectories();

  // Sanitize id to alphanumeric characters, hyphens, and underscores only
  const safeId = id.replace(/[^a-zA-Z0-9_-]/g, "");
  if (!safeId) {
    return { imageUrl: imageInput, filename: "" };
  }

  // If it's a data URL or raw base64
  if (imageInput.startsWith("data:image/") || imageInput.startsWith("data:application/")) {
    const validation = validateAndSanitizeImage(imageInput, 15 * 1024 * 1024);
    if (validation.valid && validation.extension) {
      const ext = validation.extension; // Strictly "jpg" | "png" | "webp"
      const filename = `${safeId}.${ext}`;
      const filePath = path.join(IMAGES_DIR, filename);

      const base64Data = validation.sanitizedBase64!.split(",")[1];
      fs.writeFileSync(filePath, Buffer.from(base64Data, "base64"));

      // Automatically apply school branding: logo (top-left), WeChat QR (bottom-right), and footer text
      try {
        const { applyWatermarkToFile } = await import("@/lib/watermark");
        await applyWatermarkToFile(filePath);
      } catch (wmErr) {
        console.warn("[Storage] Watermark application skipped:", wmErr);
      }

      return { imageUrl: `/api/images/${filename}`, filename };
    } else {
      console.warn(`[Storage Security Alert] Refused to persist invalid image: ${validation.error}`);
      return { imageUrl: imageInput, filename: "" };
    }
  }

  // If it's a remote URL
  if (imageInput.startsWith("http://") || imageInput.startsWith("https://")) {
    try {
      const parsedUrl = new URL(imageInput);
      const hostname = parsedUrl.hostname.toLowerCase();

      // Block SSRF loopback / private addresses
      const isPrivateOrLoopback =
        hostname === "localhost" ||
        hostname === "127.0.0.1" ||
        hostname === "::1" ||
        hostname.startsWith("10.") ||
        hostname.startsWith("192.168.") ||
        /^172\.(1[6-9]|2[0-9]|3[0-1])\./.test(hostname) ||
        hostname.startsWith("169.254.");

      if (isPrivateOrLoopback) {
        console.warn(`[Storage SSRF Blocked] Attempted fetch to private/local address: ${hostname}`);
        return { imageUrl: imageInput, filename: "" };
      }

      const res = await fetch(imageInput);
      if (res.ok) {
        const buffer = Buffer.from(await res.arrayBuffer());
        // Validate downloaded buffer via data URL representation
        const validation = validateAndSanitizeImage(
          `data:image/png;base64,${buffer.toString("base64")}`,
          15 * 1024 * 1024
        );

        if (validation.valid && validation.extension) {
          const filename = `${safeId}.${validation.extension}`;
          const filePath = path.join(IMAGES_DIR, filename);
          fs.writeFileSync(filePath, buffer);

          // Automatically apply school branding: logo (top-left), WeChat QR (bottom-right), and footer text
          try {
            const { applyWatermarkToFile } = await import("@/lib/watermark");
            await applyWatermarkToFile(filePath);
          } catch (wmErr) {
            console.warn("[Storage] Watermark application skipped:", wmErr);
          }

          return { imageUrl: `/api/images/${filename}`, filename };
        } else {
          console.warn(`[Storage Security Alert] Remote file was not a valid image: ${validation.error}`);
        }
      }
    } catch (e) {
      console.warn("[Storage] Could not fetch remote image to cache locally, storing direct URL:", e);
    }
    return { imageUrl: imageInput, filename: "" };
  }

  // Local sample asset
  return { imageUrl: imageInput, filename: "" };
}

/**
 * Record a generation event (both successes, safety refusals, and server errors)
 */
export async function logGenerationEvent(data: {
  styleId: string;
  styleName: string;
  provider: "gemini" | "openrouter" | "demo_mock" | "unknown";
  model: string;
  durationMs: number;
  status: "success" | "refused" | "error";
  rawImageUrl?: string;
  errorType?: "safety_refusal" | "provider_error" | "rate_limit" | "server_error";
  errorMessage?: string;
  refusalReason?: string;
  clientId?: string;
}): Promise<GenerationRecord> {
  const now = new Date();
  const id = `gen_${Date.now()}_${Math.random().toString(36).substring(2, 7)}`;
  const dateStr = now.toISOString().slice(0, 10); // YYYY-MM-DD

  let imageUrl = data.rawImageUrl;
  let imageFilename = "";

  if (data.status === "success" && data.rawImageUrl) {
    const persisted = await persistImageFile(id, data.rawImageUrl);
    imageUrl = persisted.imageUrl;
    imageFilename = persisted.filename;
  }

  const record: GenerationRecord = {
    id,
    timestamp: now.toISOString(),
    dateStr,
    styleId: data.styleId,
    styleName: data.styleName,
    provider: data.provider,
    model: data.model,
    durationMs: data.durationMs,
    status: data.status,
    imageUrl,
    imageFilename,
    errorType: data.errorType,
    errorMessage: data.errorMessage,
    refusalReason: data.refusalReason,
    clientId: data.clientId,
  };

  try {
    const existing = getAllRecords();
    existing.unshift(record); // Prepend so newest is first
    await saveAllRecordsAtomic(existing);
  } catch (err) {
    console.error("[Storage] Failed to append record:", err);
  }

  return record;
}

/**
 * Query statistics with date filtering
 */
export function getGenerationStats(query: {
  period?: "day" | "week" | "month" | "all" | "custom";
  fromDate?: string; // YYYY-MM-DD
  toDate?: string;   // YYYY-MM-DD
}): StatsSummary {
  const records = getAllRecords();
  const now = new Date();
  const todayStr = now.toISOString().slice(0, 10);

  const period = query.period || "week";
  let filtered = records;

  let computedFrom = query.fromDate;
  let computedTo = query.toDate;

  if (period === "day") {
    computedFrom = todayStr;
    computedTo = todayStr;
    filtered = records.filter((r) => r.dateStr === todayStr);
  } else if (period === "week") {
    const sevenDaysAgo = new Date(now.getTime() - 6 * 24 * 60 * 60 * 1000);
    computedFrom = sevenDaysAgo.toISOString().slice(0, 10);
    computedTo = todayStr;
    filtered = records.filter((r) => r.dateStr >= computedFrom! && r.dateStr <= computedTo!);
  } else if (period === "month") {
    const thirtyDaysAgo = new Date(now.getTime() - 29 * 24 * 60 * 60 * 1000);
    computedFrom = thirtyDaysAgo.toISOString().slice(0, 10);
    computedTo = todayStr;
    filtered = records.filter((r) => r.dateStr >= computedFrom! && r.dateStr <= computedTo!);
  } else if (period === "custom") {
    if (query.fromDate) {
      filtered = filtered.filter((r) => r.dateStr >= query.fromDate!);
    }
    if (query.toDate) {
      filtered = filtered.filter((r) => r.dateStr <= query.toDate!);
    }
  }

  const totalAttempts = filtered.length;
  const successCount = filtered.filter((r) => r.status === "success").length;
  const refusedCount = filtered.filter((r) => r.status === "refused").length;
  const errorCount = filtered.filter((r) => r.status === "error").length;

  const successRate = totalAttempts > 0 ? Math.round((successCount / totalAttempts) * 1000) / 10 : 0;

  const successfulRecords = filtered.filter((r) => r.status === "success");
  const totalDuration = successfulRecords.reduce((sum, r) => sum + (r.durationMs || 0), 0);
  const avgDurationMs = successfulRecords.length > 0 ? Math.round(totalDuration / successfulRecords.length) : 0;

  // Daily Trends aggregation
  const dailyMap = new Map<string, { total: number; success: number; refused: number; error: number }>();

  // If week or day or month, pre-populate all dates in range so charts are continuous
  if (computedFrom && computedTo) {
    const start = new Date(computedFrom);
    const end = new Date(computedTo);
    const cur = new Date(start);
    while (cur <= end) {
      const d = cur.toISOString().slice(0, 10);
      dailyMap.set(d, { total: 0, success: 0, refused: 0, error: 0 });
      cur.setDate(cur.getDate() + 1);
    }
  }

  for (const r of filtered) {
    const entry = dailyMap.get(r.dateStr) || { total: 0, success: 0, refused: 0, error: 0 };
    entry.total += 1;
    if (r.status === "success") entry.success += 1;
    else if (r.status === "refused") entry.refused += 1;
    else if (r.status === "error") entry.error += 1;
    dailyMap.set(r.dateStr, entry);
  }

  const dailyTrends = Array.from(dailyMap.entries())
    .map(([date, counts]) => ({ date, ...counts }))
    .sort((a, b) => a.date.localeCompare(b.date));

  // Style Breakdown
  const styleCountMap = new Map<string, { name: string; count: number }>();
  for (const r of filtered) {
    if (r.styleId) {
      const existing = styleCountMap.get(r.styleId) || { name: r.styleName || r.styleId, count: 0 };
      existing.count += 1;
      styleCountMap.set(r.styleId, existing);
    }
  }
  const styleBreakdown = Array.from(styleCountMap.entries())
    .map(([styleId, { name, count }]) => ({
      styleId,
      styleName: name,
      count,
      percentage: totalAttempts > 0 ? Math.round((count / totalAttempts) * 100) : 0,
    }))
    .sort((a, b) => b.count - a.count);

  // Provider Breakdown
  const providerCountMap = new Map<string, { provider: string; model: string; count: number }>();
  for (const r of filtered) {
    const key = `${r.provider}::${r.model || "default"}`;
    const existing = providerCountMap.get(key) || { provider: r.provider, model: r.model || "default", count: 0 };
    existing.count += 1;
    providerCountMap.set(key, existing);
  }
  const providerBreakdown = Array.from(providerCountMap.values())
    .map((item) => ({
      ...item,
      percentage: totalAttempts > 0 ? Math.round((item.count / totalAttempts) * 100) : 0,
    }))
    .sort((a, b) => b.count - a.count);

  // Filtered Images Gallery (success only)
  const images = filtered.filter((r) => r.status === "success" && r.imageUrl);

  // Filtered Error & Refusal Logs
  const errorLogs = filtered.filter((r) => r.status !== "success");

  return {
    period,
    fromDate: computedFrom,
    toDate: computedTo,
    totalAttempts,
    successCount,
    refusedCount,
    errorCount,
    successRate,
    avgDurationMs,
    dailyTrends,
    styleBreakdown,
    providerBreakdown,
    images,
    errorLogs,
  };
}