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

import { BaseSyncClient, type Resource, type ResourcesRecord } from '../../source/shared/base-client.js';

type ResourceKey = {
    resourceId: string;
    keyMaterial: Uint8Array;
};

/**
 * Encrypts each resource with its own key, where the resource ID itself is the secret.
 *
 * The caller-facing resource ID is hashed into a 32 byte AES key, and that key is hashed again
 * to get the resource ID sent to the server. The server never sees the AES key or the original
 * resource ID, so the resource ID must be unguessable (high entropy) for this to be secure.
 */
export class EncryptedClientPerResource<T = unknown> extends BaseSyncClient<T> {
    /** Maps the hashed (server-side) resource ID to the original resource ID and its key */
    private readonly resourceKeys = new Map<string, ResourceKey>();

    constructor(private readonly client: BaseSyncClient<string>) {
        super();

        this.client.on('resource-instance-updated', async (message) => {
            // Ignore updates for resources we don't hold the key for
            const resourceKey = this.resourceKeys.get(message.resourceId);
            if (!resourceKey) return;

            try {
                this.emit('resource-instance-updated', {
                    ...message,
                    resourceId: resourceKey.resourceId,
                    value: (await EncryptedClientPerResource.decrypt(resourceKey.keyMaterial, message.value)) as T,
                });
            } catch (error) {
                this.emit('error', error instanceof Error ? error : new Error(String(error)));
            }
        });
    }

    /**
     * Derives the AES key and server-side resource ID from a resource ID, and remembers the mapping.
     * @param resourceId - The resource ID (the secret)
     * @returns The hashed resource ID to send to the server
     */
    private hashResourceId(resourceId: string): string {
        const keyMaterial = sha256.hash(new TextEncoder().encode(resourceId));
        const hashedResourceId = binToHex(sha256.hash(keyMaterial));

        this.resourceKeys.set(hashedResourceId, { resourceId, keyMaterial });

        return hashedResourceId;
    }

    /**
     * Gets the original resource ID and key for a hashed resource ID.
     * @param hashedResourceId - The hashed (server-side) resource ID
     * @returns The original resource ID and its key material
     */
    private getResourceKey(hashedResourceId: string): ResourceKey {
        const resourceKey = this.resourceKeys.get(hashedResourceId);

        if (!resourceKey) {
            throw new Error(`No key found for resource ${hashedResourceId}`);
        }

        return resourceKey;
    }

    /**
     * Decrypts every instance of every resource in a record, and maps the hashed resource IDs back to the original resource IDs.
     * @param resources - The encrypted resources record, keyed by hashed resource ID
     * @returns The resources record with decrypted values, keyed by original resource ID
     */
    private async decryptResources(resources: ResourcesRecord<string>): Promise<ResourcesRecord<T>> {
        const entries = await Promise.all(
            Object.entries(resources).map(async ([hashedResourceId, resource]) => {
                const { resourceId, keyMaterial } = this.getResourceKey(hashedResourceId);

                const instances: Resource<T>[] = await Promise.all(
                    resource.instances.map(async (instance) => ({
                        ...instance,
                        resourceId,
                        value: (await EncryptedClientPerResource.decrypt(keyMaterial, 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 hashedResourceIds = resourceIds.map((resourceId) => this.hashResourceId(resourceId));

        const encryptedResources = await this.client.readMany(hashedResourceIds);

        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]) => {
                const hashedResourceId = this.hashResourceId(resourceId);
                const { keyMaterial } = this.getResourceKey(hashedResourceId);

                return [hashedResourceId, await EncryptedClientPerResource.encrypt(keyMaterial, 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(this.hashResourceId(resourceId));
    }

    public async unsubscribe(resourceId: string): Promise<void> {
        return this.client.unsubscribe(this.hashResourceId(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));
    }
}