import { sha256 } from '@noble/hashes/sha2.js';
import { bytesToHex, utf8ToBytes } from '@noble/hashes/utils.js';
import { describe, expect, it } from 'vitest';

import {
  decryptSavedWalletMaterial,
  deriveWalletMaterial,
  encryptSecretParts,
  ensureSavedWalletMaterialOwner,
  getPlainWalletMaterial,
  isValidStoredWalletMaterial,
  linkingPubKeyHex,
  restoreWalletMaterialStored,
  savedWalletMaterialExists,
  savedWalletMaterialIsEncrypted,
  savedWalletMaterialOwnerId,
  saveWalletMaterial,
  serializeWalletMaterial,
  walletMaterialLinkingKey,
} from './keys';
import { parseJsonObject, stubLocalStorage } from './test-utils';

const FIXED_MNEMONIC =
  'abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about';
const PASSWORD = 'correct horse battery staple';
const MATERIAL_STORAGE_KEY = 'sattle_wallet_material_v2';
const LEGACY_STORAGE_KEY = 'sattle_linking_key';
const LINKING_KEY_HEX = 'da2a752d3d28668288ff50bed644e4a95f726d8711fd8754060f0a9e378f84aa';
const CASH_ROOT_HEX =
  'c7a2496e9b453a67c5d2a1f04936ec1259440d45454c795a99a66269e4cd3005111e1cc966fca2fe32f054f14caceab90449e536d94cf6935ea12a087e414f60';

const materialHash = (serialized: string): string => bytesToHex(sha256(utf8ToBytes(serialized)));

const readRawMaterialRecord = (): Record<string, unknown> => {
  const raw = localStorage.getItem(MATERIAL_STORAGE_KEY);
  if (raw === null) throw new Error('expected a saved wallet-material record');
  return parseJsonObject(raw);
};

