lunarfront-app/packages/backend/src/services/vault.service.ts

import { eq, and, count, inArray, type Column } from 'drizzle-orm'
import type { PostgresJsDatabase } from 'drizzle-orm/postgres-js'
import { vaultConfig, vaultCategories, vaultCategoryPermissions, vaultEntries } from '../db/schema/vault.js'
import { userRoles } from '../db/schema/rbac.js'
import { withPagination, withSort, buildSearchCondition, paginatedResponse } from '../utils/pagination.js'
import type { PaginationInput } from '@lunarfront/shared/schemas'
import { randomBytes, createCipheriv, createDecipheriv, pbkdf2Sync } from 'crypto'
import bcrypt from 'bcryptjs'

// --- Encryption key held in memory ---
let derivedKey: Buffer | null = null

const PBKDF2_ITERATIONS = 100_000
const KEY_LENGTH = 32 // 256 bits for AES-256
const ALGORITHM = 'aes-256-gcm'

const ACCESS_RANK: Record<string, number> = { traverse: 0, view: 1, edit: 2, admin: 3 }

// --- Key Service ---

export const VaultKeyService = {
  isUnlocked(): boolean {
    return derivedKey !== null
  },

  async isInitialized(db: PostgresJsDatabase<any>): Promise<boolean> {
    const [config] = await db.select({ id: vaultConfig.id }).from(vaultConfig).limit(1)
    return !!config
  },

  async initialize(db: PostgresJsDatabase<any>, masterPassword: string): Promise<void> {
    const existing = await this.isInitialized(db)
    if (existing) throw new Error('Vault is already initialized')

    const salt = randomBytes(32).toString('hex')
    const hash = await bcrypt.hash(masterPassword, 10)
    derivedKey = pbkdf2Sync(masterPassword, salt, PBKDF2_ITERATIONS, KEY_LENGTH, 'sha512')

    await db.insert(vaultConfig).values({
      masterKeySalt: salt,
      masterKeyHash: hash,
    })
  },

  async unlock(db: PostgresJsDatabase<any>, masterPassword: string): Promise<boolean> {
    const [config] = await db.select().from(vaultConfig).limit(1)
    if (!config) throw new Error('Vault is not initialized')

    const valid = await bcrypt.compare(masterPassword, config.masterKeyHash)
    if (!valid) return false

    derivedKey = pbkdf2Sync(masterPassword, config.masterKeySalt, PBKDF2_ITERATIONS, KEY_LENGTH, 'sha512')
    return true
  },

  lock(): void {
    derivedKey = null
  },

  async changeMasterPassword(db: PostgresJsDatabase<any>, currentPassword: string, newPassword: string): Promise<void> {
    if (!derivedKey) throw new Error('Vault is locked')

    const [config] = await db.select().from(vaultConfig).limit(1)
    if (!config) throw new Error('Vault is not initialized')

    const valid = await bcrypt.compare(currentPassword, config.masterKeyHash)
    if (!valid) throw new Error('Current password is incorrect')

    // Generate new key material
    const newSalt = randomBytes(32).toString('hex')
    const newHash = await bcrypt.hash(newPassword, 10)
    const newKey = pbkdf2Sync(newPassword, newSalt, PBKDF2_ITERATIONS, KEY_LENGTH, 'sha512')

    // Re-encrypt all entries
    const entries = await db.select().from(vaultEntries)
    for (const entry of entries) {
      if (entry.encryptedValue && entry.iv) {
        // Decrypt with old key
        const plaintext = decrypt(entry.encryptedValue, entry.iv)
        // Encrypt with new key
        const { ciphertext, iv } = encryptWith(plaintext, newKey)
        await db.update(vaultEntries).set({ encryptedValue: ciphertext, iv, updatedAt: new Date() }).where(eq(vaultEntries.id, entry.id))
      }
    }

    // Update config
    await db.update(vaultConfig).set({
      masterKeySalt: newSalt,
      masterKeyHash: newHash,
      updatedAt: new Date(),
    }).where(eq(vaultConfig.id, config.id))

    derivedKey = newKey
  },
}

