Public Verification

'use node';

// convex/config/node/ed25529.ts

import { SnapchargeType } from '@/types';
import crypto, { webcrypto } from 'node:crypto';
import { JsonValueType, SARType } from '@/global/types';
import { JWK, SignJWT, jwtVerify, importJWK } from 'jose';
import {
  parseJSON,
  getTimestamp,
  errorResponse,
  getEnvVariable,
  stableStringify,
  successResponse,
} from '@/global/helpers';

import INTEGRITY_PUBLIC_KEY_V1 from '../../.well-known/uhp-integrity-20250809.public.json';
import SIGNER_PUBLIC_KEY_V1 from '../../.well-known/uhp-signer-20250809.public.json';

// Make sure Web Crypto API is available
// in this convex env
if (!globalThis.crypto) {
  // @ts-ignore
  globalThis.crypto = webcrypto;
}

export type VersionTypes = 'v1';

export type SignProofType = {
  key: JWK;
  kid: string;
  hash: string;
  identity: string;
};

export type GeneratedProof = {
  createdAt: number;
  reason: 'initial' | 'final' | 'update' | 'refund';
  signer: { token: string; hash: string; kid: string };
  integrity: { token: string; hash: string; kid: string };
};

// ===== CENSORED ENV VARS =====
// Keys are still parsed and typed, but values are hidden
export const Ed25519_SIGNER_KEY_V1 = parseJSON<JWK>(
  getEnvVariable('[REDACTED_SIGNER_KEY]')
);
export const Ed25519_INTEGRITY_KEY_V1 = parseJSON<JWK>(
  getEnvVariable('[REDACTED_INTEGRITY_KEY]')
);
export const Ed25519_KID_DATE_V1 = Number(
  getEnvVariable('[REDACTED_KID_DATE]')
);
export const CURRENT_VERSION = String(
  getEnvVariable('[REDACTED_CURRENT_VERSION]')
) as VersionTypes;
// =============================

export const UHPENRY_KID_DATES = { v1: Ed25519_KID_DATE_V1 };
export const UHPENRY_PRIVATE_KEYS = {
  v1: { integrity: Ed25519_INTEGRITY_KEY_V1, signer: Ed25519_SIGNER_KEY_V1 },
};
export const UHPENRY_PUBLIC_KEYS = {
  v1: { integrity: INTEGRITY_PUBLIC_KEY_V1, signer: SIGNER_PUBLIC_KEY_V1 },
};

/**
 * Creates a SHA-256 hash (hex) of the stable-stringified JSON data.
 *
 * @param data - The JSON value to hash.
 * @returns The hex-encoded SHA-256 hash string.
 */
export const hashDeterministicJson = (data: JsonValueType): string => {
  try {
    if (!data || typeof data !== 'object') {
      throw new Error('Invalid data for hashing. Expected a JSON object.');
    }
    const str = stableStringify(data);
    return crypto.createHash('sha256').update(str, 'utf8').digest('hex');
  } catch (error: any) {
    throw new Error(error?.message ?? 'Failed to hash JSON data');
  }
};

/**
 * Retrieves the Key ID (KID) and corresponding JWK keys for a given type and version.
 *
 * @param type - The type of key to retrieve ("integrity" or "signer").
 * @param version - Optional version identifier (defaults to CURRENT_VERSION).
 * @returns An object containing the KID string and the matching integrity/signer keys.
 * @throws If the version is unknown or the keys/KID are not found.
 */
export const getSignKID = (
  type: 'integrity' | 'signer',
  version?: VersionTypes
) => {
  try {
    const targetVersion = version ?? CURRENT_VERSION;

    // Validate type
    if (type !== 'integrity' && type !== 'signer') {
      throw new Error(
        `Invalid type: ${type}. Expected "integrity" or "signer".`
      );
    }

    // Validate keys object
    const keys = UHPENRY_PRIVATE_KEYS[targetVersion];
    if (!keys || typeof keys !== 'object') {
      throw new Error(`No valid keys found for version: ${targetVersion}`);
    }

    // Make sure required JWK properties exist
    const { kty, crv, x, d } = keys[type];
    if (kty !== 'OKP' || crv !== 'Ed25519') {
      throw new Error(
        `Invalid JWK type/curve for version ${targetVersion}. Expected OKP+Ed25519.`
      );
    }
    if (!x) {
      throw new Error(
        `Missing "x" (public key) in keys for version ${targetVersion}.`
      );
    }
    if (!d) {
      throw new Error(
        `Missing "d" (private key) in keys for version ${targetVersion}.`
      );
    }

    // Validate KID date mapping
    const kidDate = String(UHPENRY_KID_DATES[targetVersion]);
    if (!kidDate || kidDate.trim() === '') {
      throw new Error(`No KID date found for version: ${targetVersion}`);
    }

    // Build KID
    const kid = `uhp-${type}-${kidDate}.public.json`;
    return { kid, keys };
  } catch (error: any) {
    throw new Error(error?.message ?? 'Failed to get signing KID');
  }
};

