import { Buffer } from "buffer";
import { ed25519 as library } from "@noble/curves/ed25519.js";
import { sha512 } from "@noble/hashes/sha2.js";
import { mnemonicToEntropy, wordlists } from "bip39";
import { blake } from "../blake";
import { bech32 } from "../bech32";
import { AlgorithmId, CoseKey, KeyType } from "../../cip8/CoseKey";
import { CardanoCliJsonKey } from "../../types";
import { hmacSha512, pbkdf2Sha512 } from "./node"; // Default to node, build script will swap for browser
import { bytesFromString } from "../../util";

export const KEY_HASH_LENGTH = 28;

export abstract class CardanoKey {
  abstract sign(message: Buffer): Buffer;
  abstract verify(message: Buffer, signature: Buffer): boolean;
  abstract toBech32(): string;
  abstract toString(): string;
  abstract toBytes(): Buffer;
  abstract publicBytes(): Buffer;
  abstract publicKeyHash(): Buffer;
  abstract toCoseKey(): CoseKey;
  abstract toCardanoCliJson(): CardanoCliJsonKey;
  abstract withoutSignKey(): CardanoKey;

  public static fromBytes(bytes: Buffer): CardanoKey {
    if (bytes.length === 96 || bytes.length === 64) {
      return Ed25519HdKey.fromBytes(bytes);
    }
    return Ed25519Key.fromBytes(bytes);
  }

  public static fromBech32(bech32Str: string): CardanoKey {
    const decoded = bech32.decode(bech32Str, 1000);
    const bytes = Buffer.from(decoded.data);

    if (decoded.prefix === "xprv") {
      return Ed25519HdKey.fromBytes(bytes);
    }

    if (decoded.prefix === "xpub") {
      return Ed25519HdKey.fromXpubBytes(bytes);
    }

    return Ed25519Key.fromBytes(bytes);
  }

  public static fromHex(hex: string): CardanoKey {
    return this.fromBytes(Buffer.from(hex, "hex"));
  }

  public static fromString(str: string): CardanoKey {
    const decoded = bytesFromString(str, {
      bech32Limit: 1000,
      invalidMessage: "Invalid key string: expected bech32 or hex-encoded key",
    });

    if (decoded.prefix === "xprv") {
      return Ed25519HdKey.fromBytes(decoded.bytes);
    }

    if (decoded.prefix === "xpub") {
      return Ed25519HdKey.fromXpubBytes(decoded.bytes);
    }

    if (decoded.prefix === "vk_") {
      return Ed25519Key.fromPublicKey(decoded.bytes);
    }

    return this.fromBytes(decoded.bytes);
  }

  public static fromCardanoCliJson(json: CardanoCliJsonKey): CardanoKey {
    const keyType = json?.type ?? "";
    const cborHex = json?.cborHex ?? "";

    if (keyType.includes("Extended") || cborHex.startsWith("5880")) {
      return Ed25519HdKey.fromCardanoCliJson(json);
    }

    return Ed25519Key.fromCardanoCliJson(json);
  }
}

function bytesToBigIntLE(bytes: Uint8Array): bigint {
  let n = 0n;
  for (let i = bytes.length - 1; i >= 0; i--) {
    n = (n << 8n) + BigInt(bytes[i]);
  }
  return n;
}

function bigIntToBytesLE(num: bigint, len: number): Buffer {
  const out = Buffer.alloc(len);
  let n = num;
  for (let i = 0; i < len; i++) {
    out[i] = Number(n & 0xffn);
    n >>= 8n;
  }
  return out;
}

function pubFromKl(kl: Buffer): Buffer {
  const scalar = bytesToBigIntLE(kl) % library.Point.Fn.ORDER;
  return Buffer.from(library.Point.BASE.multiply(scalar).toBytes());
}

const hardenIndex = (num: number) => 0x80000000 + num;

// Stores ed25519 KeyPair and hash of publicKey
export class Ed25519Key extends CardanoKey {
  public signingKey?: Buffer;
  public verificationKey: Buffer;
  public keyHash: Buffer;

