import { ed25519, x25519 } from '@noble/curves/ed25519.js';
import { sha256 } from "@noble/hashes/sha2.js";
import { hmac } from '@noble/hashes/hmac.js';
import { ecb } from '@noble/ciphers/aes.js';
import { bytesToHex, equalBytes, hexToBytes } from "./parser";
import { IPublicKey, ISharedSecret, IStaticSecret } from './crypto.types';
import { NodeHash } from './identity.types';

const PUBLIC_KEY_SIZE = 32;
const SEED_SIZE = 32;
const HMAC_SIZE = 2;
const SHARED_SECRET_SIZE = 32;
const SIGNATURE_SIZE = 64;
const STATIC_SECRET_SIZE = 32;

// The "Public" group is a special group that all nodes are implicitly part of.
const publicSecret = hexToBytes("8b3387e9c5cdea6ac9e5edbaa115cd72", 16);

export class PublicKey implements IPublicKey {
  public key: Uint8Array;

  constructor(key: Uint8Array | string) {
    if (typeof key === 'string') {
      this.key = hexToBytes(key, PUBLIC_KEY_SIZE);
    } else if (key instanceof Uint8Array) {
      this.key = key;
    } else {
      throw new Error('Invalid type for PublicKey constructor');
    }
  }

  public toHash(): NodeHash {
    return sha256.create().update(this.key).digest()[0] as NodeHash;
  }

  public toBytes(): Uint8Array {
    return this.key;
  }

  public toString(): string {
    return bytesToHex(this.key);
  }

  public equals(other: PublicKey | Uint8Array | string): boolean {
    let otherKey: Uint8Array;
    if (other instanceof PublicKey) {
      otherKey = other.toBytes();
    } else if (other instanceof Uint8Array) {
      otherKey = other;
    } else if (typeof other === 'string') {
      otherKey = hexToBytes(other, PUBLIC_KEY_SIZE);
    } else {
      throw new Error('Invalid type for PublicKey comparison');
    }
    return equalBytes(this.key, otherKey);
  }

  public verify(message: Uint8Array, signature: Uint8Array): boolean {
    if (signature.length !== SIGNATURE_SIZE) {
      throw new Error(`Invalid signature length: expected ${SIGNATURE_SIZE} bytes, got ${signature.length}`);
    }
    return ed25519.verify(signature, message, this.key);
  }
}

export class PrivateKey {
  private secretKey: Uint8Array;
  private publicKey: PublicKey;

  constructor(seed: Uint8Array | string) {
    if (typeof seed === 'string') {
      seed = hexToBytes(seed, SEED_SIZE);
    }
    if (seed.length !== SEED_SIZE) {
      throw new Error(`Invalid seed length: expected ${SEED_SIZE} bytes, got ${seed.length}`);
    }

    const { secretKey, publicKey } = ed25519.keygen(seed); // Validate seed by generating keys
    this.secretKey = secretKey;
    this.publicKey = new PublicKey(publicKey);
  }

  public toPublicKey(): PublicKey {
    return this.publicKey;
  }

  public toBytes(): Uint8Array {
    return this.secretKey;
  }

  public toString(): string {
    return bytesToHex(this.secretKey);
  }

  public sign(message: Uint8Array): Uint8Array {
    return ed25519.sign(message, this.secretKey);
  }

  public calculateSharedSecret(other: PublicKey | Uint8Array | string): Uint8Array {
    let otherPublicKey: PublicKey;
    if (other instanceof PublicKey) {
      otherPublicKey = other;
    } else if (other instanceof Uint8Array) {
      otherPublicKey = new PublicKey(other);
    } else if (typeof other === 'string') {
      otherPublicKey = new PublicKey(other);
    } else {
      throw new Error('Invalid type for calculateSharedSecret comparison');
    }
    return x25519.getSharedSecret(this.secretKey, otherPublicKey.toBytes());
  }

  static generate(): PrivateKey {
    const { secretKey } = ed25519.keygen(); // Ensure ed25519 is initialized
    return new PrivateKey(secretKey);
  }
}

export class SharedSecret implements ISharedSecret {
  private secret: Uint8Array;

