import { createHash, createPrivateKey, createPublicKey, randomUUID, sign as cryptoSign, verify as cryptoVerify, type KeyObject } from 'crypto';
import { canonicalJson, sha256Hex } from '../decision-log';
import type { ProvenanceBlock } from '../receipts/schema';
import type { ReceiptV2, ReceiptV3, ReviewerCascadeEvidence, WorkflowConfig } from './types';

const ED25519_PKCS8_SEED_PREFIX = Buffer.from('302e020100300506032b657004220420', 'hex');
const ED25519_SPKI_PUBLIC_PREFIX = Buffer.from('302a300506032b6570032100', 'hex');

type WorkflowSigningKeys = NonNullable<WorkflowConfig['signingKeys']>;

interface ResolvedSigningKey {
  privateKey: KeyObject;
  publicKeyRaw: Buffer;
  publicKeyHex: string;
  publicKeyFingerprint: string;
}

export interface WorkflowReceiptVerifyOptions {
  publicKeyHex?: string;
  allowedSignerPublicKeysHex?: string[];
}

export function canonicalJSONStringify(value: unknown): string {
  return canonicalJson(value);
}

export function sha256(input: string): string {
  return sha256Hex(input);
}

export function workflowReceiptId(prefix = 'ag_r'): string {
  return `${prefix}_${randomUUID().replace(/-/g, '').slice(0, 16)}`;
}

export async function signReceiptPayload(payload: unknown, signingKeys?: WorkflowSigningKeys): Promise<{
  signature: string;
  public_key_fingerprint: string;
  public_key_hex: string;
}> {
  const resolved = resolveWorkflowSigningKey(signingKeys);
  const canonical = canonicalJson(payload);
  return {
    signature: cryptoSign(null, Buffer.from(canonical, 'utf8'), resolved.privateKey).toString('hex'),
    public_key_fingerprint: resolved.publicKeyFingerprint,
    public_key_hex: resolved.publicKeyHex,
  };
}

export async function buildReceipt(args: {
  workflow_id: string | null;
  outcome_name: string;
  user_id: string;
  cost_usd: number;
  parent_receipt_id: string | null;
  chain_validation_hash: string | null;
  workflow_checkpoint_idx: number | null;
  workflow_total_spend_usd_to_date: number | null;
  is_checkpoint?: boolean;
  checkpoint_label?: string | null;
  reviewer_cascade?: ReviewerCascadeEvidence | null;
  receipt_type?: ReceiptV2['receipt_type'];
  cancelled_reason?: string | null;
  signingKeys?: WorkflowSigningKeys;
  receipt_id?: string;
  signed_at?: string;
}): Promise<ReceiptV2> {
  if (!Number.isFinite(args.cost_usd) || args.cost_usd < 0) {
    throw new Error('Workflow receipt cost must be a finite non-negative number');
  }
  const unsigned: Omit<ReceiptV2, 'signature' | 'public_key_fingerprint' | 'public_key_hex'> = {
    receipt_id: args.receipt_id ?? workflowReceiptId(args.is_checkpoint ? 'ag_cp' : 'ag_r'),
    schema_version: 2,
    outcome_name: args.outcome_name,
    user_id: args.user_id,
    cost_usd: roundUsd(args.cost_usd),
    signed_at: args.signed_at ?? new Date().toISOString(),
    workflow_id: args.workflow_id,
    parent_receipt_id: args.parent_receipt_id,
    chain_validation_hash: args.chain_validation_hash,
    workflow_checkpoint_idx: args.workflow_checkpoint_idx,
    workflow_total_spend_usd_to_date:
      args.workflow_total_spend_usd_to_date === null ? null : roundUsd(args.workflow_total_spend_usd_to_date),
    is_checkpoint: args.is_checkpoint === true,
    checkpoint_label: args.checkpoint_label ?? null,
    reviewer_cascade: args.reviewer_cascade ?? null,
    receipt_type: args.receipt_type ?? (args.is_checkpoint ? 'checkpoint' : 'outcome'),
    cancelled_reason: args.cancelled_reason ?? null,
  };
  return { ...unsigned, ...(await signReceiptPayload(unsigned, args.signingKeys)) };
}

export async function buildReceiptV3(args: {
  workflow_id: string | null;
  outcome_name: string;
  user_id: string;
  cost_usd: number;
  parent_receipt_id: string | null;
  chain_validation_hash: string | null;
  workflow_checkpoint_idx: number | null;
  workflow_total_spend_usd_to_date: number | null;
  provenance: ProvenanceBlock;
  is_checkpoint?: boolean;
  checkpoint_label?: string | null;
  reviewer_cascade?: ReviewerCascadeEvidence | null;
  receipt_type?: ReceiptV2['receipt_type'];
  cancelled_reason?: string | null;
  signingKeys?: WorkflowSigningKeys;
  receipt_id?: string;
  signed_at?: string;
}): Promise<ReceiptV3> {
  const v2 = await buildReceipt(args);
  const unsigned: Omit<ReceiptV3, 'signature' | 'public_key_fingerprint' | 'public_key_hex'> = {
    ...stripReceiptSignature(v2),
    schema_version: 3,
    provenance: args.provenance,
  };
  return { ...unsigned, ...(await signReceiptPayload(unsigned, args.signingKeys)) };
}

