import { subtle, getRandomValues } from "node:crypto";
import { fromExtendedJson, toExtendedJson } from "@xo-cash/utils";
import { sha256, binToBase64 } from "@bitauth/libauth";

import { BaseSyncClient, type ResourcesRecord } from "../../src/shared/base-client.js";

/**
 * Encrypts every resource with a single shared key.
 */
export class EncryptedClient<T = unknown> extends BaseSyncClient<T> {
  private readonly encryptionKey: Uint8Array;

  constructor(private readonly client: BaseSyncClient<string>, encryptionKey: Uint8Array | string) {
    super();

    // Hash the encryption key into 32 bytes if it is a string
    this.encryptionKey = typeof encryptionKey === 'string' ? sha256.hash(new TextEncoder().encode(encryptionKey)) : encryptionKey;

    this.client.on('resource-instance-updated', async (message) => {
      try {
        this.emit('resource-instance-updated', {
          ...message,
          value: await EncryptedClient.decrypt(this.encryptionKey, message.value) as T,
        });
      } catch (error) {
        this.emit('error', error instanceof Error ? error : new Error(String(error)));
      }
    });
  }

  /**
   * Decrypts every instance of every resource in a record.
   * @param resources - The encrypted resources record
   * @returns The resources record with decrypted values
   */
  private async decryptResources(resources: ResourcesRecord<string>): Promise<ResourcesRecord<T>> {
    const entries = await Promise.all(Object.entries(resources).map(async ([ resourceId, resource ]) => {
      const instances = await Promise.all(resource.instances.map(async (instance) => ({
        ...instance,
        value: await EncryptedClient.decrypt(this.encryptionKey, instance.value) as T,
      })));

      return [ resourceId, { ...resource, instances } ] as const;
    }));

    return Object.fromEntries(entries);
  }

  public async connect(): Promise<void> {
    return this.client.connect();
  }

  public async disconnect(): Promise<void> {
    return this.client.disconnect();
  }

  public async readMany(resourceIds: string[]): Promise<ResourcesRecord<T>> {
    const encryptedResources = await this.client.readMany(resourceIds);

    return this.decryptResources(encryptedResources);
  }

  public async writeMany(resources: Record<string, T>): Promise<ResourcesRecord<T>> {
    const encryptedEntries = await Promise.all(Object.entries(resources).map(async ([ resourceId, value ]) => {
      return [ resourceId, await EncryptedClient.encrypt(this.encryptionKey, value) ] as const;
    }));

    const response = await this.client.writeMany(Object.fromEntries(encryptedEntries));

    return this.decryptResources(response);
  }

  public async subscribe(resourceId: string): Promise<void> {
    return this.client.subscribe(resourceId);
  }

  public async unsubscribe(resourceId: string): Promise<void> {
    return this.client.unsubscribe(resourceId);
  }

  /**
   * Imports raw key material as an AES-GCM key
   * @param keyMaterial - The 32 byte key material
   * @returns The AES-GCM CryptoKey
   */
  private static importKey(keyMaterial: Uint8Array) {
    return subtle.importKey('raw', Buffer.from(keyMaterial), { name: 'AES-GCM' }, false, ['encrypt', 'decrypt']);
  }

  /**
   * Encrypts a value using AES-GCM with a random 12 byte IV
   * @param keyMaterial - The encryption key material
   * @param value - The value to encrypt
   * @returns The IV and ciphertext as a base64 string
   */
  public static async encrypt(keyMaterial: Uint8Array, value: unknown): Promise<string> {
    const key = await this.importKey(keyMaterial);
    const iv = getRandomValues(new Uint8Array(12));

    // Encode the value as extended JSON and then to UTF-8 bytes
    const encodedValue = new TextEncoder().encode(toExtendedJson(value));

    const encrypted = await subtle.encrypt({ name: 'AES-GCM', iv }, key, encodedValue);

    // Prepend the IV so it can be recovered when decrypting
    return binToBase64(Buffer.concat([ iv, new Uint8Array(encrypted) ]));
  }

  /**
   * Decrypts a value encrypted by `encrypt`
   * @param keyMaterial - The encryption key material
   * @param value - The IV and ciphertext as a base64 string
   * @returns The decrypted value
   */
  public static async decrypt(keyMaterial: Uint8Array, value: string): Promise<unknown> {
    const key = await this.importKey(keyMaterial);

    // Decode the base64 first: the IV is 12 bytes, not 12 base64 characters
    const bytes = Buffer.from(value, 'base64');
    const iv = bytes.subarray(0, 12);
    const encryptedValue = bytes.subarray(12);

    const decrypted = await subtle.decrypt({ name: 'AES-GCM', iv }, key, encryptedValue);

    return fromExtendedJson(new TextDecoder().decode(decrypted));
  }
}