  constructor(secret: Uint8Array) {
    if (secret.length === SHARED_SECRET_SIZE / 2) {
      // Zero pad to the left if the secret is too short (e.g. from x25519)
      const padded = new Uint8Array(SHARED_SECRET_SIZE);
      padded.set(secret, SHARED_SECRET_SIZE - secret.length);
      secret = padded;
    }
    if (secret.length !== SHARED_SECRET_SIZE) {
      throw new Error(`Invalid shared secret length: expected ${SHARED_SECRET_SIZE} bytes, got ${secret.length}`);
    }
    this.secret = secret;
  }

  public toHash(): NodeHash {
    return this.secret[0] as NodeHash;
  }

  public toBytes(): Uint8Array {
    return this.secret;
  }

  public toString(): string {
    return bytesToHex(this.secret);
  }

  public decrypt(hmac: Uint8Array, ciphertext: Uint8Array): Uint8Array {
    if (hmac.length !== HMAC_SIZE) {
      throw new Error(`Invalid HMAC length: expected ${HMAC_SIZE} bytes, got ${hmac.length}`);
    }
    const expectedHmac = this.calculateHmac(ciphertext);
    if (!equalBytes(hmac, expectedHmac)) {
      throw new Error(`Invalid HMAC: decryption failed: expected ${bytesToHex(expectedHmac)}, got ${bytesToHex(hmac)}`);
    }
    const cipher = ecb(this.secret.slice(0, 16), { disablePadding: true });
    const plaintext = new Uint8Array(ciphertext.length);
    for (let i = 0; i < ciphertext.length; i += 16) {
      const block = ciphertext.slice(i, i + 16);
      const dec = cipher.decrypt(block);
      plaintext.set(dec, i);
    }
    // Remove trailing null bytes (0x00) due to padding
    let end = plaintext.length;
    while (end > 0 && plaintext[end - 1] === 0) {
      end--;
    }
    return plaintext.slice(0, end);
  }

  public encrypt(data: Uint8Array): { hmac: Uint8Array, ciphertext: Uint8Array } {
    const key = this.secret.slice(0, 16);
    const cipher = ecb(key, { disablePadding: true });

    const fullBlocks = Math.floor(data.length / 16);
    const remaining = data.length % 16;
    const ciphertext = new Uint8Array((fullBlocks + (remaining > 0 ? 1 : 0)) * 16);
    for (let i = 0; i < fullBlocks; i++) {
      const block = data.slice(i * 16, (i + 1) * 16);
      const enc = cipher.encrypt(block);
      ciphertext.set(enc, i * 16);
    }

    if (remaining > 0) {
      const lastBlock = new Uint8Array(16);
      lastBlock.set(data.slice(fullBlocks * 16));
      const enc = cipher.encrypt(lastBlock);
      ciphertext.set(enc, fullBlocks * 16);
    }

    const hmac = this.calculateHmac(ciphertext);
    return { hmac, ciphertext };
  }

  private calculateHmac(data: Uint8Array): Uint8Array {
    return hmac(sha256, this.secret, data).slice(0, HMAC_SIZE);
  }

  static fromName(name: string): SharedSecret {
    if (name === "Public") {
      return new SharedSecret(publicSecret);
    } else if (!/^#/.test(name)) {
      throw new Error("Only the 'Public' group or groups starting with '#' are supported");
    }
    const hash = sha256.create().update(new TextEncoder().encode(name)).digest();
    return new SharedSecret(hash.slice(0, SHARED_SECRET_SIZE));
  }
}

export class StaticSecret implements IStaticSecret {
  private secret: Uint8Array;

  constructor(secret: Uint8Array | string) {
    if (typeof secret === 'string') {
      secret = hexToBytes(secret, STATIC_SECRET_SIZE);
    }
    if (secret.length !== STATIC_SECRET_SIZE) {
      throw new Error(`Invalid static secret length: expected ${STATIC_SECRET_SIZE} bytes, got ${secret.length}`);
    }
    this.secret = secret;
  }

  public publicKey(): IPublicKey {
    const publicKey = x25519.getPublicKey(this.secret);
    return new PublicKey(publicKey);
  }

  public diffieHellman(otherPublicKey: IPublicKey): SharedSecret {
    const sharedSecret = x25519.getSharedSecret(this.secret, otherPublicKey.toBytes());
    return new SharedSecret(sharedSecret);
  }
}