// --- Encryption helpers ---

function encrypt(plaintext: string): { ciphertext: string; iv: string } {
  if (!derivedKey) throw new Error('Vault is locked')
  return encryptWith(plaintext, derivedKey)
}

function encryptWith(plaintext: string, key: Buffer): { ciphertext: string; iv: string } {
  const iv = randomBytes(12)
  const cipher = createCipheriv(ALGORITHM, key, iv)
  let encrypted = cipher.update(plaintext, 'utf8', 'hex')
  encrypted += cipher.final('hex')
  const authTag = cipher.getAuthTag().toString('hex')
  return {
    ciphertext: encrypted + ':' + authTag,
    iv: iv.toString('hex'),
  }
}

function decrypt(ciphertext: string, iv: string): string {
  if (!derivedKey) throw new Error('Vault is locked')
  const [encrypted, authTag] = ciphertext.split(':')
  const decipher = createDecipheriv(ALGORITHM, derivedKey, Buffer.from(iv, 'hex'))
  decipher.setAuthTag(Buffer.from(authTag, 'hex'))
  let decrypted = decipher.update(encrypted, 'hex', 'utf8')
  decrypted += decipher.final('utf8')
  return decrypted
}

// --- Permission Service ---

export const VaultPermissionService = {
  async hasAccess(db: PostgresJsDatabase<any>, categoryId: string, userId: string, minLevel: 'traverse' | 'view' | 'edit' | 'admin' = 'view'): Promise<boolean> {
    const level = await this.getAccessLevel(db, categoryId, userId)
    if (!level) return false
    return ACCESS_RANK[level] >= ACCESS_RANK[minLevel]
  },

  async getAccessLevel(db: PostgresJsDatabase<any>, categoryId: string, userId: string): Promise<'admin' | 'edit' | 'view' | 'traverse' | null> {
    const [category] = await db.select({ isPublic: vaultCategories.isPublic, createdBy: vaultCategories.createdBy })
      .from(vaultCategories).where(eq(vaultCategories.id, categoryId)).limit(1)
    if (!category) return null

    // Creator always has admin
    if (category.createdBy === userId) return 'admin'

    // Check direct user permission
    const [userPerm] = await db.select({ accessLevel: vaultCategoryPermissions.accessLevel }).from(vaultCategoryPermissions)
      .where(and(eq(vaultCategoryPermissions.categoryId, categoryId), eq(vaultCategoryPermissions.userId, userId)))
      .limit(1)
    if (userPerm) return userPerm.accessLevel

    // Check role-based permissions — pick highest across all roles
    const userRoleRows = await db.select({ roleId: userRoles.roleId }).from(userRoles).where(eq(userRoles.userId, userId))
    if (userRoleRows.length > 0) {
      const roleIds = userRoleRows.map((r) => r.roleId)
      const rolePerms = await db.select({ accessLevel: vaultCategoryPermissions.accessLevel }).from(vaultCategoryPermissions)
        .where(and(eq(vaultCategoryPermissions.categoryId, categoryId), inArray(vaultCategoryPermissions.roleId, roleIds)))
      let best: string | null = null
      for (const rp of rolePerms) {
        if (!best || ACCESS_RANK[rp.accessLevel] > ACCESS_RANK[best]) best = rp.accessLevel
      }
      if (best) return best as any
    }

    // Public categories give view access
    if (category.isPublic) return 'view'

    return null
  },

  async listPermissions(db: PostgresJsDatabase<any>, categoryId: string) {
    return db.select().from(vaultCategoryPermissions).where(eq(vaultCategoryPermissions.categoryId, categoryId))
  },

  async getPermissionById(db: PostgresJsDatabase<any>, permissionId: string) {
    const [perm] = await db.select().from(vaultCategoryPermissions).where(eq(vaultCategoryPermissions.id, permissionId)).limit(1)
    return perm ?? null
  },

  async setPermission(db: PostgresJsDatabase<any>, categoryId: string, roleId: string | undefined, userId: string | undefined, accessLevel: string) {
    if (roleId) {
      await db.delete(vaultCategoryPermissions).where(and(eq(vaultCategoryPermissions.categoryId, categoryId), eq(vaultCategoryPermissions.roleId, roleId)))
    }
    if (userId) {
      await db.delete(vaultCategoryPermissions).where(and(eq(vaultCategoryPermissions.categoryId, categoryId), eq(vaultCategoryPermissions.userId, userId)))
    }

    const [perm] = await db.insert(vaultCategoryPermissions).values({
      categoryId,
      roleId: roleId ?? null,
      userId: userId ?? null,
      accessLevel: accessLevel as any,
    }).returning()
    return perm
  },

  async removePermission(db: PostgresJsDatabase<any>, permissionId: string) {
    const [perm] = await db.delete(vaultCategoryPermissions).where(eq(vaultCategoryPermissions.id, permissionId)).returning()
    return perm ?? null
  },
}

