import 'dart:convert';
import 'dart:typed_data';
import 'package:pointycastle/export.dart';

/// Derives and stores the 256-bit session key from a user-supplied passphrase.
///
/// Key derivation: PBKDF2-HMAC-SHA256 with a random salt.
/// The salt is kept in memory for the session; both peers must use the SAME
/// passphrase and DIFFERENT salts are handled by transmitting the salt
/// in the handshake packet before the session starts.
class KeyManager {
  static const int _keyLen  = 32; // AES-256
  static const int _saltLen = 16;
  static const int _iters   = 100000;

  Uint8List? _key;
  Uint8List? _salt;

  /// True once [deriveKey] has been called successfully.
  bool get hasKey => _key != null;

  /// The derived 32-byte key (throws if not yet derived).
  Uint8List get key {
    final k = _key;
    if (k == null) throw StateError('Key not derived yet');
    return k;
  }

  /// The 16-byte salt used during derivation (for sharing with the peer).
  Uint8List get salt {
    final s = _salt;
    if (s == null) throw StateError('Salt not available');
    return s;
  }

  // ── Derivation ─────────────────────────────────────────────────────

  /// Derive a 256-bit key from [passphrase] using PBKDF2-HMAC-SHA256.
  ///
  /// If [salt] is provided (received from the peer during handshake),
  /// it is used directly. Otherwise a fresh random salt is generated.
  ///
  /// Throws [ArgumentError] if the passphrase is empty.
  void deriveKey(String passphrase, {Uint8List? salt}) {
    if (passphrase.isEmpty) {
      throw ArgumentError('Passphrase must not be empty');
    }

    // Use provided salt or generate one.
    final usedSalt = salt ?? _generateSalt();
    _salt = usedSalt;

    final passwordBytes = Uint8List.fromList(utf8.encode(passphrase));

    // PBKDF2-HMAC-SHA256 via PointyCastle.
    final pbkdf2 = PBKDF2KeyDerivator(HMac(SHA256Digest(), 64))
      ..init(Pbkdf2Parameters(usedSalt, _iters, _keyLen));

    _key = pbkdf2.process(passwordBytes);
  }

  /// Load a raw 32-byte key directly (for testing or pre-shared key import).
  void loadRawKey(Uint8List rawKey) {
    if (rawKey.length != _keyLen) {
      throw ArgumentError('Raw key must be $_keyLen bytes');
    }
    _key = Uint8List.fromList(rawKey);
    _salt = Uint8List(_saltLen); // zero salt when using raw key
  }

  /// Clear the key from memory (call on session teardown).
  void clear() {
    if (_key != null) {
      _key!.fillRange(0, _key!.length, 0);
      _key = null;
    }
    _salt = null;
  }

  // ── Nonce construction ─────────────────────────────────────────────

  /// Build a 16-byte AES-CTR nonce for a given packet sequence number.
  ///
  /// Layout: [salt[0..7] ∥ sessionId[0..3] ∥ seqHigh ∥ seqLow ∥ 00 00 00 00]
  /// (128-bit nonce; the counter part at the end starts at 0 per-block.)
  ///
  /// Using the salt's first 8 bytes + session ID + seq ensures that
  /// each packet has a unique nonce without transmitting the full nonce.
  Uint8List buildNonce(Uint8List sessionId, int seq) {
    final nonce = Uint8List(16);
    // First 8 bytes from salt (or zeros if salt unavailable).
    final s = _salt;
    if (s != null) {
      for (int i = 0; i < 8 && i < s.length; i++) { nonce[i] = s[i]; }
    }
    // Next 4 bytes: session ID.
    for (int i = 0; i < 4 && i < sessionId.length; i++) {
      nonce[8 + i] = sessionId[i];
    }
    // Last 4 bytes: packet sequence number (big-endian).
    nonce[12] = (seq >> 24) & 0xFF;
    nonce[13] = (seq >> 16) & 0xFF;
    nonce[14] = (seq >>  8) & 0xFF;
    nonce[15] =  seq        & 0xFF;
    return nonce;
  }

  // ── Private helpers ────────────────────────────────────────────────

  static Uint8List _generateSalt() {
    final rng = FortunaRandom();
    // Seed Fortuna with timestamp + hash-spread bytes.
    final seed = Uint8List(32);
    final ts = DateTime.now().microsecondsSinceEpoch;
    for (int i = 0; i < 8; i++) {
      seed[i] = (ts >> (i * 8)) & 0xFF;
    }
    rng.seed(KeyParameter(seed));
    return rng.nextBytes(_saltLen);
  }
}