/**
 * Creates a cryptographically signed proof token (JWT) for any Uhpenry entity.
 *
 * This proof can be used to verify the authenticity and integrity of data such as:
 * - Snapcharge transactions
 * - License files (PDF, etc.)
 * - Snapshots or other signed assets
 *
 * The generated token is signed using the provided JWK and includes:
 * - `hash`: A deterministic hash of the signed data
 * - `sub` (subject): An identity string that uniquely identifies the entity (e.g., transaction ID)
 * - `aud` (audience): Always set to 'uhpenry.com' for public verification context
 * - `iss` (issuer): Always set to 'uhpenry.com' to indicate the token origin
 *
 * @typedef {Object} SignProofType
 * @property {JWK} key       - JSON Web Key used to sign the token.
 * @property {string} kid    - Key ID for identifying the signing key.
 * @property {string} hash   - Deterministic hash of the entity's data.
 * @property {string} identity - Unique identifier for the entity this proof is about.
 *
 * @param {SignProofType} args - Signing arguments including key, kid, hash, and entity identity.
 * @returns {Promise<SARType<{ token: string; hash: string } | null>>}
 * Returns a success response with the signed token and hash, or an error response.
 */
export const signProof = async (
  args: SignProofType
): Promise<SARType<{ token: string; hash: string } | null>> => {
  try {
    const { hash, key, identity, kid } = args;

    if (!key || typeof key !== 'object') {
      throw new Error('Signing key is missing or not an object');
    }

    const { kty, crv, x, d } = key as Partial<JWK>;

    if (kty !== 'OKP' || crv !== 'Ed25519') {
      throw new Error(
        `Invalid JWK type/curve. Expected OKP+Ed25519, got kty=${kty} crv=${crv}`
      );
    }

    if (!x) throw new Error('JWK missing public parameter "x"');
    if (!d)
      throw new Error(
        'JWK missing private parameter "d" (private key required for signing)'
      );

    const privateKey = await importJWK(key, 'EdDSA');

    // Ensure it's private key
    if ((privateKey as any).type !== 'private') {
      throw new Error('Imported key is not a private key');
    }

    const token = await new SignJWT({ hash })
      .setProtectedHeader({ alg: 'EdDSA', kid })
      .setIssuedAt()
      .setIssuer('uhpenry.com')
      .setSubject(identity)
      .setAudience('uhpenry.com')
      .sign(privateKey);

    return successResponse('Proof signed successfully', { token, hash });
  } catch (error: any) {
    console.log('checkpoint 5.4', {
      message: error?.message,
      stack: error?.stack,
    });
    return errorResponse(error.message ?? 'Failed to sign proof', null);
  }
};

/**
 * Generates a two-stage cryptographic proof (integrity + signer) for a Snapcharge object.
 *
 * @param {SnapchargeType} data - The original JSON object to sign.
 * @param {'initial' | 'final' | 'update' | 'refund'} reason - Reason for creating this proof.
 * @returns {Promise<GeneratedProof>} - The integrity proof, signer proof, creation timestamp, and reason.
 *
 * @throws {Error} If signing fails for either the integrity or signer stage.
 */
