/**
 * Noise NK State Machine for Peach Sync
 *
 * Pattern: Noise NK
 * DH: X25519
 * Cipher: ChaCha20-Poly1305
 * Hash: BLAKE2s
 *
 * NK Pattern:
 *   <- s                    (pre-message: responder static public key)
 *   ...
 *   -> e, es                (initiator sends ephemeral key, performs DH with responder static)
 *   <- e, ee                (responder sends ephemeral key, performs DH)
 */

import { x25519 } from "@noble/curves/ed25519.js";
import { chacha20poly1305 } from "@noble/ciphers/chacha.js";
import { blake2s } from "@noble/hashes/blake2.js";
import { randomBytes } from "@noble/hashes/utils.js";

export interface KeyPair {
  publicKey: Uint8Array;
  privateKey: Uint8Array;
}

export interface NoiseResult {
  output: Uint8Array | null;
  canSend: boolean;
  isComplete: boolean;
}

export class NoiseNKStateMachine {
  private staticKeyPair: KeyPair;
  private ephemeralKeyPair: KeyPair | null = null;
  private remoteEphemeralKey: Uint8Array | null = null;
  private readonly isInitiator: boolean;
  private remoteStaticKey: Uint8Array | null = null;
  private sendNonce = 0;
  private recvNonce = 0;
  private sessionKey: Uint8Array | null = null;
  private handshakeComplete = false;
  private canSendMessages = false;

  constructor(staticKeyPair: KeyPair, isInitiator: boolean) {
    this.staticKeyPair = staticKeyPair;
    this.isInitiator = isInitiator;
  }

  setRemoteStaticKey(key: Uint8Array): void {
    if (this.isInitiator) {
      this.remoteStaticKey = key;
    }
  }

  private deriveSessionKey(dhOutputs: Uint8Array[]): Uint8Array {
    const totalLength = dhOutputs.reduce((sum, arr) => sum + arr.length, 0);
    const concat = new Uint8Array(totalLength);
    let offset = 0;
    for (const output of dhOutputs) {
      concat.set(output, offset);
      offset += output.length;
    }
    return blake2s(concat, { dkLen: 32 });
  }

  private performDH(localPrivate: Uint8Array, remotePublic: Uint8Array): Uint8Array {
    return x25519.getSharedSecret(localPrivate, remotePublic);
  }

  private generateEphemeralKeyPair(): KeyPair {
    const privateKey = randomBytes(32);
    const publicKey = x25519.getPublicKey(privateKey);
    return { publicKey, privateKey };
  }

  private encryptWithKey(key: Uint8Array, plaintext: Uint8Array, nonce: bigint): Uint8Array {
    const nonceBytes = new Uint8Array(12);
    const view = new DataView(nonceBytes.buffer);
    view.setBigUint64(4, nonce, true);
    const cipher = chacha20poly1305(key, nonceBytes);
    return cipher.encrypt(plaintext);
  }

  private decryptWithKey(key: Uint8Array, ciphertext: Uint8Array, nonce: bigint): Uint8Array {
    const nonceBytes = new Uint8Array(12);
    const view = new DataView(nonceBytes.buffer);
    view.setBigUint64(4, nonce, true);
    const cipher = chacha20poly1305(key, nonceBytes);
    return cipher.decrypt(ciphertext);
  }

  readMessage(data: Uint8Array): NoiseResult {
    if (this.handshakeComplete) {
      return { output: null, canSend: this.canSendMessages, isComplete: true };
    }

    if (this.isInitiator) {
      if (!this.remoteStaticKey) {
        throw new Error("Remote static key not set");
      }

      if (this.ephemeralKeyPair === null) {
        this.ephemeralKeyPair = this.generateEphemeralKeyPair();
        const e = this.ephemeralKeyPair.publicKey;
        const es = this.performDH(this.ephemeralKeyPair.privateKey, this.remoteStaticKey);
        const key = this.deriveSessionKey([es]);
        const output = this.encryptWithKey(key, data, BigInt(this.sendNonce));
        this.sendNonce++;
        return { output, canSend: true, isComplete: false };
      } else {
        this.remoteEphemeralKey = data;
        const ee = this.performDH(this.ephemeralKeyPair.privateKey, this.remoteEphemeralKey);
        this.sessionKey = this.deriveSessionKey([ee]);
        this.handshakeComplete = true;
        this.canSendMessages = true;
        return { output: null, canSend: true, isComplete: true };
      }
    } else {
      if (this.remoteEphemeralKey === null) {
        this.remoteEphemeralKey = data;
        const e = this.staticKeyPair.publicKey;
        const ee = this.performDH(this.staticKeyPair.privateKey, this.remoteEphemeralKey);
        this.sessionKey = this.deriveSessionKey([ee]);
        return { output: e, canSend: true, isComplete: false };
      } else {
        const es = this.performDH(this.staticKeyPair.privateKey, this.remoteEphemeralKey);
        const key = this.deriveSessionKey([es]);
        const output = this.decryptWithKey(key, data, BigInt(this.recvNonce));
        this.recvNonce++;
        this.sessionKey = this.deriveSessionKey([es, this.remoteEphemeralKey]);
        this.handshakeComplete = true;
        this.canSendMessages = true;
        return { output, canSend: true, isComplete: true };
      }
    }
  }

  writeMessage(payload: Uint8Array): Uint8Array {
    if (this.handshakeComplete) {
      return this.encrypt(payload);
    }

    if (this.isInitiator) {
      if (this.ephemeralKeyPair === null) {
        this.ephemeralKeyPair = this.generateEphemeralKeyPair();
        return this.ephemeralKeyPair.publicKey;
      } else if (this.remoteEphemeralKey === null) {
        if (!this.remoteStaticKey) {
          throw new Error("Remote static key not set");
        }
        const es = this.performDH(this.ephemeralKeyPair.privateKey, this.remoteStaticKey);
        const key = this.deriveSessionKey([es]);
        return this.encryptWithKey(key, payload, BigInt(this.sendNonce));
      }
      throw new Error("Invalid handshake state");
    } else {
      if (this.remoteEphemeralKey === null) {
        this.remoteEphemeralKey = payload;
        return this.staticKeyPair.publicKey;
      } else {
        const es = this.performDH(this.staticKeyPair.privateKey, this.remoteEphemeralKey);
        const key = this.deriveSessionKey([es]);
        this.sessionKey = key;
        this.handshakeComplete = true;
        this.canSendMessages = true;
        return this.encryptWithKey(key, payload, BigInt(this.sendNonce));
      }
    }
  }

  getSessionKey(): Uint8Array | null {
    return this.sessionKey;
  }

  encrypt(plaintext: Uint8Array): Uint8Array {
    if (!this.sessionKey) {
      throw new Error("Session key not available");
    }
    const ciphertext = this.encryptWithKey(this.sessionKey, plaintext, BigInt(this.sendNonce));
    this.sendNonce++;
    return ciphertext;
  }

  decrypt(ciphertext: Uint8Array): Uint8Array {
    if (!this.sessionKey) {
      throw new Error("Session key not available");
    }
    const plaintext = this.decryptWithKey(this.sessionKey, ciphertext, BigInt(this.recvNonce));
    this.recvNonce++;
    return plaintext;
  }
}