import OAuthProvider, {
  AuthorizationError,
  CimdFetchError,
  type OAuthHelpers,
  authorizationErrorRedirect,
} from '@cloudflare/workers-oauth-provider';

interface Env {
  OAUTH_KV: KVNamespace;
  OAUTH_PROVIDER: OAuthHelpers;
  PUBLIC_URL: string; // https://mcp.example.com, this Worker's own address
  ORIGIN_URL: string; // https://mcp-origin.example.com, where the MCP server answers
  ALLOWED_USERS: string; // comma-separated GitHub numeric ids
  JWT_PRIVATE_JWK: string; // Worker secret: the signing key, as a JSON JWK
  GITHUB_CLIENT_ID: string;
  GITHUB_CLIENT_SECRET: string; // Worker secret
}

type Props = { userId: string };

/* ---- Signed token for the MCP server ------------------------------------------------ */

const b64url = (bytes: ArrayBuffer | Uint8Array): string =>
  btoa(String.fromCharCode(...new Uint8Array(bytes)))
    .replace(/\+/g, '-')
    .replace(/\//g, '_')
    .replace(/=+$/, '');
const encode = (value: unknown): string => b64url(new TextEncoder().encode(JSON.stringify(value)));

async function signToken(privateJwk: JsonWebKey & { kid: string }, claims: Record<string, unknown>): Promise<string> {
  const key = await crypto.subtle.importKey('jwk', privateJwk, { name: 'ECDSA', namedCurve: 'P-256' }, false, ['sign']);
  const input = `${encode({ alg: 'ES256', typ: 'at+jwt', kid: privateJwk.kid })}.${encode(claims)}`;
  const signature = await crypto.subtle.sign({ name: 'ECDSA', hash: 'SHA-256' }, key, new TextEncoder().encode(input));
  return `${input}.${b64url(signature)}`;
}

/** The public half of the signing key, as the document the MCP server fetches. */
function jwks(env: Env): Response {
  const { kty, crv, x, y, kid } = JSON.parse(env.JWT_PRIVATE_JWK);
  return Response.json(
    { keys: [{ kty, crv, x, y, kid, use: 'sig', alg: 'ES256' }] },
    { headers: { 'Cache-Control': 'public, max-age=300' } },
  );
}

/* ---- /mcp: check the client's token (the library did), then pass the request on ------ */

const allowed = (env: Env, userId: string): boolean =>
  env.ALLOWED_USERS.split(',').map((id) => id.trim()).includes(userId);

const mcp = {
  async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
    const { userId } = ctx.props as Props;
    if (!allowed(env, userId)) return new Response('Forbidden', { status: 403 });
    const now = Math.floor(Date.now() / 1000);
    const token = await signToken(JSON.parse(env.JWT_PRIVATE_JWK), {
      iss: env.PUBLIC_URL,
      aud: `${env.ORIGIN_URL}/mcp`,
      sub: userId,
      iat: now,
      exp: now + 60,
    });
    const url = new URL(request.url);
    const upstream = new Request(new URL(url.pathname + url.search, env.ORIGIN_URL), request);
    upstream.headers.set('Authorization', `Bearer ${token}`);
    return fetch(upstream);
  },
};

/* ---- Everything else: sign-in and the key document ------------------------------------ */