  public constructor(priv: Buffer | undefined, pub: Buffer, pkh: Buffer) {
    super();
    this.verificationKey = pub;
    this.keyHash = pkh;
    if (priv) {
      if (priv.length !== 32 && priv.length !== 64) {
        throw new Error("Invalid private key length: Ed25519Key supports 32-byte seeds or 64-byte extended keys.");
      }
      this.signingKey = priv;
    }
  }

  get private(): Buffer | undefined {
    return this.signingKey;
  }

  get public(): Buffer {
    return this.verificationKey;
  }

  get pkh(): Buffer {
    return this.keyHash;
  }

  public hasPrivateKey(): boolean {
    return !!this.signingKey;
  }

  public withoutSignKey(): Ed25519Key {
    return new Ed25519Key(undefined, this.verificationKey, this.keyHash);
  }

  public static override fromCardanoCliJson(cardanoCliKey: CardanoCliJsonKey): Ed25519Key {
    if (!cardanoCliKey.cborHex || typeof cardanoCliKey.cborHex !== "string") {
      throw new Error("Ed25519Key.fromCardanocliJson: keyJson.cborHex to be present got " + cardanoCliKey.cborHex);
    }

    const keyType = cardanoCliKey.type || "";
    const isVerificationKey = keyType.includes("VerificationKey");
    const isExtendedKey = keyType.includes("Extended");

    if (isExtendedKey) {
      throw new Error("Ed25519Key.fromCardanoCliJson: extended key types are not supported here. Use CardanoKey.fromCardanoCliJson or Ed25519HdKey.fromCardanoCliJson.");
    }

    let keyCbor: string;
    if (cardanoCliKey.cborHex.startsWith("5820") && cardanoCliKey.cborHex.length === 68) {
      keyCbor = cardanoCliKey.cborHex.slice(4);
      const keyBytes = Buffer.from(keyCbor, "hex");
      if (isVerificationKey) {
        return Ed25519Key.fromPublicKey(keyBytes);
      }
      return Ed25519Key.fromPrivateKey(keyBytes);
    } else {
      throw new Error("Ed25519Key.fromCardanocliJson: Expected non-extended CBOR bytes prefix '5820' for 32 bytes");
    }
  }

  public toCardanoCliJson(): CardanoCliJsonKey {
    if (this.signingKey) {
      if (this.signingKey.length !== 32) {
        throw new Error("Ed25519Key.toCardanoCliJson: unsupported private key length for non-extended key. Use Ed25519HdKey for extended keys.");
      }

      return {
        type: "PaymentSigningKeyShelley_ed25519",
        description: "Payment Signing Key",
        cborHex: "5820" + this.signingKey.toString("hex"),
      };
    }

    return {
      type: "PaymentVerificationKeyShelley_ed25519",
      description: "Payment Verification Key",
      cborHex: "5820" + this.verificationKey.toString("hex"),
    };
  }

  public static generate() {
    const privKey = Buffer.from(library.utils.randomSecretKey());
    return Ed25519Key.fromPrivateKey(privKey);
  }

  public static override fromBytes(bytes: Buffer): Ed25519Key {
    if (bytes.length === 32 || bytes.length === 64) {
      return Ed25519Key.fromPrivateKey(bytes);
    }
    return Ed25519Key.fromPublicKey(bytes);
  }

  public static override fromBech32(bech32Str: string): Ed25519Key {
    const { data } = bech32.decode(bech32Str);
    return Ed25519Key.fromBytes(Buffer.from(data));
  }

  public static override fromHex(hex: string): Ed25519Key {
    return Ed25519Key.fromBytes(Buffer.from(hex, "hex"));
  }

  public static override fromString(str: string): Ed25519Key {
    const decoded = bytesFromString(str, {
      invalidMessage: "Invalid key string: expected bech32 or hex-encoded key",
    });
    return Ed25519Key.fromBytes(decoded.bytes);
  }

  public toBytes(): Buffer {
    return this.signingKey || this.verificationKey;
  }

  public publicBytes(): Buffer {
    return this.verificationKey;
  }

  public publicKeyHash(): Buffer {
    return this.keyHash;
  }

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

  public toBech32(): string {
    if (this.signingKey) {
      return this.bech32PrivateKey();
    }
    return this.bech32PublicKey();
  }