export function verifyAnyReceipt(receipt: ReceiptV2 | ReceiptV3, options: WorkflowReceiptVerifyOptions = {}): boolean {
  return verifyReceipt(receipt, options);
}

export function verifyReceipt(receipt: ReceiptV2 | ReceiptV3, options: WorkflowReceiptVerifyOptions = {}): boolean {
  const { signature, public_key_hex: embeddedPublicKeyHex, public_key_fingerprint: fingerprint, ...unsigned } = receipt as ReceiptV2 & { public_key_hex?: string };
  if (!signature || !/^[0-9a-fA-F]{128}$/.test(String(signature))) return false;
  const publicKeyHex = String(options.publicKeyHex || embeddedPublicKeyHex || '').toLowerCase();
  if (!/^[0-9a-f]{64}$/.test(publicKeyHex)) return false;
  const allowed = normalizeAllowedSignerKeys(options.allowedSignerPublicKeysHex);
  if (allowed && !allowed.has(publicKeyHex)) return false;
  const publicKeyRaw = Buffer.from(publicKeyHex, 'hex');
  if (fingerprint && computePublicKeyFingerprint(publicKeyRaw) !== String(fingerprint)) return false;
  const publicKey = createPublicKey({ key: Buffer.concat([ED25519_SPKI_PUBLIC_PREFIX, publicKeyRaw]), format: 'der', type: 'spki' });
  try {
    return cryptoVerify(null, Buffer.from(canonicalJson(unsigned), 'utf8'), publicKey, Buffer.from(signature, 'hex'));
  } catch {
    return false;
  }
}

export function roundUsd(value: number): number {
  return Math.round(value * 1_000_000) / 1_000_000;
}

function stripReceiptSignature(receipt: ReceiptV2): Omit<ReceiptV2, 'signature' | 'public_key_fingerprint' | 'public_key_hex'> {
  const { signature: _signature, public_key_fingerprint: _fingerprint, public_key_hex: _publicKeyHex, ...unsigned } = receipt;
  return unsigned;
}

function resolveWorkflowSigningKey(signingKeys?: WorkflowSigningKeys): ResolvedSigningKey {
  if (signingKeys) return keyFromSeed(signingKeys.privateKey, signingKeys.publicKey);
  const raw = process.env.AGENTGUARD_SIGNING_KEY_ED25519_PRIVATE;
  if (!raw) throw new Error('Workflow receipt signing key is required');
  return keyFromEnvironment(raw);
}

function keyFromSeed(seed: Uint8Array, expectedPublicKey?: Uint8Array): ResolvedSigningKey {
  if (seed.length !== 32) throw new Error(`Ed25519 private seed must be 32 bytes, got ${seed.length} bytes`);
  const privateKey = createPrivateKey({ key: Buffer.concat([ED25519_PKCS8_SEED_PREFIX, Buffer.from(seed)]), format: 'der', type: 'pkcs8' });
  const publicKeyRaw = rawPublicKeyFromPrivate(privateKey);
  if (expectedPublicKey && Buffer.compare(publicKeyRaw, Buffer.from(expectedPublicKey)) !== 0) {
    throw new Error('Workflow receipt public key does not match private seed');
  }
  return resolved(privateKey, publicKeyRaw);
}

function keyFromEnvironment(raw: string): ResolvedSigningKey {
  const value = raw.trim();
  let privateKey: KeyObject;
  if (value.startsWith('-----BEGIN')) {
    privateKey = createPrivateKey(value);
  } else if (value.startsWith('{')) {
    privateKey = createPrivateKey({ key: JSON.parse(value), format: 'jwk' });
  } else {
    const clean = value.replace(/^0x/, '');
    let bytes: Buffer;
    if (/^[0-9a-fA-F]+$/.test(clean) && clean.length % 2 === 0) bytes = Buffer.from(clean, 'hex');
    else bytes = Buffer.from(value, 'base64');
    privateKey = bytes.length === 32
      ? createPrivateKey({ key: Buffer.concat([ED25519_PKCS8_SEED_PREFIX, bytes]), format: 'der', type: 'pkcs8' })
      : createPrivateKey({ key: bytes, format: 'der', type: 'pkcs8' });
  }
  return resolved(privateKey, rawPublicKeyFromPrivate(privateKey));
}

function rawPublicKeyFromPrivate(privateKey: KeyObject): Buffer {
  const der = createPublicKey(privateKey).export({ format: 'der', type: 'spki' });
  return Buffer.from(der).subarray(-32);
}

function resolved(privateKey: KeyObject, publicKeyRaw: Buffer): ResolvedSigningKey {
  return {
    privateKey,
    publicKeyRaw,
    publicKeyHex: publicKeyRaw.toString('hex'),
    publicKeyFingerprint: computePublicKeyFingerprint(publicKeyRaw),
  };
}

function computePublicKeyFingerprint(publicKeyRaw: Buffer): string {
  return createHash('sha256').update(publicKeyRaw).digest('hex').slice(0, 16);
}

function normalizeAllowedSignerKeys(explicit?: string[]): Set<string> | null {
  const values = [...(explicit ?? [])];
  const envAllowed = process.env.AGENTGUARD_ALLOWED_SIGNER_PUBKEYS_HEX;
  if (envAllowed) values.push(...envAllowed.split(','));
  if (values.length === 0) return null;
  const clean = values.map((value) => String(value).trim().toLowerCase()).filter((value) => /^[0-9a-f]{64}$/.test(value));
  return new Set(clean);
}