import { MYINFO_SCOPES, mapMyinfoSnapshot } from "./application-assessment";
import { issueMyinfoEvidence } from "./myinfo-evidence";
import "server-only";
import * as oidc from "openid-client";
import { EncryptJWT, jwtDecrypt, importJWK, exportJWK } from "jose";
import { z } from "zod";
import type { MyinfoConfig } from "./myinfo-config";
import { mapMyinfoProfile } from "./myinfo-profile";

export const transactionCookie = "maxcredit-myinfo";
export const transactionOptions = {
  httpOnly: true,
  secure: true,
  sameSite: "lax" as const,
  path: "/api/myinfo",
  maxAge: 600,
};
const transactionSchema = z.object({
  state: z.string().min(32),
  nonce: z.string().min(32),
  verifier: z.string().min(43),
  dpop: z.object({
    kty: z.literal("EC"),
    crv: z.literal("P-256"),
    x: z.string(),
    y: z.string(),
    d: z.string(),
  }),
  origin: z.string(),
  clientId: z.string(),
  issuer: z.string(),
});
type Transaction = z.infer<typeof transactionSchema>;

async function client(settings: MyinfoConfig) {
  const signing = (await importJWK(settings.signing, "ES256")) as CryptoKey;
  const config = await oidc.discovery(
    new URL(settings.issuer),
    settings.clientId,
    {
      token_endpoint_auth_method: "private_key_jwt",
      id_token_signed_response_alg: "ES256",
      userinfo_signed_response_alg: "ES256",
    },
    oidc.PrivateKeyJwt(
      { key: signing, kid: settings.signing.kid },
      {
        [oidc.modifyAssertion]: (header) => {
          header.typ = "JWT";
        },
      },
    ),
    { timeout: 10, execute: [oidc.enableNonRepudiationChecks] },
  );
  oidc.enableDecryptingResponses(config, ["A256CBC-HS512", "A256GCM"], {
    key: (await importJWK(settings.encryption, "ECDH-ES+A256KW")) as CryptoKey,
    kid: settings.encryption.kid,
    alg: "ECDH-ES+A256KW",
  });
  return config;
}
export async function sealTransaction(
  settings: MyinfoConfig,
  transaction: Transaction,
) {
  return new EncryptJWT(transaction)
    .setProtectedHeader({ alg: "dir", enc: "A256GCM" })
    .setIssuedAt()
    .setExpirationTime("10m")
    .encrypt(Buffer.from(settings.secret, "hex"));
}
export async function openTransaction(settings: MyinfoConfig, cookie: string) {
  const { payload } = await jwtDecrypt(
    cookie,
    Buffer.from(settings.secret, "hex"),
    {
      keyManagementAlgorithms: ["dir"],
      contentEncryptionAlgorithms: ["A256GCM"],
      requiredClaims: ["exp", "iat"],
      maxTokenAge: "10m",
    },
  );
  const transaction = transactionSchema.parse(payload);
  if (
    transaction.origin !== settings.origin ||
    transaction.clientId !== settings.clientId ||
    transaction.issuer !== settings.issuer
  )
    throw new Error("Invalid transaction");
  return transaction;
}
export async function startMyinfo(settings: MyinfoConfig) {
  const config = await client(settings);
  const keyPair = await oidc.randomDPoPKeyPair("ES256", { extractable: true });
  const transaction = {
    state: oidc.randomState(),
    nonce: oidc.randomNonce(),
    verifier: oidc.randomPKCECodeVerifier(),
    dpop: await exportJWK(keyPair.privateKey),
    origin: settings.origin,
    clientId: settings.clientId,
    issuer: settings.issuer,
  };
  const url = await oidc.buildAuthorizationUrlWithPAR(
    config,
    {
      response_type: "code",
      scope: MYINFO_SCOPES,
      redirect_uri: settings.redirectUri,
      state: transaction.state,
      nonce: transaction.nonce,
      code_challenge: await oidc.calculatePKCECodeChallenge(
        transaction.verifier,
      ),
      code_challenge_method: "S256",
    },
    { DPoP: oidc.getDPoPHandle(config, keyPair) },
  );
  return {
    url: url.href,
    cookie: await sealTransaction(
      settings,
      transactionSchema.parse(transaction),
    ),
  };
}
export async function finishMyinfo(
  settings: MyinfoConfig,
  url: URL,
  cookie: string,
  evidenceSecret: string,
) {
  const transaction = await openTransaction(settings, cookie);
  if (
    url.origin !== settings.origin ||
    url.pathname !== "/api/myinfo/callback" ||
    url.searchParams.getAll("state").length !== 1 ||
    url.searchParams.get("state") !== transaction.state
  )
    throw new Error("Invalid callback");
  const config = await client(settings);
  const { kty, crv, x, y } = transaction.dpop;
  const DPoP = oidc.getDPoPHandle(config, {
    privateKey: (await importJWK(transaction.dpop, "ES256")) as CryptoKey,
    publicKey: (await importJWK({ kty, crv, x, y }, "ES256")) as CryptoKey,
  });
  const tokens = await oidc.authorizationCodeGrant(
    config,
    url,
    {
      pkceCodeVerifier: transaction.verifier,
      expectedState: transaction.state,
      expectedNonce: transaction.nonce,
      idTokenExpected: true,
    },
    undefined,
    { DPoP },
  );
  const subject = tokens.claims()?.sub;
  if (!subject || tokens.token_type.toLowerCase() !== "dpop")
    throw new Error("Invalid token");
  const info = await oidc.fetchUserInfo(config, tokens.access_token, subject, {
    DPoP,
  });
  return {
    ...mapMyinfoProfile(info.person_info),
    evidence: await issueMyinfoEvidence(
      mapMyinfoSnapshot(info.person_info, settings.mode),
      evidenceSecret,
      settings.origin,
    ),
  };
}