  public sign(message: Buffer): Buffer {
    if (!this.signingKey) {
      throw new Error("No private key available for signing");
    }
    return Buffer.from(library.sign(message, this.signingKey));
  }

  public static fromPrivateKey(privKey: Buffer) {
    let pubKey: Uint8Array;
    if (privKey.length === 32) {
      pubKey = library.getPublicKey(privKey);
    } else {
      throw new Error("Invalid private key length: Ed25519Key supports only 32-byte keys.");
    }
    const pkh = blake.hash28(Buffer.from(pubKey));
    return new Ed25519Key(privKey, Buffer.from(pubKey), Buffer.from(pkh));
  }

  public static fromRaw(privateKey: Buffer, publicKey: Buffer): Ed25519Key {
    const pkh = blake.hash28(publicKey);
    return new Ed25519Key(privateKey, publicKey, Buffer.from(pkh));
  }

  public static fromPublicKey(pubKey: Buffer) {
    const pkh = blake.hash28(pubKey);
    return new Ed25519Key(undefined, pubKey, Buffer.from(pkh));
  }

  public static fromPrivateKeyHex(privKey: string) {
    return Ed25519Key.fromPrivateKey(Buffer.from(privKey, "hex"));
  }

  public bech32Pkh(prefix: string = "stake"): string {
    return bech32.encode(prefix, this.pkh);
  }

  public bech32PublicKey(prefix: string = "vk_"): string {
    return bech32.encode(prefix, this.public);
  }

  public bech32PrivateKey(prefix: string = "sk_"): string {
    return bech32.encode(prefix, this.private!);
  }

  public verify(message: Buffer, signature: Buffer): boolean {
    return library.verify(signature, message, this.verificationKey);
  }

  public toCoseKey(): CoseKey {
    const key = new CoseKey(KeyType.OKP, AlgorithmId.EDSA);
    key.addHeader(-1, 6);
    key.addHeader(-2, this.public);
    return key;
  }

  public static fromCoseKey(key: CoseKey): Ed25519Key {
    if (key.algorithmId != AlgorithmId.EDSA) {
      throw new Error("Invalid algorithm Id: expected \"" + AlgorithmId.EDSA + "\" got: \"" + key.algorithmId + "\"");
    }
    return this.fromPublicKey(key.pubKeyBytes);
  }

  public toJSON() {
    return {
      private: this.private ? Buffer.from(this.private).toString("hex") : undefined,
      public: Buffer.from(this.public).toString("hex"),
      pkh: Buffer.from(this.pkh).toString("hex"),
    };
  }

  public json() {
    return this.toJSON();
  }

  public static fromJson(json: any): Ed25519Key {
    if (!json || typeof json !== "object") {
      throw new Error("Invalid JSON format for Ed25519Key: Input must be a non-null object.");
    }
    if (!json.public || !json.pkh) {
      throw new Error("Invalid JSON format for Ed25519Key: Missing required fields (public or pkh).");
    }

    const pub = Buffer.from(json.public, "hex");
    const pkh = Buffer.from(json.pkh, "hex");
    const priv = json.private ? Buffer.from(json.private, "hex") : undefined;
    return new Ed25519Key(priv, pub, pkh);
  }
}

export class Ed25519HdKey extends CardanoKey {
  private kl?: Buffer;
  private kr?: Buffer;
  private pub: Buffer;
  private chainCode: Buffer;

  public constructor(chainCode: Buffer, pub: Buffer, kl?: Buffer, kr?: Buffer) {
    super();
    if (pub.length !== 32) throw new Error("Invalid pub length: must be 32 bytes");
    if (chainCode.length !== 32) throw new Error("Invalid chainCode length: must be 32 bytes");
    if ((kl && !kr) || (!kl && kr)) throw new Error("Both kl and kr must be provided for a private Ed25519HdKey");
    if (kl && kl.length !== 32) throw new Error("Invalid kl length: must be 32 bytes");
    if (kr && kr.length !== 32) throw new Error("Invalid kr length: must be 32 bytes");
    this.kl = kl;
    this.kr = kr;
    this.pub = pub;
    this.chainCode = chainCode;
  }