export const generateSignedProof = async (
  Snapcharge: SnapchargeType,
  reason: 'initial' | 'final' | 'update' | 'refund'
): Promise<GeneratedProof> => {
  try {
    if (!Snapcharge || typeof Snapcharge !== 'object') {
      throw new Error(`Invalid Snapcharge: must be a non-null object`);
    }
    if (!Snapcharge._id) {
      throw new Error(`Invalid Snapcharge: missing or invalid "_id"`);
    }

    if (!['initial', 'final', 'refund'].includes(reason)) {
      throw new Error(`Invalid reason: ${reason}`);
    }

    if (Snapcharge?.verifications) {
      throw new Error(
        'Verifications field should be removed from Snapcharge when generating signed proof'
      );
    }

    // Get deterministic hash of the Snapcharge data (hash A)
    let hash: string;
    try {
      hash = hashDeterministicJson(Snapcharge);
    } catch (err: any) {
      throw new Error(`Failed to hash Snapcharge: ${err?.message ?? err}`);
    }
    if (typeof hash !== 'string' || !hash.trim()) {
      throw new Error(`Generated hash for Snapcharge is invalid`);
    }

    // Retrieve keys for integrity and signer
    const { keys: integrityKeys, kid: integrityKid } = getSignKID('integrity');
    const { keys: signerKeys, kid: signerKid } = getSignKID('signer');

    if (
      !integrityKeys?.integrity ||
      typeof integrityKeys.integrity !== 'object'
    ) {
      throw new Error(`Missing or invalid integrity private key`);
    }
    if (!signerKeys?.signer || typeof signerKeys.signer !== 'object') {
      throw new Error(`Missing or invalid signer private key`);
    }

    // Sign hash A with the integrity key
    const integrityParams = {
      hash,
      identity: Snapcharge._id,
      key: integrityKeys.integrity,
      kid: integrityKid,
    };
    const integrityResponse = await signProof(integrityParams);
    if (!integrityResponse?.success || !integrityResponse.payload) {
      throw new Error(
        `Integrity signing failed: ${
          integrityResponse?.message ?? 'unknown error'
        }`
      );
    }
    if (!integrityResponse.payload.token || !integrityResponse.payload.hash) {
      throw new Error(`Integrity signing returned incomplete payload`);
    }

    const integrity = {
      ...integrityResponse.payload,
      kid: integrityKid,
    };

    // Get deterministic hash of the integrity object (hash B)
    let integrityHash: string;
    try {
      integrityHash = hashDeterministicJson(integrity);
    } catch (err: any) {
      throw new Error(`Failed to hash integrity proof: ${err?.message ?? err}`);
    }
    if (typeof integrityHash !== 'string' || !integrityHash.trim()) {
      throw new Error(`Generated hash for integrity proof is invalid`);
    }

    // Sign hash B with the signer key
    const signerParams = {
      hash: integrityHash,
      identity: Snapcharge._id,
      key: signerKeys.signer,
      kid: signerKid,
    };
    const signerResponse = await signProof(signerParams);
    if (!signerResponse?.success || !signerResponse.payload) {
      throw new Error(
        `Signer signing failed: ${signerResponse?.message ?? 'unknown error'}`
      );
    }
    if (!signerResponse.payload.token || !signerResponse.payload.hash) {
      throw new Error(`Signer signing returned incomplete payload`);
    }

    const signer = {
      ...signerResponse.payload,
      kid: signerKid,
    };

    // Return the proofs along with metadata
    return {
      integrity,
      signer,
      createdAt: getTimestamp(),
      reason,
    };
  } catch (error: any) {
    throw new Error(error?.message ?? 'Failed to generate signed proof');
  }
};

/**
 * Retrieves the public keys object based on a given Key ID (kid).
 *
 * Expected kid format: 'uhp-{integrity|signer}-{yyyymmdd}.public.json'
 *
 * @param {string} kid - The Key ID string from the token header.
 * @returns {Record<'integrity' | 'signer', JWK>} - The corresponding public keys for the version.
 * @throws Will throw an error if the version or keys are not found.
 */
export const getPublicKey = (kid: string): JWK => {
  try {
    if (typeof kid !== 'string' || !kid.trim()) {
      throw new Error(`Invalid KID: must be a non-empty string`);
    }

    // Ensure .public.json suffix
    if (!kid.endsWith('.public.json')) {
      throw new Error(
        `Invalid kid format (missing .public.json suffix): ${kid}`
      );
    }

    // Extract the main part before '.public.json'
    const mainPart = kid.split('.').shift();
    if (!mainPart) {
      throw new Error(`Invalid kid format (empty main part): ${kid}`);
    }

    // Expected format: 'uhp-{type}-{date}'
    const parts = mainPart.split('-');
    if (parts.length !== 3 || parts[0] !== 'uhp') {
      throw new Error(`Unexpected kid structure: ${kid}`);
    }

    const [, type, dateStr] = parts;
    const date = Number(dateStr);

    if (type !== 'integrity' && type !== 'signer') {
      throw new Error(`Invalid key type in kid: ${type}`);
    }
    if (!Number.isInteger(date) || date <= 0) {
      throw new Error(`Invalid date in kid: ${dateStr}`);
    }

    // Find version matching the date from UHPENRY_KID_DATES
    const versionEntry = Object.entries(UHPENRY_KID_DATES).find(
      ([, dateValue]) => dateValue === date
    ) as [VersionTypes, number] | undefined;

    if (!versionEntry) {
      throw new Error(`No version found for date: ${date}`);
    }

    const [version] = versionEntry;

    if (!UHPENRY_PUBLIC_KEYS[version]) {
      throw new Error(`No public keys object found for version: ${version}`);
    }

    const publicKey = UHPENRY_PUBLIC_KEYS[version][type] as JWK;

    if (!publicKey || typeof publicKey !== 'object') {
      throw new Error(
        `No public key found for type "${type}" in version "${version}"`
      );
    }

    // Ensure JWK shape is valid for Ed25519
    const { kty, crv, x } = publicKey as Partial<JWK>;
    if (kty !== 'OKP' || crv !== 'Ed25519') {
      throw new Error(
        `Invalid JWK type/curve for version ${version}. Expected OKP+Ed25519.`
      );
    }

    if (!x) {
      throw new Error(`Public JWK missing "x" parameter in version ${version}`);
    }

    return publicKey;
  } catch (error: any) {
    throw new Error(error?.message ?? 'Failed to get public key from KID');
  }
};

