import { concat, Hex, hexToBigInt, size, sliceHex, toHex } from 'viem';
import { base64ToHex, hexToBase64 } from './base64';

const SIGNATURE_SIZE = 65;
const SECRET_SIZE = 32;
const NONCE_SIZE = 32;
const COORDINATE_SIZE = 32;

export type DistributionSecret =
  | {
      type: 'pre_generated';
      subSig: Hex;
      subSecret: Hex;
    }
  | {
      type: 'secret';
      secret: Hex;
    }
  | {
      type: 'real_time';
      sig: Hex;
      nonce: bigint;
      coordinate: Hex;
    };

export function parseDistributionSecret(secret: string): DistributionSecret {
  const secretHex = base64ToHex(secret);
  const secretSize = size(secretHex);

  if (secretSize === COORDINATE_SIZE + SIGNATURE_SIZE + NONCE_SIZE) {
    return {
      type: 'real_time',
      sig: sliceHex(secretHex, 0, SIGNATURE_SIZE),
      nonce: hexToBigInt(
        sliceHex(secretHex, SIGNATURE_SIZE, SIGNATURE_SIZE + NONCE_SIZE)
      ),
      coordinate: sliceHex(secretHex, SIGNATURE_SIZE + NONCE_SIZE),
    };
  } else if (secretSize === SIGNATURE_SIZE + SECRET_SIZE) {
    const subSig = sliceHex(secretHex, 0, SIGNATURE_SIZE);
    const subSecret = sliceHex(secretHex, SIGNATURE_SIZE);

    return {
      type: 'pre_generated',
      subSig,
      subSecret,
    };
  } else if (secretSize === SECRET_SIZE) {
    return {
      type: 'secret',
      secret: secretHex,
    };
  } else {
    throw new Error('Invalid distribution secret size');
  }
}

export function serializeDistributionSecret(params: DistributionSecret) {
  if (params.type === 'pre_generated') {
    return hexToBase64(concat([params.subSig, params.subSecret]));
  } else if (params.type === 'secret') {
    return hexToBase64(params.secret);
  } else if (params.type === 'real_time') {
    return hexToBase64(
      concat([params.sig, toHex(params.nonce, { size: 32 }), params.coordinate])
    );
  } else {
    throw new Error('Invalid distribution secret type');
  }
}