  public hasPrivateKey(): boolean {
    return !!(this.kl && this.kr);
  }

  public withoutSignKey(): Ed25519HdKey {
    return new Ed25519HdKey(Buffer.from(this.chainCode), Buffer.from(this.pub));
  }

  public static hardenIndex(num: number) {
    return hardenIndex(num);
  }

  public child(index: number | number[], harden = false): Ed25519HdKey {
    const indexes = typeof index === "number" ? [index] : index;
    if (!Array.isArray(indexes)) throw new Error("Index must be a number or an array of numbers");
    for (const i of indexes) {
      if (!Number.isInteger(i) || i < 0 || i >= 2 ** 32) {
        throw new Error(`Invalid index: ${i} must be an integer between 0 and 2^32 - 1`);
      }
    }
    let current: Ed25519HdKey = this;
    for (const i of indexes) {
      current = current.deriveChild(harden ? hardenIndex(i) : i);
    }
    return current;
  }

  private deriveChild(i: number): Ed25519HdKey {
    const isHardened = i >= 0x80000000;
    const serI = Buffer.alloc(4);
    serI.writeUInt32LE(i, 0);

    let zPrefix: Uint8Array;
    let cPrefix: Uint8Array;
    let serP: Buffer;
    if (isHardened) {
      if (!this.kl || !this.kr) throw new Error("Cannot derive hardened child from xpub key");
      zPrefix = new Uint8Array([0x00]);
      cPrefix = new Uint8Array([0x01]);
      serP = Buffer.concat([this.kl, this.kr]);
    } else {
      zPrefix = new Uint8Array([0x02]);
      cPrefix = new Uint8Array([0x03]);
      serP = this.pub;
    }

    const zData = Buffer.concat([Buffer.from(zPrefix), serP, serI]);
    const z = hmacSha512(this.chainCode, zData);

    const cData = Buffer.concat([Buffer.from(cPrefix), serP, serI]);
    const iFull = hmacSha512(this.chainCode, cData);
    const cChild = iFull.subarray(32, 64);

    const zL = z.subarray(0, 28);
    const zR = z.subarray(32, 64);

    const zlBig = bytesToBigIntLE(zL) * 8n;
    const parentPoint = library.Point.fromBytes(this.pub);
    const deltaPoint = library.Point.BASE.multiply(zlBig % library.Point.Fn.ORDER);
    const pubChild = Buffer.from(parentPoint.add(deltaPoint).toBytes());

    if (!this.kl || !this.kr) {
      return new Ed25519HdKey(cChild, pubChild);
    }

    const klBig = bytesToBigIntLE(this.kl);
    const klChild = bigIntToBytesLE((klBig + zlBig) % (1n << 256n), 32);
    const zrBig = bytesToBigIntLE(zR);
    const krBig = bytesToBigIntLE(this.kr);
    const krChild = bigIntToBytesLE((krBig + zrBig) % (1n << 256n), 32);
    return new Ed25519HdKey(cChild, pubChild, klChild, krChild);
  }

  public static fromMnemonicString(mnemonic: string): Ed25519HdKey {
    let entropyHex: string | undefined;
    try {
      entropyHex = mnemonicToEntropy(mnemonic);
    } catch (_) {
      for (const wl of Object.values(wordlists)) {
        try {
          entropyHex = mnemonicToEntropy(mnemonic, wl as string[]);
          break;
        } catch (_) {}
      }
    }
    if (!entropyHex) {
      throw new Error("Invalid mnemonic");
    }
    const entropyBuffer = Buffer.from(entropyHex, "hex");
    return this.fromBip39Entropy(entropyBuffer, "");
  }

  public static fromEntropy(entropy: Buffer, password: string = ""): Ed25519HdKey {
    return this.fromBip39Entropy(entropy, password);
  }

  private static fromBip39Entropy(entropy: Buffer, password: string): Ed25519HdKey {
    const passBytes = Buffer.from(password, "utf8");
    const root = pbkdf2Sha512(passBytes, entropy, 4096, 96);

    const kl = Buffer.from(root.subarray(0, 32));
    const kr = root.subarray(32, 64);
    const chainCode = root.subarray(64, 96);

    kl[0] &= 0xf8;
    kl[31] &= 0x5f;
    kl[31] |= 0x40;

    return new Ed25519HdKey(chainCode, pubFromKl(kl), kl, kr);
  }

