import fs from 'node:fs';
import path from 'node:path';
import crypto from 'node:crypto';
import { exportJWK, importPKCS8, importSPKI, type JWK, type KeyLike, SignJWT } from 'jose';
import { config } from '../config.js';

interface KeyMaterial {
  kid: string;
  privatePem: string;
  publicPem: string;
  privateKey: KeyLike;
  publicJwk: JWK;
}

let cached: KeyMaterial | null = null;

async function loadOrGenerate(): Promise<KeyMaterial> {
  if (cached) return cached;
  const dir = path.resolve(config.OAUTH_KEY_DIR);
  fs.mkdirSync(dir, { recursive: true });
  const privPath = path.join(dir, 'oauth-private.pem');
  const pubPath = path.join(dir, 'oauth-public.pem');
  const kidPath = path.join(dir, 'oauth.kid');

  let privatePem: string;
  let publicPem: string;
  let kid: string;

  if (fs.existsSync(privPath) && fs.existsSync(pubPath) && fs.existsSync(kidPath)) {
    privatePem = fs.readFileSync(privPath, 'utf8');
    publicPem = fs.readFileSync(pubPath, 'utf8');
    kid = fs.readFileSync(kidPath, 'utf8').trim();
  } else {
    const { publicKey, privateKey } = crypto.generateKeyPairSync('rsa', { modulusLength: 2048 });
    privatePem = privateKey.export({ type: 'pkcs8', format: 'pem' }) as string;
    publicPem = publicKey.export({ type: 'spki', format: 'pem' }) as string;
    kid = crypto.randomBytes(8).toString('hex');
    fs.writeFileSync(privPath, privatePem, { mode: 0o600 });
    fs.writeFileSync(pubPath, publicPem, { mode: 0o644 });
    fs.writeFileSync(kidPath, kid);
  }

  const privateKey = await importPKCS8(privatePem, 'RS256');
  const publicKey = await importSPKI(publicPem, 'RS256');
  const publicJwk: JWK = { ...(await exportJWK(publicKey)), kid, alg: 'RS256', use: 'sig' };

  cached = { kid, privatePem, publicPem, privateKey, publicJwk };
  return cached;
}

export async function getJWKS(): Promise<{ keys: JWK[] }> {
  const k = await loadOrGenerate();
  return { keys: [k.publicJwk] };
}

export async function signAccessToken(input: {
  subject: string;
  audience: string;
  issuer: string;
  scope?: string;
  ttlSeconds?: number;
}): Promise<string> {
  const k = await loadOrGenerate();
  const ttl = input.ttlSeconds ?? 3600;
  return await new SignJWT({ scope: input.scope ?? '' })
    .setProtectedHeader({ alg: 'RS256', kid: k.kid, typ: 'JWT' })
    .setIssuer(input.issuer)
    .setSubject(input.subject)
    .setAudience(input.audience)
    .setIssuedAt()
    .setExpirationTime(`${ttl}s`)
    .sign(k.privateKey);
}