import { env } from "cloudflare:workers";
import { z } from "zod";
import { run } from "../db";
import { accountDb } from "../scoped-db";
import { forgetLocator, registerLocator } from "../scoped-db/locators";
import { randomToken, sha256Hex } from "../crypto";
import {
  CODE_PREFIX,
  holderOf,
  REFRESH_PREFIX,
  REFRESH_TTL_MS,
  refreshTokenInsert,
  secretFor,
  storeRefreshToken,
  type StoredGrant,
} from "./oauth-store";
import { consumeAuthToken, storeAuthToken } from "../auth-tokens";
import { loadActiveUserById, PAT_PREFIX, toViewer, viewerFromSessionCookie } from "../auth";

import { hasPermission, loadRepo, loadRepoById, repoPermission } from "../repos";
import { enforceRateLimit, requestIp } from "../limits";
import { MCP_AUDIENCE, parseTokenPermissions } from "../token-scope";
import { deleteAccessToken, mintFineGrainedToken, purgeExpiredOAuthTokens, rotateAccessToken } from "./grants";
import { OAuthError, redirectMatches, registerClient, resolveClient } from "./oauth-clients";
import type { OAuthClient } from "./oauth-clients";
import { consentPage, errorPage, signInPage } from "./oauth-consent";
import { EMAIL_UNVERIFIED_MESSAGE, canWriteAsViewer } from "~/lib/email-verification";
import { formatScope, parseScope, SCOPES_SUPPORTED } from "./scopes";
import type { Viewer } from "~/lib/types";

const ACCESS_TTL_MS = 60 * 60 * 1000;
const CODE_TTL_MS = 60 * 1000;
const PENDING_TTL_SECONDS = 600;
const PKCE_VERIFIER = /^[A-Za-z0-9._~-]{43,128}$/;

export const MCP_PATH = "/mcp";

export function resourceUrl(origin: string): string {
  return `${origin}${MCP_PATH}`;
}

export function protectedResourceMetadataUrl(origin: string): string {
  return `${origin}/.well-known/oauth-protected-resource${MCP_PATH}`;
}

function json(body: unknown, status = 200): Response {
  return Response.json(body, {
    status,
    headers: { "cache-control": "no-store", "access-control-allow-origin": "*" },
  });
}

export function protectedResourceMetadata(origin: string): Response {
  return json({
    resource: resourceUrl(origin),
    authorization_servers: [origin],
    bearer_methods_supported: ["header"],
    resource_name: `${env.FORGE_NAME} MCP server`,
  });
}

export function authorizationServerMetadata(origin: string): Response {
  return json({
    issuer: origin,
    authorization_endpoint: `${origin}/oauth/authorize`,
    token_endpoint: `${origin}/oauth/token`,
    registration_endpoint: `${origin}/oauth/register`,
    revocation_endpoint: `${origin}/oauth/revoke`,
    response_types_supported: ["code"],
    response_modes_supported: ["query"],
    grant_types_supported: ["authorization_code", "refresh_token"],
    code_challenge_methods_supported: ["S256"],
    token_endpoint_auth_methods_supported: ["none"],
    revocation_endpoint_auth_methods_supported: ["none"],
    scopes_supported: SCOPES_SUPPORTED,
    client_id_metadata_document_supported: true,
    authorization_response_iss_parameter_supported: true,
  });
}

/** `forgeAccess`: the user opted in to the token also working for git, LFS and REST, not just /mcp. */

type PendingAuthorization = {
  userId: string;
  clientId: string;
  clientName: string;
  redirectUri: string;
  state: string | null;
  codeChallenge: string;
  resource: string;
  grant: StoredGrant;
};

function redirectWith(redirectUri: string, params: Record<string, string | null>): Response {
  const target = new URL(redirectUri);
  for (const [key, value] of Object.entries(params)) if (value !== null) target.searchParams.set(key, value);
  return new Response(null, { status: 302, headers: { location: target.toString(), "cache-control": "no-store" } });
}

async function resolveRequestedRepos(scope: ReturnType<typeof parseScope>, viewer: Viewer) {
  const granted: Array<{ id: string; fullName: string }> = [];
  const unavailable: string[] = [];
  for (const { owner, name } of scope.repos) {
    const repo = await loadRepo(owner, name);
    const permission = repo ? await repoPermission(repo, viewer) : "none";
    if (repo && hasPermission(permission, "read"))
      granted.push({ id: repo.id, fullName: `${repo.owner_name}/${repo.name}` });
    else unavailable.push(`${owner}/${name}`);
  }
  return { granted, unavailable };
}