const escapeHtml = (value: string): string => value.replace(/[&<>"']/g, (c) => `&#${c.charCodeAt(0)};`);

const defaultHandler = {
  async fetch(request: Request, env: Env): Promise<Response> {
    const url = new URL(request.url);
    try {
      if (url.pathname === '/.well-known/jwks.json') return jwks(env);
      if (url.pathname === '/authorize') return await authorize(request, env);
      if (url.pathname === '/callback') return await callback(request, env);
    } catch (error) {
      if (error instanceof AuthorizationError && error.redirectTo) return Response.redirect(error.redirectTo, 302);
      if (error instanceof AuthorizationError || error instanceof CimdFetchError) {
        return new Response('This sign-in request is not valid. Start again from the client.', { status: 400 });
      }
      throw error;
    }
    return new Response('Not found', { status: 404 });
  },
};

async function authorize(request: Request, env: Env): Promise<Response> {
  const oauth = env.OAUTH_PROVIDER;

  if (request.method === 'GET') {
    const authRequest = await oauth.parseAuthRequest(request);
    const details = await oauth.describeConsent(authRequest);
    const consent = await oauth.beginConsent(authRequest);
    consent.headers.set('Content-Type', 'text/html; charset=utf-8');
    const origin = details.clientDomain
      ? `Published by <strong>${escapeHtml(details.clientDomain)}</strong>.`
      : 'This app registered itself; its name is not verified.';
    return new Response(
      `<!doctype html><meta charset="utf-8"><title>Authorize</title>
<h1>Allow ${escapeHtml(details.clientName)} to use your vault?</h1>
<p>${origin} Access will be sent to <strong>${escapeHtml(details.redirectHost)}</strong>.</p>
<form method="post"><input type="hidden" name="handle" value="${escapeHtml(consent.handle)}">
<button name="decision" value="approve">Allow</button> <button name="decision" value="deny">Deny</button></form>`,
      { headers: consent.headers },
    );
  }

  const form = await request.formData();
  const handle = String(form.get('handle'));
  if (form.get('decision') !== 'approve') {
    const denied = await oauth.denyConsent(request, handle);
    return new Response(null, { status: 302, headers: denied.headers });
  }
  const approved = await oauth.approveConsent(request, handle);
  const verifier = crypto.randomUUID() + crypto.randomUUID();
  const { state, headers } = await oauth.beginUpstream(approved.request, {
    data: { verifier },
    headers: approved.headers,
  });
  const challenge = b64url(await crypto.subtle.digest('SHA-256', new TextEncoder().encode(verifier)));
  const github = new URL('https://github.com/login/oauth/authorize');
  github.search = new URLSearchParams({
    client_id: env.GITHUB_CLIENT_ID,
    redirect_uri: `${env.PUBLIC_URL}/callback`,
    state,
    code_challenge: challenge,
    code_challenge_method: 'S256',
  }).toString();
  headers.set('Location', github.toString());
  return new Response(null, { status: 302, headers });
}

async function callback(request: Request, env: Env): Promise<Response> {
  const oauth = env.OAUTH_PROVIDER;
  const params = new URL(request.url).searchParams;
  const { request: original, data, headers } = await oauth.finishUpstream<{ verifier: string }>(request);

  const denied = (): Response => {
    headers.set('Location', authorizationErrorRedirect(original, 'access_denied'));
    return new Response(null, { status: 302, headers });
  };
  if (params.get('error') || !params.get('code')) return denied();

  const tokenResponse = await fetch('https://github.com/login/oauth/access_token', {
    method: 'POST',
    headers: { Accept: 'application/json', 'Content-Type': 'application/json' },
    body: JSON.stringify({
      client_id: env.GITHUB_CLIENT_ID,
      client_secret: env.GITHUB_CLIENT_SECRET,
      code: params.get('code'),
      redirect_uri: `${env.PUBLIC_URL}/callback`,
      code_verifier: data.verifier,
    }),
  });
  const { access_token } = (await tokenResponse.json()) as { access_token?: string };
  if (!access_token) return denied();

  const userResponse = await fetch('https://api.github.com/user', {
    headers: { Authorization: `Bearer ${access_token}`, 'User-Agent': 'mcp-front' },
  });
  const user = (await userResponse.json()) as { id?: number };
  const userId = String(user.id ?? '');
  if (!userId || !allowed(env, userId)) return denied();

  const { redirectTo } = await oauth.completeAuthorization({
    request: original,
    userId,
    metadata: {},
    scope: original.scope,
    props: { userId } satisfies Props,
  });
  headers.set('Location', redirectTo);
  return new Response(null, { status: 302, headers });
}

export default new OAuthProvider<Env>({
  apiRoute: '/mcp',
  apiHandler: mcp,
  defaultHandler,
  authorizeEndpoint: '/authorize',
  tokenEndpoint: '/oauth/token',
  clientRegistrationEndpoint: '/oauth/register',
  clientIdMetadataDocumentEnabled: true,
  accessTokenTTL: 3600,
  scopesSupported: ['mcp'],
  requiredScopes: ['mcp'],
  resourceMetadata: {
    resource: 'https://mcp.example.com/mcp',
    authorization_servers: ['https://mcp.example.com'],
  },
});