/**
 * Account Encryption Extension
 *
 * Prisma Client Extension that transparently encrypts/decrypts OAuth tokens
 * in the Account table using VaultService (OpenBao Transit or AES-256-GCM fallback).
 *
 * Migrated from $use() middleware (removed in Prisma 5+) to $extends() API.
 */

import { Logger } from "@nestjs/common";
import { Prisma } from "@prisma/client";
import type { VaultService } from "../vault/vault.service";
import { TransitKey } from "../vault/vault.constants";

const TOKEN_FIELDS = ["accessToken", "refreshToken", "idToken"] as const;

interface AccountData extends Record<string, unknown> {
  accessToken?: string | null;
  refreshToken?: string | null;
  idToken?: string | null;
  encryptionVersion?: string | null;
}

export function createAccountEncryptionExtension(vaultService: VaultService) {
  const logger = new Logger("AccountEncryptionExtension");

  return Prisma.defineExtension({
    name: "account-encryption",
    query: {
      account: {
        async create({ args, query }) {
          await encryptTokens(args.data as AccountData, vaultService);
          return query(args);
        },
        async update({ args, query }) {
          await encryptTokens(args.data as AccountData, vaultService);
          return query(args);
        },
        async updateMany({ args, query }) {
          await encryptTokens(args.data as AccountData, vaultService);
          return query(args);
        },
        async upsert({ args, query }) {
          await encryptTokens(args.create as AccountData, vaultService);
          await encryptTokens(args.update as AccountData, vaultService);
          return query(args);
        },
        async findUnique({ args, query }) {
          const result = await query(args);
          if (result && typeof result === "object") {
            await decryptTokens(result as AccountData, vaultService, logger);
          }
          return result;
        },
        async findFirst({ args, query }) {
          const result = await query(args);
          if (result && typeof result === "object") {
            await decryptTokens(result as AccountData, vaultService, logger);
          }
          return result;
        },
        async findMany({ args, query }) {
          const results = await query(args);
          for (const account of results) {
            await decryptTokens(account as AccountData, vaultService, logger);
          }
          return results;
        },
      },
    },
  });
}

async function encryptTokens(data: AccountData, vaultService: VaultService): Promise<void> {
  let encrypted = false;
  let encryptionVersion: "aes" | "vault" | null = null;

  for (const field of TOKEN_FIELDS) {
    const value = data[field];
    if (value == null) continue;
    if (typeof value === "string" && isEncrypted(value)) continue;

    if (typeof value === "string") {
      const ciphertext = await vaultService.encrypt(value, TransitKey.ACCOUNT_TOKENS);
      data[field] = ciphertext;
      encrypted = true;
      encryptionVersion = ciphertext.startsWith("vault:v1:") ? "vault" : "aes";
    }
  }

  if (encrypted && encryptionVersion) {
    data.encryptionVersion = encryptionVersion;
  }
}

async function decryptTokens(
  account: AccountData,
  vaultService: VaultService,
  _logger: Logger
): Promise<void> {
  const shouldDecrypt =
    account.encryptionVersion === "aes" || account.encryptionVersion === "vault";

  for (const field of TOKEN_FIELDS) {
    const value = account[field];
    if (value == null) continue;

    if (typeof value === "string") {
      if (shouldDecrypt) {
        try {
          account[field] = await vaultService.decrypt(value, TransitKey.ACCOUNT_TOKENS);
        } catch (error) {
          const errorMsg = error instanceof Error ? error.message : "Unknown error";
          throw new Error(
            `Failed to decrypt account credentials. Please reconnect this account. Details: ${errorMsg}`
          );
        }
      } else if (!account.encryptionVersion && isEncrypted(value)) {
        try {
          account[field] = await vaultService.decrypt(value, TransitKey.ACCOUNT_TOKENS);
        } catch (error) {
          const errorMsg = error instanceof Error ? error.message : "Unknown error";
          throw new Error(
            `Failed to decrypt account credentials. Please reconnect this account. Details: ${errorMsg}`
          );
        }
      }
    }
  }
}

function isEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") return false;
  if (isAESEncrypted(value)) return true;
  if (value.startsWith("vault:v1:")) return true;
  return false;
}

function isAESEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") return false;
  const parts = value.split(":");
  if (parts.length !== 3) return false;
  return parts.every((part) => /^[0-9a-f]+$/i.test(part));
}
