src/db/r2.ts

import {
  DeleteObjectCommand,
  GetObjectCommand,
  PutObjectCommand,
  S3Client,
} from '@aws-sdk/client-s3';
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
import * as fs from 'fs';
import * as path from 'path';
import { pipeline } from 'stream/promises';
import { Readable } from 'stream';
import { config } from '../config';

const R2_PREFIX = 'r2://';
let client: S3Client | null = null;

export interface R2ObjectRef {
  bucket: string;
  key: string;
}

function getR2Client(): S3Client {
  if (!client) {
    client = new S3Client({
      region: config.r2Region,
      endpoint: config.r2Endpoint,
      forcePathStyle: true,
      credentials: {
        accessKeyId: config.r2AccessKeyId,
        secretAccessKey: config.r2SecretAccessKey,
      },
    });
  }
  return client;
}

export function isR2ObjectRef(value: string): boolean {
  return value.startsWith(R2_PREFIX);
}

export function createR2ObjectRef(bucket: string, key: string): string {
  return `${R2_PREFIX}${bucket}/${key.replace(/^\/+/, '')}`;
}

export function parseR2ObjectRef(value: string): R2ObjectRef {
  if (!isR2ObjectRef(value)) {
    throw new Error('Not an R2 object reference');
  }

  const withoutScheme = value.slice(R2_PREFIX.length);
  const slashIndex = withoutScheme.indexOf('/');
  if (slashIndex <= 0 || slashIndex === withoutScheme.length - 1) {
    throw new Error('Invalid R2 object reference');
  }

  return {
    bucket: withoutScheme.slice(0, slashIndex),
    key: withoutScheme.slice(slashIndex + 1),
  };
}

export async function putR2Buffer(
  bucket: string,
  key: string,
  body: Buffer,
  contentType: string,
): Promise<string> {
  await getR2Client().send(
    new PutObjectCommand({
      Bucket: bucket,
      Key: key,
      Body: body,
      ContentLength: body.length,
      ContentType: contentType,
    }),
  );
  return createR2ObjectRef(bucket, key);
}

export async function putR2File(
  bucket: string,
  key: string,
  localPath: string,
  contentType: string,
): Promise<string> {
  const stat = await fs.promises.stat(localPath);
  await getR2Client().send(
    new PutObjectCommand({
      Bucket: bucket,
      Key: key,
      Body: fs.createReadStream(localPath),
      ContentLength: stat.size,
      ContentType: contentType,
    }),
  );
  return createR2ObjectRef(bucket, key);
}

export async function downloadR2File(
  objectRef: string,
  localPath: string,
): Promise<void> {
  const { bucket, key } = parseR2ObjectRef(objectRef);
  const response = await getR2Client().send(
    new GetObjectCommand({ Bucket: bucket, Key: key }),
  );
  if (!response.Body) throw new Error(`R2 object body is empty: ${objectRef}`);

  await fs.promises.mkdir(path.dirname(localPath), { recursive: true });
  await pipeline(response.Body as Readable, fs.createWriteStream(localPath));
}

export async function readR2Text(objectRef: string): Promise<string | null> {
  const { bucket, key } = parseR2ObjectRef(objectRef);
  try {
    const response = await getR2Client().send(
      new GetObjectCommand({ Bucket: bucket, Key: key }),
    );
    if (!response.Body) return null;
    return await response.Body.transformToString('utf-8');
  } catch (error: unknown) {
    const name = error instanceof Error ? error.name : '';
    const status = (error as { $metadata?: { httpStatusCode?: number } })?.$metadata
      ?.httpStatusCode;
    if (name === 'NoSuchKey' || name === 'NotFound' || status === 404) return null;
    throw error;
  }
}

export async function deleteR2Object(objectRef: string): Promise<void> {
  const { bucket, key } = parseR2ObjectRef(objectRef);
  await getR2Client().send(
    new DeleteObjectCommand({ Bucket: bucket, Key: key }),
  );
}

export async function createR2SignedDownloadUrl(
  objectRef: string,
  expiresInSeconds: number,
  downloadFileName?: string,
): Promise<string> {
  const { bucket, key } = parseR2ObjectRef(objectRef);
  const disposition = downloadFileName
    ? `attachment; filename*=UTF-8''${encodeURIComponent(downloadFileName)}`
    : undefined;
  return getSignedUrl(
    getR2Client(),
    new GetObjectCommand({
      Bucket: bucket,
      Key: key,
      ResponseContentDisposition: disposition,
    }),
    { expiresIn: expiresInSeconds },
  );
}