async function startAuthorization(request: Request, url: URL): Promise<Response> {
  const params = url.searchParams;
  const clientId = params.get("client_id");
  const redirectUri = params.get("redirect_uri");
  if (!clientId || !redirectUri) return errorPage("client_id and redirect_uri are required");
  let client: OAuthClient;
  try {
    client = await resolveClient(clientId, requestIp(request));
  } catch (error) {
    return errorPage(error instanceof OAuthError ? error.message : "Unknown client");
  }
  if (!redirectMatches(client.redirectUris, redirectUri))
    return errorPage("redirect_uri is not registered for this client");
  const state = params.get("state");
  const fail = (error: string, description: string) =>
    redirectWith(redirectUri, { error, error_description: description, state, iss: url.origin });
  if (params.get("response_type") !== "code") return fail("unsupported_response_type", "Only code is supported");
  const challenge = params.get("code_challenge");
  if (!challenge || params.get("code_challenge_method") !== "S256") {
    return fail("invalid_request", "PKCE with S256 is required");
  }
  const resource = params.get("resource") ?? resourceUrl(url.origin);
  if (resource.replace(/\/$/, "") !== resourceUrl(url.origin)) return fail("invalid_target", "Unknown resource");

  const viewer = await viewerFromSessionCookie(request);
  if (!viewer) return signInPage(`${url.pathname}${url.search}`);
  if (!canWriteAsViewer(viewer)) return errorPage(EMAIL_UNVERIFIED_MESSAGE, 403);
  const scope = parseScope(params.get("scope"));
  const { granted, unavailable } = await resolveRequestedRepos(scope, viewer);
  if (scope.repos.length > 0 && granted.length === 0)
    return fail("invalid_scope", "None of the requested repositories are accessible");
  const pending: PendingAuthorization = {
    userId: viewer.id,
    clientId: client.id,
    clientName: client.name,
    redirectUri,
    state,
    codeChallenge: challenge,
    resource,
    grant: {
      permissions: scope.permissions,
      allRepos: scope.repos.length === 0,
      repoIds: granted.map((repo) => repo.id),
      repoNames: granted.map((repo) => repo.fullName),
      agent: client.name,
      forgeAccess: false,
    },
  };
  const requestId = randomToken("authz");
  // AUTH-A8: pending consent lives in auth_tokens so approval consumes it atomically (one code per consent).
  await storeAuthToken("mcp-authz", requestId, viewer.id, JSON.stringify(pending), PENDING_TTL_SECONDS);
  return consentPage({
    requestId,
    clientName: client.name,
    clientId: client.id,
    redirectHost: new URL(redirectUri).origin,
    viewerName: viewer.name,
    permissions: scope.permissions,
    repos: pending.grant.repoNames,
    unavailable,
    allRepos: pending.grant.allRepos,
    expiresInMinutes: REFRESH_TTL_MS / 60_000,
    forgeRequested: scope.forge,
  });
}

/** The consent form is a browser POST, so it must be same-origin and bound to the signed-in user. */
function sameOrigin(request: Request, url: URL): boolean {
  const site = request.headers.get("sec-fetch-site");
  if (site) return site === "same-origin";
  return request.headers.get("origin") === url.origin;
}

const pendingShape = z.object({
  userId: z.string(),
  clientId: z.string(),
  clientName: z.string(),
  redirectUri: z.string(),
  state: z.string().nullable(),
  codeChallenge: z.string(),
  resource: z.string(),
  grant: z.object({
    permissions: z.record(z.string(), z.string()),
    allRepos: z.boolean(),
    repoIds: z.array(z.string()),
    repoNames: z.array(z.string()),
    agent: z.string(),
    forgeAccess: z.boolean(),
  }),
});

function parsePending(raw: string): PendingAuthorization {
  const parsed = pendingShape.parse(JSON.parse(raw));
  return {
    ...parsed,
    grant: { ...parsed.grant, permissions: parseTokenPermissions(JSON.stringify(parsed.grant.permissions)) },
  };
}

