// Copyright 2020 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only

import { createWriteStream } from 'fs';
import { isNumber } from 'lodash';
import type { Readable } from 'stream';
import { Transform } from 'stream';
import { pipeline } from 'stream/promises';
import { ensureFile } from 'fs-extra';
import * as log from '../logging/log';
import * as Errors from '../types/errors';
import { strictAssert } from '../util/assert';
import {
  AttachmentSizeError,
  mightBeOnBackupTier,
  type AttachmentType,
  AttachmentVariant,
} from '../types/Attachment';
import * as Bytes from '../Bytes';
import {
  deriveBackupMediaKeyMaterial,
  type BackupMediaKeyMaterialType,
  deriveBackupMediaThumbnailInnerEncryptionKeyMaterial,
} from '../Crypto';
import {
  getAttachmentCiphertextLength,
  safeUnlinkSync,
  splitKeys,
  type ReencryptedAttachmentV2,
  decryptAndReencryptLocally,
  measureSize,
} from '../AttachmentCrypto';
import type { ProcessedAttachment } from './Types.d';
import type { WebAPIType } from './WebAPI';
import { createName, getRelativePath } from '../util/attachmentPath';
import { MediaTier } from '../types/AttachmentDownload';
import { getBackupKey } from '../services/backups/crypto';
import { backupsService } from '../services/backups';
import {
  getMediaIdForAttachment,
  getMediaIdForAttachmentThumbnail,
} from '../services/backups/util/mediaId';
import { MAX_BACKUP_THUMBNAIL_SIZE } from '../types/VisualAttachment';
import { missingCaseError } from '../util/missingCaseError';
import { IV_LENGTH, MAC_LENGTH } from '../types/Crypto';

const DEFAULT_BACKUP_CDN_NUMBER = 3;

export function getCdnKey(attachment: ProcessedAttachment): string {
  const cdnKey = attachment.cdnId || attachment.cdnKey;
  strictAssert(cdnKey, 'Attachment was missing cdnId or cdnKey');
  return cdnKey;
}

function getBackupMediaOuterEncryptionKeyMaterial(
  attachment: AttachmentType
): BackupMediaKeyMaterialType {
  const mediaId = getMediaIdForAttachment(attachment);
  const backupKey = getBackupKey();
  return deriveBackupMediaKeyMaterial(backupKey, mediaId.bytes);
}

function getBackupThumbnailInnerEncryptionKeyMaterial(
  attachment: AttachmentType
): BackupMediaKeyMaterialType {
  const mediaId = getMediaIdForAttachmentThumbnail(attachment);
  const backupKey = getBackupKey();
  return deriveBackupMediaThumbnailInnerEncryptionKeyMaterial(
    backupKey,
    mediaId.bytes
  );
}
function getBackupThumbnailOuterEncryptionKeyMaterial(
  attachment: AttachmentType
): BackupMediaKeyMaterialType {
  const mediaId = getMediaIdForAttachmentThumbnail(attachment);
  const backupKey = getBackupKey();
  return deriveBackupMediaKeyMaterial(backupKey, mediaId.bytes);
}

export async function getCdnNumberForBackupTier(
  attachment: ProcessedAttachment
): Promise<number> {
  strictAssert(
    attachment.backupLocator,
    'Attachment was missing backupLocator'
  );
  let backupCdnNumber = attachment.backupLocator.cdnNumber;

  if (backupCdnNumber == null) {
    const mediaId = getMediaIdForAttachment(attachment);
    const backupCdnInfo = await backupsService.getBackupCdnInfo(mediaId.string);
    if (backupCdnInfo.isInBackupTier) {
      backupCdnNumber = backupCdnInfo.cdnNumber;
    } else {
      backupCdnNumber = DEFAULT_BACKUP_CDN_NUMBER;
    }
  }
  return backupCdnNumber;
}