describe('WalletMaterialV2 persistence', () => {
  it('round-trips canonical plaintext material in the new namespace', async () => {
    // Given complete wallet material and an unrelated legacy record
    stubLocalStorage();
    localStorage.setItem(
      LEGACY_STORAGE_KEY,
      JSON.stringify({ enc: false, value: '07'.repeat(32) }),
    );
    const material = deriveWalletMaterial(FIXED_MNEMONIC);

    // When the material is saved without a password
    await saveWalletMaterial(material);

    // Then only the new namespace carries the exact complete material
    expect(savedWalletMaterialExists()).toBe(true);
    expect(savedWalletMaterialIsEncrypted()).toBe(false);
    expect(getPlainWalletMaterial()).toEqual(material);
    expect(readRawMaterialRecord()).toEqual({
      enc: false,
      value: serializeWalletMaterial(material),
      materialHash: materialHash(serializeWalletMaterial(material)),
      ownerId: linkingPubKeyHex(walletMaterialLinkingKey(material)),
      version: 2,
    });
    expect(localStorage.getItem(LEGACY_STORAGE_KEY)).not.toBeNull();
    expect(localStorage.getItem(MATERIAL_STORAGE_KEY)).not.toContain(FIXED_MNEMONIC);
    expect(localStorage.getItem(MATERIAL_STORAGE_KEY)).not.toContain('"seed"');
  });

  it('round-trips canonical password-encrypted material without plaintext fields', async () => {
    // Given complete wallet material
    stubLocalStorage();
    const material = deriveWalletMaterial(FIXED_MNEMONIC);

    // When the material is password protected
    await saveWalletMaterial(material, PASSWORD);

    // Then only authenticated decryption returns both parts
    expect(savedWalletMaterialIsEncrypted()).toBe(true);
    expect(getPlainWalletMaterial()).toBeNull();
    expect(await decryptSavedWalletMaterial(PASSWORD)).toEqual(material);
    await expect(decryptSavedWalletMaterial('wrong password')).rejects.toThrow();
    const raw = JSON.stringify(readRawMaterialRecord());
    expect(raw).not.toContain(LINKING_KEY_HEX);
    expect(raw).not.toContain(CASH_ROOT_HEX);
    expect(raw).not.toContain(FIXED_MNEMONIC);
  });

  it('strips restored owner claims until the material proves its real owner', () => {
    // Given plaintext material carrying a foreign owner claim
    stubLocalStorage();
    const material = deriveWalletMaterial(FIXED_MNEMONIC);
    const foreignOwner = linkingPubKeyHex(new Uint8Array(32).fill(9));

    // When the record is restored and subsequently proven by plaintext unlock
    restoreWalletMaterialStored({
      enc: false,
      value: serializeWalletMaterial(material),
      materialHash: materialHash(serializeWalletMaterial(material)),
      ownerId: foreignOwner,
      version: 2,
    });
    const beforeProof = readRawMaterialRecord();
    const unlocked = getPlainWalletMaterial();
    if (unlocked === null) throw new Error('expected restored plaintext material');
    ensureSavedWalletMaterialOwner(unlocked);

    // Then the claim was absent before proof and only the derived owner was stamped
    expect(beforeProof).toEqual({
      enc: false,
      value: serializeWalletMaterial(material),
      materialHash: materialHash(serializeWalletMaterial(material)),
    });
    expect(savedWalletMaterialOwnerId()).toBe(linkingPubKeyHex(walletMaterialLinkingKey(material)));
  });

  it('rejects a plaintext owner mismatch without returning or stamping material', () => {
    // Given canonical material stamped with a foreign current owner
    stubLocalStorage();
    const material = deriveWalletMaterial(FIXED_MNEMONIC);
    localStorage.setItem(
      MATERIAL_STORAGE_KEY,
      JSON.stringify({
        enc: false,
        value: serializeWalletMaterial(material),
        materialHash: materialHash(serializeWalletMaterial(material)),
        ownerId: linkingPubKeyHex(new Uint8Array(32).fill(9)),
        version: 2,
      }),
    );
    const before = localStorage.getItem(MATERIAL_STORAGE_KEY);

    // When it is read or offered for proof
    const unlocked = getPlainWalletMaterial();
    const stamp = () => ensureSavedWalletMaterialOwner(material);

    // Then neither operation trusts or rewrites the foreign record
    expect(unlocked).toBeNull();
    expect(stamp).toThrow('different wallet');
    expect(localStorage.getItem(MATERIAL_STORAGE_KEY)).toBe(before);
  });

  it('rejects an encrypted owner mismatch after authentication', async () => {
    // Given valid encrypted material whose owner marker was replaced
    stubLocalStorage();
    const material = deriveWalletMaterial(FIXED_MNEMONIC);
    await saveWalletMaterial(material, PASSWORD);
    const foreign = {
      ...readRawMaterialRecord(),
      ownerId: linkingPubKeyHex(new Uint8Array(32).fill(9)),
    };
    localStorage.setItem(MATERIAL_STORAGE_KEY, JSON.stringify(foreign));
    const before = localStorage.getItem(MATERIAL_STORAGE_KEY);

    // When the authenticated material contradicts the current owner marker
    const unlock = decryptSavedWalletMaterial(PASSWORD);

    // Then no material is returned or marker changed
    await expect(unlock).rejects.toThrow('different wallet');
    expect(localStorage.getItem(MATERIAL_STORAGE_KEY)).toBe(before);
  });

  it.each([
    ['legacy key-only plaintext', LINKING_KEY_HEX, 2],
    ['wrong outer version', serializeWalletMaterial(deriveWalletMaterial(FIXED_MNEMONIC)), 1],
    [
      'malformed material plaintext',
      JSON.stringify({ version: 2, linkingKeyHex: LINKING_KEY_HEX, cashRootHex: '00' }),
      2,
    ],
  ])('rejects %s records at the storage boundary', (_name, value, version) => {
    // Given an invalid record under the v2 material namespace
    stubLocalStorage();
    const material = deriveWalletMaterial(FIXED_MNEMONIC);
    const record = {
      enc: false,
      value,
      materialHash: materialHash(value),
      ownerId: linkingPubKeyHex(walletMaterialLinkingKey(material)),
      version,
    };
    localStorage.setItem(MATERIAL_STORAGE_KEY, JSON.stringify(record));

    // When storage helpers parse it
    // Then the record is unusable as wallet material
    expect(isValidStoredWalletMaterial(record)).toBe(false);
    expect(savedWalletMaterialExists()).toBe(false);
    expect(getPlainWalletMaterial()).toBeNull();
  });

  it('rejects authenticated ciphertext containing malformed material', async () => {
    // Given valid AES-GCM parts whose plaintext is not a v2 material envelope
    stubLocalStorage();
    const parts = await encryptSecretParts(LINKING_KEY_HEX, PASSWORD);
    const material = deriveWalletMaterial(FIXED_MNEMONIC);
    localStorage.setItem(
      MATERIAL_STORAGE_KEY,
      JSON.stringify({
        enc: true,
        ...parts,
        materialHash: materialHash(LINKING_KEY_HEX),
        ownerId: linkingPubKeyHex(walletMaterialLinkingKey(material)),
        version: 2,
      }),
    );

    // When the password successfully authenticates the malformed plaintext
    const unlock = decryptSavedWalletMaterial(PASSWORD);

    // Then it still fails closed as invalid material
    await expect(unlock).rejects.toThrow('valid v2 wallet material');
  });
});