aigen/server/utils/storage.ts

import { existsSync, readdirSync, readFileSync, realpathSync, statSync } from 'node:fs'
import { statfs } from 'node:fs/promises'
import { execFile } from 'node:child_process'
import { dirname, join, resolve } from 'node:path'
import { promisify } from 'node:util'
import { libraryRoot, ownerRoot } from '~/server/utils/library'

const execFileAsync = promisify(execFile)

const OVERLAY_TYPES = new Set(['overlay', 'aufs', 'fuse-overlayfs', 'overlayfs'])

export interface StorageBreakdown {
  clips: number
  stills: number
  tmp: number
  other: number
}

export interface VolumeStats {
  ok: boolean
  totalBytes: number | null
  usedBytes: number | null
  freeBytes: number | null
  overlay: boolean
}

export interface LibraryStorage {
  libraryBytes: number
  breakdown: StorageBreakdown
  volume: VolumeStats
  counted: string
}

function emptyBreakdown(): StorageBreakdown {
  return { clips: 0, stills: 0, tmp: 0, other: 0 }
}

function bucketFor(prefix: string): keyof StorageBreakdown {
  if (prefix === 'files') return 'clips'
  if (prefix === 'stills') return 'stills'
  if (prefix === 'tmp') return 'tmp'
  return 'other'
}

function walkDir(dir: string, breakdown: StorageBreakdown, prefix: string) {
  let entries
  try {
    entries = readdirSync(dir, { withFileTypes: true })
  } catch {
    return
  }
  for (const entry of entries) {
    if (entry.isSymbolicLink()) continue
    const full = join(dir, entry.name)
    if (entry.isDirectory()) {
      walkDir(full, breakdown, prefix || entry.name)
      continue
    }
    if (!entry.isFile()) continue
    let size = 0
    try {
      size = statSync(full).size
    } catch {
      continue
    }
    breakdown[bucketFor(prefix)] += size
  }
}

function ownerLibraryBytes(owner: string) {
  const breakdown = emptyBreakdown()
  const root = ownerRoot(owner)
  if (existsSync(root)) walkDir(root, breakdown, '')
  return {
    libraryBytes: breakdown.clips + breakdown.stills + breakdown.tmp + breakdown.other,
    breakdown
  }
}

function decodeMountField(value: string) {
  return value.replace(/\\([0-7]{3})/g, (_, oct) => String.fromCharCode(parseInt(oct, 8)))
}

function mountFor(target: string): { mountPoint: string; fstype: string; source: string } | null {
  if (!existsSync('/proc/self/mountinfo')) return null
  let text = ''
  try {
    text = readFileSync('/proc/self/mountinfo', 'utf8')
  } catch {
    return null
  }
  const normalized = target.replace(/\\/g, '/').replace(/\/$/, '') || '/'
  let best: { mountPoint: string; fstype: string; source: string } | null = null
  for (const line of text.split('\n')) {
    if (!line.trim()) continue
    const parts = line.split(' - ')
    if (parts.length < 2) continue
    const left = parts[0].split(' ')
    const right = parts[1].split(' ')
    const mountPoint = decodeMountField(left[4] || '')
    const fstype = right[0] || ''
    const source = decodeMountField(right[1] || '')
    if (!mountPoint) continue
    const prefix = mountPoint === '/' ? '/' : `${mountPoint.replace(/\/$/, '')}/`
    const hit = normalized === mountPoint || normalized.startsWith(prefix) || `${normalized}/` === prefix
    if (!hit) continue
    if (!best || mountPoint.length > best.mountPoint.length) {
      best = { mountPoint, fstype, source }
    }
  }
  return best
}

function isOverlayFs(fstype: string, source: string) {
  const type = (fstype || '').toLowerCase()
  const src = (source || '').toLowerCase()
  return OVERLAY_TYPES.has(type) || src === 'overlay'
}

async function volumeFromStatfs(path: string): Promise<Omit<VolumeStats, 'overlay'> | null> {
  try {
    const stats = await statfs(path)
    const bsize = Number(stats.bsize)
    const total = Number(stats.blocks) * bsize
    const free = Number(stats.bavail) * bsize
    if (!Number.isFinite(total) || total <= 0) return null
    return {
      ok: true,
      totalBytes: total,
      usedBytes: Math.max(0, total - (Number.isFinite(free) ? free : 0)),
      freeBytes: Number.isFinite(free) ? Math.max(0, free) : null
    }
  } catch {
    return null
  }
}

async function volumeFromDf(path: string): Promise<Omit<VolumeStats, 'overlay'> | null> {
  try {
    const { stdout } = await execFileAsync('df', ['-kP', path], { timeout: 4000, windowsHide: true })
    const lines = String(stdout || '').trim().split(/\r?\n/).filter(Boolean)
    const row = lines[lines.length - 1]
    if (!row || lines.length < 2) return null
    const cols = row.trim().split(/\s+/)
    if (cols.length < 6) return null
    const totalK = Number(cols[cols.length - 5])
    const usedK = Number(cols[cols.length - 4])
    const availK = Number(cols[cols.length - 3])
    if (!Number.isFinite(totalK) || totalK <= 0) return null
    const total = totalK * 1024
    const free = Number.isFinite(availK) ? availK * 1024 : null
    const used = Number.isFinite(usedK) ? usedK * 1024 : (free == null ? null : Math.max(0, total - free))
    return { ok: true, totalBytes: total, usedBytes: used, freeBytes: free }
  } catch {
    return null
  }
}

function volumeProbePath() {
  const root = libraryRoot()
  const candidates = [root, dirname(resolve(root))]
  for (const candidate of candidates) {
    if (candidate && existsSync(candidate)) {
      try {
        return realpathSync(candidate)
      } catch {
        return resolve(candidate)
      }
    }
  }
  return resolve(root)
}

async function volumeStats(): Promise<VolumeStats> {
  const path = volumeProbePath()
  const mount = mountFor(path.replace(/\\/g, '/'))
  const overlay = mount ? isOverlayFs(mount.fstype, mount.source) : false
  const fromStatfs = await volumeFromStatfs(path)
  if (fromStatfs) return { ...fromStatfs, overlay }
  const fromDf = await volumeFromDf(path)
  if (fromDf) return { ...fromDf, overlay }
  return { ok: false, totalBytes: null, usedBytes: null, freeBytes: null, overlay }
}

export async function libraryStorage(owner: string): Promise<LibraryStorage> {
  const { libraryBytes, breakdown } = ownerLibraryBytes(owner)
  return {
    libraryBytes,
    breakdown,
    volume: await volumeStats(),
    counted: "This account's clips, stills, thumbs, catalogs, queues, and temp files. Shared pending JSON and other accounts are not included."
  }
}