import { env } from "cloudflare:workers";
import { z } from "zod";
import { randomId, sha256Hex } from "../crypto";
import { allows } from "./permissions";
import { jobTokenAccess } from "./job-token";
import { isCloudJob, loadJobExecution } from "./jobs";
import { repoDb } from "../scoped-db";

const MAX_ATTEMPT_DEPTH = 100;

const TOKEN_TTL_SECONDS = 5 * 60;
const MAX_AUDIENCE_LENGTH = 200;
const JWKS_MAX_AGE_SECONDS = 3600;

const rsaPrivateJwk = z.object({
  kty: z.literal("RSA"),
  n: z.string(),
  e: z.string(),
  d: z.string(),
  p: z.string(),
  q: z.string(),
  dp: z.string(),
  dq: z.string(),
  qi: z.string(),
});

type SigningKey = {
  key: CryptoKey;
  publicJwk: { kty: "RSA"; n: string; e: string; alg: "RS256"; use: "sig"; kid: string };
};

let cached: Promise<SigningKey | null> | null = null;

/** The RSA signing key is a private JWK held in Secrets Store and never leaves the Worker. */
async function loadSigningKey(): Promise<SigningKey | null> {
  let raw: string;
  try {
    raw = await env.OIDC_SIGNING_KEY.get();
  } catch (error) {
    console.error("OIDC signing key unavailable", error);
    return null;
  }
  let document: unknown;
  try {
    document = JSON.parse(raw);
  } catch {
    return null;
  }
  const parsed = rsaPrivateJwk.safeParse(document);
  if (!parsed.success) return null;
  const key = await crypto.subtle.importKey(
    "jwk",
    { ...parsed.data, alg: "RS256", ext: false },
    { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" },
    false,
    ["sign"],
  );
  const kid = (await sha256Hex(parsed.data.n)).slice(0, 16);
  return { key, publicJwk: { kty: "RSA", n: parsed.data.n, e: parsed.data.e, alg: "RS256", use: "sig", kid } };
}

function signingKey(): Promise<SigningKey | null> {
  cached ??= loadSigningKey().then((key) => {
    if (!key) cached = null;
    return key;
  });
  return cached;
}

function base64Url(input: string | Uint8Array): string {
  const bytes = typeof input === "string" ? new TextEncoder().encode(input) : input;
  let binary = "";
  for (const byte of bytes) binary += String.fromCharCode(byte);
  return btoa(binary).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}

function issuer(): string {
  return env.PUBLIC_URL.replace(/\/$/, "");
}

async function sign(claims: Record<string, string | number>, key: SigningKey): Promise<string> {
  const header = base64Url(JSON.stringify({ alg: "RS256", typ: "JWT", kid: key.publicJwk.kid }));
  const body = base64Url(JSON.stringify(claims));
  const signature = await crypto.subtle.sign(
    "RSASSA-PKCS1-v1_5",
    key.key,
    new TextEncoder().encode(`${header}.${body}`),
  );
  return `${header}.${body}.${base64Url(new Uint8Array(signature))}`;
}

/** 1 for a first run, +1 per rerun in its `rerun_of` chain, so relying parties can tell attempts apart. */
async function runAttempt(repoId: string, runId: string): Promise<number> {
  const row = await repoDb(repoId).first<{ attempt: number }>(
    `WITH RECURSIVE chain(id, depth) AS (
       SELECT rerun_of, 1 FROM workflow_runs WHERE id = ?1
       UNION ALL SELECT w.rerun_of, chain.depth + 1 FROM workflow_runs w JOIN chain ON w.id = chain.id WHERE chain.depth < ?2
     ) SELECT COUNT(id) + 1 AS attempt FROM chain`,
    runId,
    MAX_ATTEMPT_DEPTH,
  );
  return row?.attempt ?? 1;
}

async function issueToken(request: Request, url: URL): Promise<Response> {
  const access = await jobTokenAccess(request);
  if (!access) return Response.json({ error: "invalid job token" }, { status: 401 });
  if (!allows(access.permissions, "id-token", "write")) {
    return Response.json({ error: "this job lacks the id-token: write permission" }, { status: 403 });
  }
  const execution = await loadJobExecution(access.jobId);
  if (!execution || execution.job.status !== "in_progress") {
    return Response.json({ error: "job is not running" }, { status: 409 });
  }
  const key = await signingKey();
  if (!key) return Response.json({ error: "OIDC is not configured on this instance" }, { status: 501 });
  const { repo, run, spec } = execution;
  const repository = `${repo.owner_name}/${repo.name}`;
  const environment = spec.environment?.name ?? "";
  const subject = environment
    ? `repo:${repository}:environment:${environment}`
    : run.event === "pull_request"
      ? `repo:${repository}:pull_request`
      : `repo:${repository}:ref:${run.ref}`;
  const audience = (url.searchParams.get("audience") ?? `${issuer()}/${repo.owner_name}`).slice(0, MAX_AUDIENCE_LENGTH);
  const now = Math.floor(Date.now() / 1000);
  const workflowRef = `${repository}/${run.workflow_path}@${run.ref}`;
  const token = await sign(
    {
      iss: issuer(),
      sub: subject,
      aud: audience,
      jti: randomId(),
      iat: now,
      nbf: now,
      exp: now + TOKEN_TTL_SECONDS,
      repository,
      repository_owner: repo.owner_name,
      repository_id: repo.id,
      repository_owner_id: repo.owner_id,
      repository_visibility: repo.visibility,
      ref: run.ref,
      ref_type: run.ref.startsWith("refs/tags/") ? "tag" : "branch",
      sha: run.sha,
      workflow: run.workflow_name,
      workflow_ref: workflowRef,
      job_workflow_ref: workflowRef,
      event_name: run.event,
      run_id: run.id,
      run_number: run.run_number,
      run_attempt: await runAttempt(run.repo_id, run.id),
      actor_id: run.actor_id ?? "",
      environment,
      runner_environment: isCloudJob(spec.runsOn) ? "gitfor-hosted" : "self-hosted",
    },
    key,
  );
  return Response.json({ count: token.length, value: token });
}

async function jwks(): Promise<Response> {
  const key = await signingKey();
  return Response.json(
    { keys: key ? [key.publicJwk] : [] },
    { headers: { "cache-control": `public, max-age=${JWKS_MAX_AGE_SECONDS}` } },
  );
}

function discovery(): Response {
  return Response.json(
    {
      issuer: issuer(),
      jwks_uri: `${issuer()}/.well-known/jwks`,
      subject_types_supported: ["public", "pairwise"],
      response_types_supported: ["id_token"],
      id_token_signing_alg_values_supported: ["RS256"],
      scopes_supported: ["openid"],
      claims_supported: ["sub", "aud", "exp", "iat", "iss", "jti", "nbf", "ref", "repository", "sha", "environment"],
    },
    { headers: { "cache-control": `public, max-age=${JWKS_MAX_AGE_SECONDS}` } },
  );
}

export async function handleOidcRoutes(request: Request, url: URL): Promise<Response | null> {
  if (url.pathname === "/.well-known/openid-configuration") return discovery();
  if (url.pathname === "/.well-known/jwks") return jwks();
  if (url.pathname === "/api/actions/oidc/token" && request.method === "GET") return issueToken(request, url);
  return null;
}