// src/lib/krux.ts
// Krux KEF (Krux Encryption Format) implementation
import * as pako from 'pako';
import { base43Decode, base43Encode } from './base43';
import { getWalletFingerprint } from './bip32';
export const VERSIONS: Record<number, {
  name: string;
  compress: boolean;
  auth: number;
}> = {
  // We only implement the GCM versions as they are the only ones compatible with WebCrypto
  20: { name: "AES-GCM", compress: false, auth: 4 },
  21: { name: "AES-GCM +c", compress: true, auth: 4 },
};

const GCM_IV_LENGTH = 12;

function toArrayBuffer(data: Uint8Array): ArrayBuffer {
  // Create a new ArrayBuffer and copy the contents
  const buffer = new ArrayBuffer(data.byteLength);
  new Uint8Array(buffer).set(data);
  return buffer;
}

export function unwrap(envelope: Uint8Array): { label: string; labelBytes: Uint8Array, version: number; iterations: number; payload: Uint8Array } {
  if (envelope.length < 5) throw new Error("Invalid KEF envelope: too short");
  const lenId = envelope[0];
  if (!(0 <= lenId && lenId <= 252)) throw new Error("Invalid label length in KEF envelope");
  if (1 + lenId + 4 > envelope.length) throw new Error("Invalid KEF envelope: insufficient data");

  const labelBytes = envelope.subarray(1, 1 + lenId);
  const label = new TextDecoder().decode(labelBytes);
  const version = envelope[1 + lenId];
  if (!VERSIONS[version]) {
    throw new Error(`Unsupported KEF version: ${version}`);
  }

  const iterStart = 2 + lenId;
  let iters = (envelope[iterStart] << 16) | (envelope[iterStart + 1] << 8) | envelope[iterStart + 2];
  const iterations = iters <= 10000 ? iters * 10000 : iters;
  const payload = envelope.subarray(5 + lenId);
  return { label, labelBytes, version, iterations, payload };
}

import { pbkdf2HmacSha256 } from './pbkdf2';
import { entropyToMnemonic, mnemonicToEntropy } from './seedblend';

// ... (rest of the file is the same until KruxCipher)

export class KruxCipher {
  private keyPromise: Promise<CryptoKey>;

  constructor(passphrase: string, salt: Uint8Array, iterations: number) {
    this.keyPromise = (async () => {
      // Use pure-JS PBKDF2 implementation which has been validated against Krux's test vector
      const derivedKeyBytes = await pbkdf2HmacSha256(passphrase, salt, iterations, 32);

      // Import the derived bytes as an AES-GCM key
      return crypto.subtle.importKey(
        "raw",
        toArrayBuffer(derivedKeyBytes),
        { name: "AES-GCM", length: 256 },
        false,
        ["encrypt", "decrypt"]
      );
    })();
  }

  // Encrypt function is unused in SeedBlender, but kept for completeness
  async encrypt(plaintext: Uint8Array, version = 20, iv?: Uint8Array): Promise<Uint8Array> {
    const v = VERSIONS[version];
    if (!v) throw new Error(`Unsupported KEF version: ${version}`);

    let dataToEncrypt = plaintext;
    if (v.compress) {
      dataToEncrypt = pako.deflate(plaintext);
    }

    let ivBytes = iv ? new Uint8Array(iv) : crypto.getRandomValues(new Uint8Array(GCM_IV_LENGTH));

    const key = await this.keyPromise;
    const plaintextBuffer = toArrayBuffer(dataToEncrypt);
    const ivBuffer = toArrayBuffer(ivBytes);
    const tagLengthBits = v.auth * 8;

    const encrypted = await crypto.subtle.encrypt({ name: "AES-GCM", iv: ivBuffer, tagLength: tagLengthBits }, key, plaintextBuffer);
    const encryptedBytes = new Uint8Array(encrypted);
    const ciphertext = encryptedBytes.slice(0, encryptedBytes.length - v.auth);
    const tag = encryptedBytes.slice(encryptedBytes.length - v.auth);

    const combined = new Uint8Array(ivBytes.length + ciphertext.length + tag.length);
    combined.set(ivBytes, 0);
    combined.set(ciphertext, ivBytes.length);
    combined.set(tag, ivBytes.length + ciphertext.length);
    return combined;
  }

  async decrypt(payload: Uint8Array, version: number): Promise<Uint8Array> {
    const v = VERSIONS[version];
    if (!v) throw new Error(`Unsupported KEF version: ${version}`);
    if (payload.length < GCM_IV_LENGTH + v.auth) throw new Error("Payload too short for AES-GCM");

    const iv = payload.slice(0, GCM_IV_LENGTH);
    const ciphertext = payload.slice(GCM_IV_LENGTH, payload.length - v.auth);
    const tag = payload.slice(payload.length - v.auth);

    const key = await this.keyPromise;

    try {
      const ciphertextWithTag = new Uint8Array(ciphertext.length + tag.length);
      ciphertextWithTag.set(ciphertext, 0);
      ciphertextWithTag.set(tag, ciphertext.length);

      const decryptedBuffer = await crypto.subtle.decrypt(
        { name: "AES-GCM", iv: toArrayBuffer(iv), tagLength: v.auth * 8 }, key, toArrayBuffer(ciphertextWithTag)
      );

      let decrypted = new Uint8Array(decryptedBuffer);

      if (v.compress) {
        decrypted = pako.inflate(decrypted);
      }
      return decrypted;
    } catch (error) {
      console.error("Krux decryption internal error:", error);
      throw new Error("Krux decryption failed - wrong passphrase or corrupted data");
    }
  }
}