export async function downloadAttachment(
  server: WebAPIType,
  attachment: ProcessedAttachment,
  options: {
    variant?: AttachmentVariant;
    disableRetries?: boolean;
    timeout?: number;
    mediaTier?: MediaTier;
    logPrefix?: string;
  } = { variant: AttachmentVariant.Default }
): Promise<ReencryptedAttachmentV2 & { size?: number }> {
  const logId = `downloadAttachment/${options.logPrefix ?? ''}`;

  const { digest, key, size } = attachment;

  strictAssert(digest, `${logId}: missing digest`);
  strictAssert(key, `${logId}: missing key`);
  strictAssert(isNumber(size), `${logId}: missing size`);

  const mediaTier =
    options?.mediaTier ??
    (mightBeOnBackupTier(attachment) ? MediaTier.BACKUP : MediaTier.STANDARD);

  let downloadResult: Awaited<ReturnType<typeof downloadToDisk>>;

  if (mediaTier === MediaTier.STANDARD) {
    strictAssert(
      options.variant !== AttachmentVariant.ThumbnailFromBackup,
      'Thumbnails can only be downloaded from backup tier'
    );
    const cdnKey = getCdnKey(attachment);
    const { cdnNumber } = attachment;

    const downloadStream = await server.getAttachment({
      cdnKey,
      cdnNumber,
      options,
    });
    downloadResult = await downloadToDisk({ downloadStream, size });
  } else {
    const mediaId =
      options.variant === AttachmentVariant.ThumbnailFromBackup
        ? getMediaIdForAttachmentThumbnail(attachment)
        : getMediaIdForAttachment(attachment);

    const cdnNumber = await getCdnNumberForBackupTier(attachment);
    const cdnCredentials =
      await backupsService.credentials.getCDNReadCredentials(cdnNumber);

    const backupDir = await backupsService.api.getBackupDir();
    const mediaDir = await backupsService.api.getMediaDir();

    const downloadStream = await server.getAttachmentFromBackupTier({
      mediaId: mediaId.string,
      backupDir,
      mediaDir,
      headers: cdnCredentials.headers,
      cdnNumber,
      options,
    });
    downloadResult = await downloadToDisk({
      downloadStream,
      size: getAttachmentCiphertextLength(
        options.variant === AttachmentVariant.ThumbnailFromBackup
          ? // be generous, accept downloads of up to twice what we expect for thumbnail
            MAX_BACKUP_THUMBNAIL_SIZE * 2
          : size
      ),
    });
  }

  const { relativePath: downloadedRelativePath, downloadSize } = downloadResult;

  const cipherTextAbsolutePath =
    window.Signal.Migrations.getAbsoluteAttachmentPath(downloadedRelativePath);

  try {
    switch (options.variant) {
      case AttachmentVariant.Default:
      case undefined: {
        const { aesKey, macKey } = splitKeys(Bytes.fromBase64(key));
        return await decryptAndReencryptLocally({
          type: 'standard',
          ciphertextPath: cipherTextAbsolutePath,
          idForLogging: logId,
          aesKey,
          macKey,
          size,
          theirDigest: Bytes.fromBase64(digest),
          outerEncryption:
            mediaTier === 'backup'
              ? getBackupMediaOuterEncryptionKeyMaterial(attachment)
              : undefined,
          getAbsoluteAttachmentPath:
            window.Signal.Migrations.getAbsoluteAttachmentPath,
        });
      }
      case AttachmentVariant.ThumbnailFromBackup: {
        strictAssert(
          mediaTier === 'backup',
          'Thumbnail must be downloaded from backup tier'
        );
        const thumbnailEncryptionKeys =
          getBackupThumbnailInnerEncryptionKeyMaterial(attachment);
        // backup thumbnails don't get trimmed, so we just calculate the size as the
        // ciphertextSize, less IV and MAC
        const calculatedSize = downloadSize - IV_LENGTH - MAC_LENGTH;
        return {
          ...(await decryptAndReencryptLocally({
            type: 'backupThumbnail',
            ciphertextPath: cipherTextAbsolutePath,
            idForLogging: logId,
            size: calculatedSize,
            ...thumbnailEncryptionKeys,
            outerEncryption:
              getBackupThumbnailOuterEncryptionKeyMaterial(attachment),
            getAbsoluteAttachmentPath:
              window.Signal.Migrations.getAbsoluteAttachmentPath,
          })),
          size: calculatedSize,
        };
      }
      default: {
        throw missingCaseError(options.variant);
      }
    }
  } finally {
    safeUnlinkSync(cipherTextAbsolutePath);
  }
}

async function downloadToDisk({
  downloadStream,
  size,
}: {
  downloadStream: Readable;
  size: number;
}): Promise<{ relativePath: string; downloadSize: number }> {
  const relativeTargetPath = getRelativePath(createName());
  const absoluteTargetPath =
    window.Signal.Migrations.getAbsoluteAttachmentPath(relativeTargetPath);
  await ensureFile(absoluteTargetPath);
  const writeStream = createWriteStream(absoluteTargetPath);
  const targetSize = getAttachmentCiphertextLength(size);
  let downloadSize = 0;

  try {
    await pipeline(
      downloadStream,
      checkSize(targetSize),
      measureSize(bytesSeen => {
        downloadSize = bytesSeen;
      }),
      writeStream
    );
  } catch (error) {
    try {
      safeUnlinkSync(absoluteTargetPath);
    } catch (cleanupError) {
      log.error(
        'downloadToDisk: Error while cleaning up',
        Errors.toLogFormat(cleanupError)
      );
    }

    throw error;
  }

  return { relativePath: relativeTargetPath, downloadSize };
}

// A simple transform that throws if it sees more than maxBytes on the stream.
function checkSize(expectedBytes: number) {
  let totalBytes = 0;

  // TODO (DESKTOP-7046): remove size buffer
  const maximumSizeBeforeError = expectedBytes * 1.05;
  return new Transform({
    transform(chunk, encoding, callback) {
      totalBytes += chunk.byteLength;
      if (totalBytes > maximumSizeBeforeError) {
        callback(
          new AttachmentSizeError(
            `checkSize: Received ${totalBytes} bytes, max is ${maximumSizeBeforeError}`
          )
        );
        return;
      }

      if (totalBytes > expectedBytes) {
        log.warn(
          `checkSize: Received ${totalBytes} bytes, expected ${expectedBytes}`
        );
      }

      this.push(chunk, encoding);
      callback();
    },
  });
}