import type { CryptoProvider, StorageProvider } from '@shade/core';
import {
  toBase64,
  fromBase64,
  serializeIdentityKeyPair,
  deserializeIdentityKeyPair,
  serializeSignedPreKey,
  deserializeSignedPreKey,
  serializeOneTimePreKey,
  deserializeOneTimePreKey,
  serializeSessionState,
  deserializeSessionState,
} from '@shade/core';

/**
 * Shade backup format v1.
 *
 * A passphrase-encrypted blob containing the entire local Shade state:
 * identity, signed prekeys, one-time prekeys, sessions, trusted identities.
 *
 * Encryption: AES-256-GCM. Key derivation from passphrase:
 * HKDF-SHA256(passphrase, randomSalt, "ShadeBackupKey", 32)
 *
 * Note: HKDF is NOT a proper password KDF. For a real password-based backup
 * you should use Argon2id. This implementation uses HKDF as a simple
 * placeholder because CryptoProvider doesn't expose Argon2. Upgrade path
 * is to add `argon2id` to CryptoProvider and swap it in here. For now,
 * DOCUMENT THIS and require the user to choose a high-entropy passphrase.
 */

const BACKUP_VERSION = 1;
const BACKUP_INFO = new TextEncoder().encode('ShadeBackupKey');
const SALT_SIZE = 32;

export interface BackupPayload {
  version: number;
  identity: string | null;
  registrationId: number;
  signedPreKeys: Array<{ keyId: number; data: string }>;
  oneTimePreKeys: Array<{ keyId: number; data: string }>;
  sessions: Array<{ address: string; state: string }>;
  trustedIdentities: Array<{ address: string; key: string }>;
}

export interface BackupBlob {
  version: number;
  salt: string; // base64
  nonce: string; // base64
  ciphertext: string; // base64
}

/**
 * Export the entire Shade state as an encrypted backup blob.
 *
 * @param crypto CryptoProvider instance
 * @param storage Storage to read from
 * @param passphrase User-chosen passphrase (HIGH ENTROPY REQUIRED)
 * @param knownAddresses Addresses whose sessions should be included.
 *        Pass an empty list to back up identity + prekeys only (no sessions).
 */
export async function exportBackup(
  crypto: CryptoProvider,
  storage: StorageProvider,
  passphrase: string,
  knownAddresses: string[] = [],
): Promise<BackupBlob> {
  if (passphrase.length < 12) {
    throw new Error('Passphrase must be at least 12 characters');
  }

  // Gather state
  const identity = await storage.getIdentityKeyPair();
  const registrationId = await storage.getLocalRegistrationId();

  const signedPreKeys: Array<{ keyId: number; data: string }> = [];
  // SignedPreKeys: we can't enumerate them from the interface; try IDs 1-100
  for (let i = 1; i <= 100; i++) {
    const spk = await storage.getSignedPreKey(i);
    if (spk) signedPreKeys.push({ keyId: i, data: serializeSignedPreKey(spk) });
  }

  const oneTimePreKeys: Array<{ keyId: number; data: string }> = [];
  const otpkCount = await storage.getOneTimePreKeyCount();
  // Try a reasonable ID range
  for (let i = 1; i <= otpkCount + 1000; i++) {
    const otpk = await storage.getOneTimePreKey(i);
    if (otpk) oneTimePreKeys.push({ keyId: i, data: serializeOneTimePreKey(otpk) });
    if (oneTimePreKeys.length >= otpkCount) break;
  }

  const sessions: Array<{ address: string; state: string }> = [];
  for (const address of knownAddresses) {
    const state = await storage.getSession(address);
    if (state) sessions.push({ address, state: serializeSessionState(state) });
  }

  // Trusted identities: can't enumerate; caller must provide addresses
  const trustedIdentities: Array<{ address: string; key: string }> = [];

  const payload: BackupPayload = {
    version: BACKUP_VERSION,
    identity: identity ? serializeIdentityKeyPair(identity) : null,
    registrationId,
    signedPreKeys,
    oneTimePreKeys,
    sessions,
    trustedIdentities,
  };

  const plaintext = new TextEncoder().encode(JSON.stringify(payload));

  // Derive encryption key from passphrase via HKDF
  const salt = crypto.randomBytes(SALT_SIZE);
  const key = await crypto.hkdf(
    new TextEncoder().encode(passphrase),
    salt,
    BACKUP_INFO,
    32,
  );

  const { ciphertext, nonce } = await crypto.aesGcmEncrypt(key, plaintext);

  // Zero the derived key after use
  crypto.zeroize(key);

  return {
    version: BACKUP_VERSION,
    salt: toBase64(salt),
    nonce: toBase64(nonce),
    ciphertext: toBase64(ciphertext),
  };
}

/**
 * Import a backup blob, decrypt with the passphrase, and write all state
 * to the given storage (overwriting existing state).
 */
export async function importBackup(
  crypto: CryptoProvider,
  storage: StorageProvider,
  blob: BackupBlob,
  passphrase: string,
): Promise<void> {
  if (blob.version !== BACKUP_VERSION) {
    throw new Error(`Unsupported backup version: ${blob.version}`);
  }

  const salt = fromBase64(blob.salt);
  const nonce = fromBase64(blob.nonce);
  const ciphertext = fromBase64(blob.ciphertext);

  // Derive the same key
  const key = await crypto.hkdf(
    new TextEncoder().encode(passphrase),
    salt,
    BACKUP_INFO,
    32,
  );

  let plaintext: Uint8Array;
  try {
    plaintext = await crypto.aesGcmDecrypt(key, ciphertext, nonce);
  } catch {
    crypto.zeroize(key);
    throw new Error('Wrong passphrase or corrupted backup');
  }

  crypto.zeroize(key);

  const payload = JSON.parse(new TextDecoder().decode(plaintext)) as BackupPayload;

  // Restore identity
  if (payload.identity) {
    await storage.saveIdentityKeyPair(deserializeIdentityKeyPair(payload.identity));
  }
  await storage.saveLocalRegistrationId(payload.registrationId);

  // Restore signed prekeys
  for (const spk of payload.signedPreKeys) {
    await storage.saveSignedPreKey(deserializeSignedPreKey(spk.data));
  }

  // Restore one-time prekeys
  for (const otpk of payload.oneTimePreKeys) {
    await storage.saveOneTimePreKey(deserializeOneTimePreKey(otpk.data));
  }

  // Restore sessions
  for (const s of payload.sessions) {
    await storage.saveSession(s.address, deserializeSessionState(s.state));
  }

  // Restore trust
  for (const t of payload.trustedIdentities) {
    await storage.saveTrustedIdentity(t.address, fromBase64(t.key));
  }
}

/** Serialize a backup blob to a compact single-string form (for copy/paste or QR). */
export function backupToString(blob: BackupBlob): string {
  return `shade-backup:v${blob.version}:${blob.salt}:${blob.nonce}:${blob.ciphertext}`;
}

/** Parse a backup string back into a BackupBlob. */
export function backupFromString(str: string): BackupBlob {
  const parts = str.split(':');
  if (parts.length !== 5 || parts[0] !== 'shade-backup') {
    throw new Error('Invalid backup string format');
  }
  const version = parseInt(parts[1]!.replace('v', ''), 10);
  return {
    version,
    salt: parts[2]!,
    nonce: parts[3]!,
    ciphertext: parts[4]!,
  };
}