export function hexToBytes(hex: string): Uint8Array {
  const cleaned = hex.trim().replace(/\s/g, '').replace(/^KEF:/i, '');
  if (!/^[0-9a-fA-F]+$/.test(cleaned)) throw new Error("Invalid hex string");
  if (cleaned.length % 2 !== 0) throw new Error("Hex string must have even length");
  const bytes = new Uint8Array(cleaned.length / 2);
  for (let i = 0; i < bytes.length; i++) {
    bytes[i] = parseInt(cleaned.substr(i * 2, 2), 16);
  }
  return bytes;
}

export async function decryptFromKrux(params: { kefData: string; passphrase: string; }): Promise<{ mnemonic: string; label: string; version: number; iterations: number }> {
  const { kefData, passphrase } = params;

  // STEP 1: Validate and decode data format (FIRST!)
  let bytes: Uint8Array;
  try {
    bytes = hexToBytes(kefData);
  } catch (e) {
    try {
      bytes = base43Decode(kefData);
    } catch (e2) {
      throw new Error("Invalid Krux data: Not a valid Hex or Base43 string.");
    }
  }

  // STEP 2: Unwrap and validate envelope structure
  let label: string, labelBytes: Uint8Array, version: number, iterations: number, payload: Uint8Array;
  try {
    const unwrapped = unwrap(bytes);
    label = unwrapped.label;
    labelBytes = unwrapped.labelBytes;
    version = unwrapped.version;
    iterations = unwrapped.iterations;
    payload = unwrapped.payload;
  } catch (e: any) {
    throw new Error("Invalid Krux data: Not a valid Hex or Base43 string.");
  }

  // STEP 3: Check passphrase (only after data structure is validated)
  if (!passphrase) {
    throw new Error("Passphrase is required for Krux decryption");
  }

  // STEP 4: Decrypt
  const cipher = new KruxCipher(passphrase, labelBytes, iterations);
  const decrypted = await cipher.decrypt(payload, version);
  const mnemonic = await entropyToMnemonic(decrypted);

  return { mnemonic, label, version, iterations };
}

export function bytesToHex(bytes: Uint8Array): string {
  return Array.from(bytes).map(b => b.toString(16).padStart(2, '0')).join('').toUpperCase();
}

export async function encryptToKrux(params: {
  mnemonic: string;
  passphrase: string;
}): Promise<{ kefBase43: string; label: string; version: number; iterations: number }> {
  const { mnemonic, passphrase } = params;

  if (!passphrase) throw new Error("Passphrase is required");

  const label = getWalletFingerprint(mnemonic);
  const iterations = 100000;
  const version = 20;

  const mnemonicBytes = await mnemonicToEntropy(mnemonic);
  const cipher = new KruxCipher(passphrase, new TextEncoder().encode(label), iterations);
  const payload = await cipher.encrypt(mnemonicBytes, version);
  const kef = wrap(label, version, iterations, payload);
  const kefBase43 = base43Encode(kef);

  // Debug logging disabled in production to prevent seed recovery via console history
  if (import.meta.env.DEV) {
    console.debug('KEF encryption completed', { version, iterations });
  }

  return { kefBase43, label, version, iterations };
}

export function wrap(label: string, version: number, iterations: number, payload: Uint8Array): Uint8Array {
  const labelBytes = new TextEncoder().encode(label);
  const idLen = labelBytes.length;

  // ADD THIS:
  if (idLen > 252) {
    throw new Error('Label too long');
  }

  // Convert iterations to 3 bytes (Big-Endian)
  // Scale down if > 10000 (Krux format: stores scaled value)
  let scaledIter: number;
  if (iterations >= 10000 && iterations % 10000 === 0) {
    // Divisible by 10000 - store scaled
    scaledIter = Math.floor(iterations / 10000);
  } else {
    // Store as-is (handles edge cases like 10001)
    scaledIter = iterations;
  }
  const iterBytes = new Uint8Array(3);
  iterBytes[0] = (scaledIter >> 16) & 0xFF;
  iterBytes[1] = (scaledIter >> 8) & 0xFF;
  iterBytes[2] = scaledIter & 0xFF;

  // Calculate total length
  const totalLength = 1 + idLen + 1 + 3 + payload.length;
  const envelope = new Uint8Array(totalLength);

  let offset = 0;
  envelope[offset++] = idLen;
  envelope.set(labelBytes, offset);
  offset += idLen;
  envelope[offset++] = version;
  envelope.set(iterBytes, offset);
  offset += 3;
  envelope.set(payload, offset);

  return envelope;
}