  public ed25519KeyRaw() {
    if (!this.kl || !this.kr) throw new Error("No private key material available");
    return Buffer.concat([this.kl, this.kr]);
  }

  public ed25519Key(): Ed25519Key {
    if (this.kl && this.kr) {
      return Ed25519Key.fromRaw(Buffer.concat([this.kl, this.kr]), this.pub);
    }
    return Ed25519Key.fromPublicKey(this.pub);
  }

  public sign(message: Buffer): Buffer {
    if (!this.kl || !this.kr) {
      throw new Error("No private key available for signing");
    }
    const prefixHash = sha512(Buffer.concat([this.kr, message]));
    const nonce = bytesToBigIntLE(prefixHash) % library.Point.Fn.ORDER;
    const R = library.Point.BASE.multiply(nonce);
    const rBytes = Buffer.from(R.toBytes());
    const pubBytes = this.pub;
    const hHash = sha512(Buffer.concat([rBytes, pubBytes, message]));
    const h = bytesToBigIntLE(hHash) % library.Point.Fn.ORDER;
    const scalar = bytesToBigIntLE(this.kl) % library.Point.Fn.ORDER;
    const S = (nonce + h * scalar) % library.Point.Fn.ORDER;
    const sBytes = bigIntToBytesLE(S, 32);
    return Buffer.concat([rBytes, sBytes]);
  }

  public verify(message: Buffer, signature: Buffer): boolean {
    if (signature.length !== 64) return false;
    return library.verify(signature, message, this.pub);
  }

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

  public publicBytes(): Buffer {
    return this.pub;
  }

  public publicKeyHash(): Buffer {
    return Buffer.from(blake.hash28(this.pub));
  }

  public toCoseKey(): CoseKey {
    const key = new CoseKey(KeyType.OKP, AlgorithmId.EDSA);
    key.addHeader(-1, 6);
    key.addHeader(-2, this.pub);
    return key;
  }

  public toCardanoCliJson(): CardanoCliJsonKey {
    if (this.kl && this.kr) {
      const privateKey = Buffer.concat([this.kl, this.kr]);
      const extendedSigningBytes = Buffer.concat([privateKey, this.pub, this.chainCode]);
      return {
        type: "PaymentExtendedSigningKeyShelley_ed25519_bip32",
        description: "Payment Signing Key",
        cborHex: "5880" + extendedSigningBytes.toString("hex"),
      };
    }

    return {
      type: "PaymentExtendedVerificationKeyShelley_ed25519_bip32",
      description: "Payment Verification Key",
      cborHex: "5840" + this.xpubBytes().toString("hex"),
    };
  }

  public static override fromCardanoCliJson(cardanoCliKey: CardanoCliJsonKey): Ed25519HdKey {
    if (!cardanoCliKey.cborHex || typeof cardanoCliKey.cborHex !== "string") {
      throw new Error("Ed25519HdKey.fromCardanoCliJson: keyJson.cborHex to be present got " + cardanoCliKey.cborHex);
    }

    const keyType = cardanoCliKey.type || "";

    if (keyType.includes("ExtendedSigningKey") || cardanoCliKey.cborHex.startsWith("5880")) {
      if (!cardanoCliKey.cborHex.startsWith("5880") || cardanoCliKey.cborHex.length !== 260) {
        throw new Error("Ed25519HdKey.fromCardanoCliJson: Expected CBOR bytes prefix '5880' for 128-byte extended signing key");
      }

      const raw = Buffer.from(cardanoCliKey.cborHex.slice(4), "hex");
      const kl = Buffer.from(raw.subarray(0, 32));
      const kr = Buffer.from(raw.subarray(32, 64));
      const pub = Buffer.from(raw.subarray(64, 96));
      const chainCode = Buffer.from(raw.subarray(96, 128));
      return new Ed25519HdKey(chainCode, pub, kl, kr);
    }

    if (keyType.includes("ExtendedVerificationKey") || cardanoCliKey.cborHex.startsWith("5840")) {
      if (!cardanoCliKey.cborHex.startsWith("5840") || cardanoCliKey.cborHex.length !== 132) {
        throw new Error("Ed25519HdKey.fromCardanoCliJson: Expected CBOR bytes prefix '5840' for 64-byte extended verification key");
      }

      const xpub = Buffer.from(cardanoCliKey.cborHex.slice(4), "hex");
      return Ed25519HdKey.fromXpubBytes(xpub);
    }

    throw new Error("Ed25519HdKey.fromCardanoCliJson: unsupported key type for HD key: " + keyType);
  }