// --- Category Service ---

export const VaultCategoryService = {
  async create(db: PostgresJsDatabase<any>, input: { name: string; description?: string; isPublic?: boolean; createdBy?: string }) {
    const [category] = await db.insert(vaultCategories).values({
      name: input.name,
      description: input.description ?? null,
      isPublic: input.isPublic ?? false,
      createdBy: input.createdBy ?? null,
    }).returning()
    return category
  },

  async getById(db: PostgresJsDatabase<any>, id: string) {
    const [category] = await db.select().from(vaultCategories).where(eq(vaultCategories.id, id)).limit(1)
    return category ?? null
  },

  async listAccessible(db: PostgresJsDatabase<any>, userId: string) {
    const all = await db.select().from(vaultCategories).orderBy(vaultCategories.name)
    const accessible = []
    for (const cat of all) {
      const level = await VaultPermissionService.getAccessLevel(db, cat.id, userId)
      if (level) accessible.push(cat)
    }
    return accessible
  },

  async update(db: PostgresJsDatabase<any>, id: string, input: { name?: string; description?: string; isPublic?: boolean }) {
    const [category] = await db.update(vaultCategories).set({ ...input, updatedAt: new Date() }).where(eq(vaultCategories.id, id)).returning()
    return category ?? null
  },

  async delete(db: PostgresJsDatabase<any>, id: string) {
    const [category] = await db.delete(vaultCategories).where(eq(vaultCategories.id, id)).returning()
    return category ?? null
  },
}

// --- Entry Service ---