async function finishAuthorization(request: Request, url: URL): Promise<Response> {
  if (!sameOrigin(request, url)) return errorPage("Cross-site authorization request blocked", 403);
  const viewer = await viewerFromSessionCookie(request);
  if (!viewer) return errorPage("Your session expired; start again from the client", 401);
  const form = await request.formData();
  const requestId = String(form.get("request_id") ?? "");
  const consumed = requestId ? await consumeAuthToken("mcp-authz", requestId) : null;
  if (!consumed || consumed.userId !== viewer.id) return errorPage("This authorization request expired; start again");
  const pending = parsePending(consumed.payload);
  if (form.get("decision") !== "approve" || !canWriteAsViewer(viewer)) {
    return redirectWith(pending.redirectUri, { error: "access_denied", state: pending.state, iss: url.origin });
  }
  const grant: StoredGrant = { ...pending.grant, forgeAccess: form.get("forge_access") === "on" };
  const code = secretFor(CODE_PREFIX, viewer.id);
  const codeHash = await sha256Hex(code);
  await accountDb(viewer.id).run(
    `INSERT INTO oauth_codes (code_hash, client_id, user_id, redirect_uri, code_challenge, resource, grant_json, expires_at)
     VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
    codeHash,
    pending.clientId,
    viewer.id,
    pending.redirectUri,
    pending.codeChallenge,
    pending.resource,
    JSON.stringify(grant),
    Date.now() + CODE_TTL_MS,
  );
  await registerLocator("oauth_code", codeHash, viewer.id);
  return redirectWith(pending.redirectUri, { code, state: pending.state, iss: url.origin });
}

async function pkceChallenge(verifier: string): Promise<string> {
  const digest = new Uint8Array(await crypto.subtle.digest("SHA-256", new TextEncoder().encode(verifier)));
  return btoa(String.fromCharCode(...digest))
    .replaceAll("+", "-")
    .replaceAll("/", "_")
    .replace(/=+$/, "");
}

const grantShape = z.object({
  permissions: z.record(z.string(), z.string()),
  allRepos: z.boolean(),
  repoIds: z.array(z.string()),
  repoNames: z.array(z.string()),
  agent: z.string(),
  forgeAccess: z.boolean().default(false),
});

function parseGrant(raw: string): StoredGrant {
  const parsed = grantShape.parse(JSON.parse(raw));
  return { ...parsed, permissions: parseTokenPermissions(JSON.stringify(parsed.permissions)) };
}

function tokenResponse(accessToken: string, refreshToken: string, grant: StoredGrant): Response {
  return json({
    access_token: accessToken,
    token_type: "Bearer",
    expires_in: ACCESS_TTL_MS / 1000,
    refresh_token: refreshToken,
    scope: formatScope(grant.permissions, grant.repoNames, grant.forgeAccess),
  });
}

async function issueTokens(userId: string, clientId: string, resource: string, grant: StoredGrant): Promise<Response> {
  await purgeExpiredOAuthTokens(userId);
  const access = await mintFineGrainedToken({
    userId,
    name: `MCP: ${grant.agent}`,
    scopes: `oauth:${clientId}`,
    grant,
    ttlMs: ACCESS_TTL_MS,
    clientId,
    kind: "oauth",
    audience: grant.forgeAccess ? null : MCP_AUDIENCE,
  });
  return tokenResponse(access.token, await storeRefreshToken(userId, clientId, access.id, resource, grant), grant);
}

type CodeRow = {
  client_id: string;
  user_id: string;
  redirect_uri: string;
  code_challenge: string;
  resource: string;
  grant_json: string;
  expires_at: number;
};

async function exchangeCode(form: FormData): Promise<Response> {
  const code = String(form.get("code") ?? "");
  const verifier = String(form.get("code_verifier") ?? "");
  const clientId = String(form.get("client_id") ?? "");
  const hash = await sha256Hex(code);
  const holder = await holderOf(code, CODE_PREFIX, "oauth_code", hash);
  if (!holder) throw new OAuthError("invalid_grant", "Authorization code is invalid or expired");
  const db = accountDb(holder);
  // Validate before consuming, so a request with a stolen code but no verifier can't burn the real client's code.
  const row = await db.first<CodeRow>("SELECT * FROM oauth_codes WHERE code_hash = ?", hash);
  if (!row || row.expires_at < Date.now())
    throw new OAuthError("invalid_grant", "Authorization code is invalid or expired");
  if (row.client_id !== clientId) throw new OAuthError("invalid_grant", "Code was issued to another client");
  if (String(form.get("redirect_uri") ?? "") !== row.redirect_uri)
    throw new OAuthError("invalid_grant", "redirect_uri mismatch");
  if (!PKCE_VERIFIER.test(verifier) || (await pkceChallenge(verifier)) !== row.code_challenge) {
    throw new OAuthError("invalid_grant", "PKCE verification failed");
  }
  const resource = form.get("resource");
  if (resource !== null && String(resource).replace(/\/$/, "") !== row.resource.replace(/\/$/, "")) {
    throw new OAuthError("invalid_target", "Resource does not match the authorization request");
  }
  if (!(await db.first("DELETE FROM oauth_codes WHERE code_hash = ? RETURNING code_hash", hash)))
    throw new OAuthError("invalid_grant", "Authorization code is invalid or expired");
  await forgetLocator("oauth_code", hash);
  if (!(await loadActiveUserById(row.user_id))) throw new OAuthError("invalid_grant", "This account is disabled");
  return issueTokens(row.user_id, row.client_id, row.resource, parseGrant(row.grant_json));
}

type RefreshRow = {
  client_id: string;
  user_id: string;
  access_token_id: string;
  resource: string;
  grant_json: string;
  expires_at: number;
  used_at: number | null;
};

/** Re-checks every granted repository against the user's current access. */
async function stillGranted(userId: string, grant: StoredGrant): Promise<StoredGrant | null> {
  const user = await loadActiveUserById(userId);
  if (!user) return null;
  const viewer = toViewer(user);
  const repoIds: string[] = [];
  const repoNames: string[] = [];
  for (const repoId of grant.repoIds) {
    const repo = await loadRepoById(repoId);
    if (!repo || !hasPermission(await repoPermission(repo, viewer), "read")) continue;
    repoIds.push(repo.id);
    repoNames.push(`${repo.owner_name}/${repo.name}`);
  }
  if (!grant.allRepos && repoIds.length === 0) return null;
  return { ...grant, repoIds, repoNames };
}

/**
 * Refresh tokens rotate. Claiming one is a single conditional UPDATE, so two concurrent refreshes
 * cannot both win; presenting an already-used token is treated as theft and revokes the grant (the
 * access_tokens trigger takes its refresh tokens and delegations with it).
 */
async function refreshTokens(form: FormData): Promise<Response> {
  const presented = String(form.get("refresh_token") ?? "");
  const hash = await sha256Hex(presented);
  const clientId = String(form.get("client_id") ?? "");
  const holderId = await holderOf(presented, REFRESH_PREFIX, "oauth_refresh", hash);
  if (!holderId) throw new OAuthError("invalid_grant", "Refresh token is invalid, expired or already used");
  const db = accountDb(holderId);
  // Refused before the token is claimed, so a disabled user's grant is neither rotated nor revoked: re-enabling
  // the account brings it back.
  const holder = await db.first<{ user_id: string }>(
    "SELECT user_id FROM oauth_refresh_tokens WHERE token_hash = ?",
    hash,
  );
  if (holder && !(await loadActiveUserById(holder.user_id)))
    throw new OAuthError("invalid_grant", "This account is disabled");
  const row = await db.first<RefreshRow>(
    `UPDATE oauth_refresh_tokens SET used_at = ?1
     WHERE token_hash = ?2 AND used_at IS NULL AND client_id = ?3 AND expires_at > ?1 RETURNING *`,
    Date.now(),
    hash,
    clientId,
  );
  if (!row) {
    const known = await db.first<RefreshRow>("SELECT * FROM oauth_refresh_tokens WHERE token_hash = ?", hash);
    if (known?.used_at !== null && known?.used_at !== undefined)
      await deleteAccessToken(known.access_token_id, known.user_id);
    throw new OAuthError("invalid_grant", "Refresh token is invalid, expired or already used");
  }
  try {
    const grant = await stillGranted(row.user_id, parseGrant(row.grant_json));
    if (!grant) {
      await deleteAccessToken(row.access_token_id, row.user_id);
      throw new OAuthError("invalid_grant", "None of the granted repositories are accessible any more");
    }
    const next = await refreshTokenInsert(row.user_id, row.client_id, row.access_token_id, row.resource, grant);
    const access = await rotateAccessToken(
      row.access_token_id,
      row.user_id,
      grant.repoIds,
      ACCESS_TTL_MS,
      next.statement,
    );
    await registerLocator("oauth_refresh", next.hash, row.user_id);
    await run("UPDATE oauth_clients SET last_used_at = ? WHERE id = ?", Date.now(), row.client_id);
    return tokenResponse(access.token, next.refresh, grant);
  } catch (error) {
    if (error instanceof OAuthError) throw error;
    if (!(await db.first("SELECT 1 AS ok FROM access_tokens WHERE id = ?", row.access_token_id))) {
      throw new OAuthError("invalid_grant", "This grant was revoked");
    }
    // The rotation batch is one transaction, so nothing was committed and the presented token stays valid.
    await db.run("UPDATE oauth_refresh_tokens SET used_at = NULL WHERE token_hash = ?", hash);
    throw error;
  }
}

async function tokenEndpoint(request: Request): Promise<Response> {
  await enforceRateLimit("AUTH_LIMITER", `oauth-token:${requestIp(request)}`);
  const form = await request.formData().catch(() => null);
  if (!form) throw new OAuthError("invalid_request", "Body must be application/x-www-form-urlencoded");
  const grantType = form.get("grant_type");
  if (grantType === "authorization_code") return exchangeCode(form);
  if (grantType === "refresh_token") return refreshTokens(form);
  throw new OAuthError("unsupported_grant_type", "Use authorization_code or refresh_token");
}

/** RFC 7009: always 200, whether or not the token existed. */
async function revokeEndpoint(request: Request): Promise<Response> {
  await enforceRateLimit("AUTH_LIMITER", `oauth-revoke:${requestIp(request)}`);
  const form = await request.formData().catch(() => null);
  const token = String(form?.get("token") ?? "");
  const hash = await sha256Hex(token);
  const refreshHolder = await holderOf(token, REFRESH_PREFIX, "oauth_refresh", hash);
  const refresh = refreshHolder
    ? await accountDb(refreshHolder).first<RefreshRow>("SELECT * FROM oauth_refresh_tokens WHERE token_hash = ?", hash)
    : null;
  const accessHolder = refresh === null ? await holderOf(token, PAT_PREFIX, "access_token", hash) : null;
  const access =
    refresh === null
      ? accessHolder &&
        (await accountDb(accessHolder).first<{ id: string; user_id: string }>(
          "SELECT id, user_id FROM access_tokens WHERE token_hash = ? AND scopes LIKE 'oauth:%'",
          hash,
        ))
      : { id: refresh.access_token_id, user_id: refresh.user_id };
  if (access) await deleteAccessToken(access.id, access.user_id);
  return new Response(null, { status: 200, headers: { "access-control-allow-origin": "*" } });
}

function oauthErrorResponse(error: OAuthError): Response {
  return json({ error: error.error, error_description: error.message }, error.status);
}

function preflight(): Response {
  return new Response(null, {
    status: 204,
    headers: {
      "access-control-allow-origin": "*",
      "access-control-allow-methods": "GET, POST, OPTIONS",
      "access-control-allow-headers": "authorization, content-type, mcp-protocol-version",
      "access-control-max-age": "86400",
    },
  });
}

/** OAuth 2.1 authorization server endpoints; null when the path is not ours. */
export async function handleOAuthRoutes(request: Request, url: URL): Promise<Response | null> {
  const path = url.pathname;
  if (path.startsWith("/.well-known/oauth-protected-resource")) return protectedResourceMetadata(url.origin);
  if (path === "/.well-known/oauth-authorization-server" || path === "/.well-known/openid-configuration") {
    return authorizationServerMetadata(url.origin);
  }
  if (!path.startsWith("/oauth/")) return null;
  if (request.method === "OPTIONS") return preflight();
  try {
    if (path === "/oauth/authorize") {
      if (request.method === "GET") return await startAuthorization(request, url);
      if (request.method === "POST") return await finishAuthorization(request, url);
    }
    if (request.method !== "POST") return null;
    if (path === "/oauth/token") return await tokenEndpoint(request);
    if (path === "/oauth/register") return await registerClient(request);
    if (path === "/oauth/revoke") return await revokeEndpoint(request);
  } catch (error) {
    if (error instanceof OAuthError) return oauthErrorResponse(error);
    throw error;
  }
  return null;
}