  public static override fromBytes(bytes: Buffer): Ed25519HdKey {
    if (bytes.length === 64) return this.fromXpubBytes(bytes);
    if (bytes.length !== 96) throw new Error("Bytes must be 96 (xprv) or 64 (xpub) length");

    const kl = Buffer.from(bytes.subarray(0, 32));
    const kr = bytes.subarray(32, 64);
    const chainCode = bytes.subarray(64, 96);

    kl[0] &= 0xf8;
    kl[31] &= 0x5f;
    kl[31] |= 0x40;

    return new Ed25519HdKey(chainCode, pubFromKl(kl), kl, kr);
  }

  public static override fromBech32(bech32Str: string): Ed25519HdKey {
    const decoded = bech32.decode(bech32Str, 300);
    if (decoded.prefix === "xprv") return this.fromBytes(decoded.data);
    if (decoded.prefix === "xpub") return this.fromXpubBytes(decoded.data);
    throw new Error(`Invalid key prefix: expected "xprv" or "xpub", got "${decoded.prefix}"`);
  }

  public static override fromHex(hex: string): Ed25519HdKey {
    return this.fromBytes(Buffer.from(hex, "hex"));
  }

  public static override fromString(str: string): Ed25519HdKey {
    const decoded = bytesFromString(str, {
      bech32Limit: 300,
      invalidMessage: "Invalid key string: expected bech32 xprv/xpub or hex-encoded key",
    });

    if (decoded.prefix) {
      if (decoded.prefix === "xprv") return this.fromBytes(decoded.bytes);
      if (decoded.prefix === "xpub") return this.fromXpubBytes(decoded.bytes);
      throw new Error(`Invalid key prefix: expected "xprv" or "xpub", got "${decoded.prefix}"`);
    }

    return this.fromBytes(decoded.bytes);
  }

  public static fromXprvString(xprv: string): Ed25519HdKey {
    const decoded = bech32.decode(xprv, 300);
    if (decoded.prefix !== "xprv") throw new Error(`Invalid xprv prefix: expected "xprv", got "${decoded.prefix}"`);
    return this.fromBytes(decoded.data);
  }

  public static fromXprvBytes(bytes: Buffer): Ed25519HdKey {
    return this.fromBytes(bytes);
  }

  public static fromXpubBytes(bytes: Buffer): Ed25519HdKey {
    if (bytes.length !== 64) throw new Error("xpub bytes must be 64 length");
    return new Ed25519HdKey(Buffer.from(bytes.subarray(32, 64)), Buffer.from(bytes.subarray(0, 32)));
  }

  public static fromXpubString(xpub: string): Ed25519HdKey {
    const decoded = bech32.decode(xpub, 300);
    if (decoded.prefix !== "xpub") throw new Error(`Invalid xpub prefix: expected "xpub", got "${decoded.prefix}"`);
    return this.fromXpubBytes(decoded.data);
  }

  public toBytes(): Buffer {
    return this.kl && this.kr ? Buffer.concat([this.kl, this.kr, this.chainCode]) : this.xpubBytes();
  }

  public toString(): string {
    return this.kl && this.kr ? bech32.encode("xprv", this.toBytes()) : this.xpubString();
  }

  public xpubBytes(): Buffer {
    return Buffer.concat([this.pub, this.chainCode]);
  }

  public xpubString(): string {
    return bech32.encode("xpub", this.xpubBytes());
  }

  public bytes(): Buffer {
    return this.toBytes();
  }

  public hex(): string {
    return this.toBytes().toString("hex");
  }
}