export const VaultEntryService = {
  async create(db: PostgresJsDatabase<any>, input: {
    categoryId: string; name: string; username?: string; url?: string; notes?: string; createdBy?: string
  }, secret?: string) {
    let encryptedValue: string | null = null
    let iv: string | null = null
    if (secret) {
      const encrypted = encrypt(secret)
      encryptedValue = encrypted.ciphertext
      iv = encrypted.iv
    }

    const [entry] = await db.insert(vaultEntries).values({
      categoryId: input.categoryId,
      name: input.name,
      username: input.username ?? null,
      url: input.url ?? null,
      notes: input.notes ?? null,
      encryptedValue,
      iv,
      createdBy: input.createdBy ?? null,
      updatedBy: input.createdBy ?? null,
    }).returning()
    return entry
  },

  async getById(db: PostgresJsDatabase<any>, id: string) {
    const [entry] = await db.select().from(vaultEntries).where(eq(vaultEntries.id, id)).limit(1)
    return entry ?? null
  },

  async listByCategory(db: PostgresJsDatabase<any>, categoryId: string, params: PaginationInput) {
    const baseWhere = eq(vaultEntries.categoryId, categoryId)
    const searchCondition = params.q ? buildSearchCondition(params.q, [vaultEntries.name, vaultEntries.username]) : undefined
    const where = searchCondition ? and(baseWhere, searchCondition) : baseWhere

    const sortableColumns: Record<string, Column> = {
      name: vaultEntries.name,
      username: vaultEntries.username,
      created_at: vaultEntries.createdAt,
      updated_at: vaultEntries.updatedAt,
    }

    let query = db.select({
      id: vaultEntries.id,
      categoryId: vaultEntries.categoryId,
      name: vaultEntries.name,
      username: vaultEntries.username,
      url: vaultEntries.url,
      notes: vaultEntries.notes,
      createdBy: vaultEntries.createdBy,
      updatedBy: vaultEntries.updatedBy,
      createdAt: vaultEntries.createdAt,
      updatedAt: vaultEntries.updatedAt,
      // Expose whether a secret is set, but never the value
      hasSecret: vaultEntries.encryptedValue,
    }).from(vaultEntries).where(where).$dynamic()
    query = withSort(query, params.sort, params.order, sortableColumns, vaultEntries.name)
    query = withPagination(query, params.page, params.limit)

    const [data, [{ total }]] = await Promise.all([
      query,
      db.select({ total: count() }).from(vaultEntries).where(where),
    ])

    // Transform hasSecret from encrypted value to boolean
    const sanitized = data.map((e) => ({ ...e, hasSecret: !!e.hasSecret }))

    return paginatedResponse(sanitized, total, params.page, params.limit)
  },

  async reveal(id: string, db: PostgresJsDatabase<any>): Promise<string | null> {
    if (!VaultKeyService.isUnlocked()) throw new Error('Vault is locked')
    const [entry] = await db.select({ encryptedValue: vaultEntries.encryptedValue, iv: vaultEntries.iv })
      .from(vaultEntries).where(eq(vaultEntries.id, id)).limit(1)
    if (!entry?.encryptedValue || !entry.iv) return null
    return decrypt(entry.encryptedValue, entry.iv)
  },

  async update(db: PostgresJsDatabase<any>, id: string, input: {
    name?: string; username?: string; url?: string; notes?: string; updatedBy?: string
  }, newSecret?: string) {
    const updates: Record<string, any> = { ...input, updatedAt: new Date() }

    if (newSecret !== undefined) {
      if (newSecret) {
        const encrypted = encrypt(newSecret)
        updates.encryptedValue = encrypted.ciphertext
        updates.iv = encrypted.iv
      } else {
        updates.encryptedValue = null
        updates.iv = null
      }
    }

    const [entry] = await db.update(vaultEntries).set(updates).where(eq(vaultEntries.id, id)).returning()
    return entry ?? null
  },

  async delete(db: PostgresJsDatabase<any>, id: string) {
    const [entry] = await db.delete(vaultEntries).where(eq(vaultEntries.id, id)).returning()
    return entry ?? null
  },

  async search(db: PostgresJsDatabase<any>, query: string, params: PaginationInput) {
    const searchCondition = buildSearchCondition(query, [vaultEntries.name, vaultEntries.username])
    if (!searchCondition) return paginatedResponse([], 0, params.page, params.limit)

    let q = db.select({
      id: vaultEntries.id,
      categoryId: vaultEntries.categoryId,
      name: vaultEntries.name,
      username: vaultEntries.username,
      url: vaultEntries.url,
      notes: vaultEntries.notes,
      createdBy: vaultEntries.createdBy,
      updatedBy: vaultEntries.updatedBy,
      createdAt: vaultEntries.createdAt,
      updatedAt: vaultEntries.updatedAt,
      hasSecret: vaultEntries.encryptedValue,
    }).from(vaultEntries).where(searchCondition).$dynamic()
    q = withPagination(q, params.page, params.limit)

    const [data, [{ total }]] = await Promise.all([
      q,
      db.select({ total: count() }).from(vaultEntries).where(searchCondition),
    ])

    const sanitized = data.map((e) => ({ ...e, hasSecret: !!e.hasSecret }))
    return paginatedResponse(sanitized, total, params.page, params.limit)
  },
}