/**
 * Verifies the signed proof object.
 * Throws error if verification fails.
 *
 * @param proof - The proof object generated by generateSignedProof
 * @param originalData - The original Snapcharge JSON data
 * @returns {Promise<SARType<boolean>>} - true if verification succeeds
 */
export const verifySignedProof = async (
  proof: GeneratedProof,
  originalData: SnapchargeType
): Promise<SARType<boolean>> => {
  try {
    // Verify signer token first using public key
    const signerKid = proof.signer.kid;
    const signerPublicKeyJwk = getPublicKey(signerKid);

    const signerKey = await importJWK(signerPublicKeyJwk);

    const signOption = {
      issuer: 'uhpenry.com',
      audience: 'uhpenry.com',
      algorithms: ['EdDSA'],
      requiredClaims: ['sub', 'aud', 'iss'],
      subject: originalData._id,
    };

    const { payload: signerPayload } = await jwtVerify(
      proof.signer.token,
      signerKey,
      signOption
    );

    // The signer payload contains the integrity hash
    const signerHash = signerPayload.hash as string;

    // Verify integrity token using public key
    const integrityKid = proof.integrity.kid;
    const integrityPublicKeyJwk = getPublicKey(integrityKid);
    const integrityKey = await importJWK(integrityPublicKeyJwk);

    const { payload: integrityPayload } = await jwtVerify(
      proof.integrity.token,
      integrityKey,
      signOption
    );

    // The integrity payload contains the original data hash
    const integrityHash = integrityPayload.hash as string;

    // Recalculate original data hash and compare with integrity hash
    const recalculatedOriginalHash = hashDeterministicJson(originalData);
    if (recalculatedOriginalHash !== integrityHash) {
      throw new Error('Original data hash does not match integrity hash');
    }

    // Recalculate integrity object hash and compare with signer hash
    const recalculatedIntegrityHash = hashDeterministicJson(proof.integrity);

    if (recalculatedIntegrityHash !== signerHash) {
      throw new Error('Integrity object hash does not match signer hash');
    }

    return successResponse('Signed proof verified successfully', true);
  } catch (error: any) {
    return errorResponse(
      error.message ?? 'Failed to verify signed proof',
      false
    );
  }
};

/**
 * Verifies the most recent proof for a given Snapcharge transaction.
 *
 * This function extracts the latest verification record from the provided `Snapcharge` object
 * and validates it using `verifySignedProof`. It ensures that:
 * - At least one verification record exists.
 * - The latest verification record is not an "initial" verification.
 * - The signed proof is successfully verified.
 *
 * @param {SnapchargeType} Snapcharge - The Snapcharge object containing transaction data and verifications.
 *
 * @returns {Promise<SARType<boolean>}
 * Returns a `successResponse` if verification passes, otherwise an `errorResponse` with the failure reason.
 *
 * @throws {Error} If there are no verifications, the latest proof is invalid,
 * or verification fails at any stage.
 */
export const verifySnapchargeProof = async (
  Snapcharge: SnapchargeType
): Promise<SARType<boolean>> => {
  try {
    const { verifications = [], ...SnapchargeJson } = Snapcharge;
    if (verifications.length === 0) {
      throw new Error('No verification data available.');
    }

    const proof = verifications.at(-1);
    if (!proof) {
      throw new Error('No valid verification record found.');
    }

    if (proof.reason === 'initial') {
      throw new Error(
        'Initial verification is not valid for this action. Please wait until payment and fee processing are complete, then try again.'
      );
    }

    const response = await verifySignedProof(proof, SnapchargeJson);
    if (!response.success || !response.payload) {
      throw new Error(response.message || 'Verification failed.');
    }

    return successResponse(response.message ?? '', true);
  } catch (error: any) {
    return errorResponse(error?.message ?? '', false);
  }
};