import "server-only";
import { createHash, randomBytes } from "node:crypto";
import { db, filter, findById, findFirst, orNull, withLock, type RecordModel } from "./db";
import {
  isAllowedRedirectUri,
  isCodeChallenge,
  redirectUriMatches,
  requestedPermission,
  scopeOf,
  verifyPkce,
  withParams,
  type GrantPermission,
} from "./oauth-rules";
import { serverOrigin } from "./request-origin";
import { coveredRepos, type TokenScope } from "./tokens";
import { toUser } from "./users";

/*
 * An OAuth 2.1 authorization server for MCP clients (AI agents), as the MCP authorization spec
 * describes it:
 *
 *   /.well-known/oauth-protected-resource[/mcp]   which server issues tokens for /mcp (RFC 9728)
 *   /.well-known/oauth-authorization-server       the endpoints below (RFC 8414)
 *   POST /oauth/register                          dynamic client registration (RFC 7591)
 *   /oauth/authorize                              the user signs in and approves (PKCE S256 only)
 *   POST /oauth/token                             authorization_code and refresh_token grants
 *   POST /oauth/revoke                            token revocation (RFC 7009)
 *
 * Clients are public (no secrets). A grant carries the same permissions as an API token: read or
 * read & write, for all of the user's repositories or selected ones. Access tokens last an hour;
 * refresh tokens rotate on every use and expire after 30 days without one.
 */

export const ACCESS_TOKEN_PREFIX = "vat_";
const REFRESH_TOKEN_PREFIX = "vrt_";
const ACCESS_TOKEN_SECONDS = 3600;
const REFRESH_TOKEN_MS = 30 * 86400e3;
const CODE_MS = 5 * 60e3;
/** Registered clients nobody approves within this long are deleted. */
const UNUSED_CLIENT_MS = 86400e3;

const sha256 = (s: string) => createHash("sha256").update(s).digest("hex");
const newSecret = (prefix: string) => `${prefix}${randomBytes(30).toString("base64url")}`;

/** The authorization server's identifier, which clients check against the metadata they fetched. */
export const issuer = serverOrigin;

/** The MCP endpoint, which is what tokens are issued for (the "resource" of RFC 8707). */
export async function mcpResource() {
  return `${await serverOrigin()}/mcp`;
}

/** A client may name the MCP endpoint or the server itself as the resource it wants a token for. */
async function isOurResource(resource: string) {
  const strip = (s: string) => s.replace(/\/+$/, "");
  return [await mcpResource(), await serverOrigin()].includes(strip(resource));
}

export class OAuthError extends Error {
  constructor(
    readonly code: string,
    message: string,
    readonly status = 400,
  ) {
    super(message);
  }
}

const invalidGrant = (message: string) => new OAuthError("invalid_grant", message);

// ---------------------------------------------------------------------------
// Clients

export type OAuthClient = { id: string; name: string; uri: string | null; redirectUris: string[] };

const toClient = (r: RecordModel): OAuthClient => ({ id: r.id, name: r.name, uri: orNull(r.client_uri), redirectUris: r.redirect_uris ?? [] });

export async function getClient(clientId: string) {
  // Record ids are 15 lowercase letters and digits; anything else can't be one.
  if (!/^[a-z0-9]{15}$/.test(clientId)) return null;
  const row = await findById("oauth_clients", clientId);
  return row ? toClient(row) : null;
}

/** Registers a client from its RFC 7591 metadata and answers with the registration. */
export async function registerClient(metadata: Record<string, unknown>) {
  const uris = metadata.redirect_uris;
  if (!Array.isArray(uris) || uris.length === 0 || uris.length > 10 || !uris.every((u) => typeof u === "string" && isAllowedRedirectUri(u))) {
    throw new OAuthError(
      "invalid_redirect_uri",
      "redirect_uris must list 1 to 10 URIs: https, http on localhost, or an app's own scheme, without fragments.",
    );
  }
  const grantTypes = metadata.grant_types;
  if (grantTypes !== undefined && !(Array.isArray(grantTypes) && grantTypes.includes("authorization_code"))) {
    throw new OAuthError("invalid_client_metadata", "Only the authorization_code grant (with refresh_token) is supported.");
  }
  const responseTypes = metadata.response_types;
  if (responseTypes !== undefined && !(Array.isArray(responseTypes) && responseTypes.includes("code"))) {
    throw new OAuthError("invalid_client_metadata", "Only the code response type is supported.");
  }
  const name = typeof metadata.client_name === "string" ? metadata.client_name.trim().slice(0, 100) : "";
  const uri = typeof metadata.client_uri === "string" && /^https?:\/\//.test(metadata.client_uri) ? metadata.client_uri.slice(0, 500) : "";

  const now = Date.now();
  const row = await (await db()).collection("oauth_clients").create({
    name: name || new URL(uris[0]).hostname || "Unnamed application",
    client_uri: uri,
    redirect_uris: uris,
    created_at: now,
  });
  pruneUnusedClients().catch((e) => console.error("[oauth] could not prune clients:", e));
  // Every client is public: whatever authentication method it asked for, it gets "none" (RFC 7591
  // lets the server replace requested values).
  return {
    client_id: row.id,
    client_id_issued_at: Math.floor(now / 1000),
    client_name: row.name,
    ...(uri ? { client_uri: uri } : {}),
    redirect_uris: uris,
    grant_types: ["authorization_code", "refresh_token"],
    response_types: ["code"],
    token_endpoint_auth_method: "none",
  };
}

async function pruneUnusedClients() {
  const { items } = await (await db()).collection("oauth_clients").getList(1, 100, {
    filter: filter("last_authorized_at = 0 && created_at < {:cutoff}", { cutoff: Date.now() - UNUSED_CLIENT_MS }),
    fields: "id",
    skipTotal: true,
  });
  for (const { id } of items) await (await db()).collection("oauth_clients").delete(id);
}

// ---------------------------------------------------------------------------
// Authorizing

export type AuthorizeRequest = {
  client: OAuthClient;
  redirectUri: string;
  state: string | null;
  codeChallenge: string;
  /** What the client asked for; the user may narrow it. */
  permission: GrantPermission;
};

/**
 * Checks the query of an authorization request. Problems with the client or its redirect URI are
 * shown to the user (`fatal`): sending anything to an unverified address could leak to an attacker.
 * Other problems go back to the client (`redirect`), as OAuth specifies.
 */
export async function parseAuthorizeRequest(
  params: Record<string, string | undefined>,
): Promise<{ request: AuthorizeRequest } | { fatal: string } | { redirect: string }> {
  const client = await getClient(params.client_id ?? "");
  if (!client) return { fatal: "This application isn't registered here (anymore). Start connecting it again from the application." };
  const redirectUri = params.redirect_uri ?? "";
  if (!redirectUriMatches(client.redirectUris, redirectUri)) {
    return { fatal: "The application asked to return to an address it didn't register, so Vectrum won't send you there." };
  }
  const state = params.state || null;
  const fail = async (error: string, description: string) => ({
    redirect: withParams(redirectUri, { error, error_description: description, state, iss: await issuer() }),
  });
  if (params.response_type !== "code") return fail("unsupported_response_type", "Only response_type=code is supported.");
  const codeChallenge = params.code_challenge ?? "";
  if (params.code_challenge_method !== "S256" || !isCodeChallenge(codeChallenge)) {
    return fail("invalid_request", "PKCE is required, with code_challenge_method=S256.");
  }
  if (params.resource && !(await isOurResource(params.resource))) {
    return fail("invalid_target", `Tokens can only be issued for ${await mcpResource()}.`);
  }
  return { request: { client, redirectUri, state, codeChallenge, permission: requestedPermission(params.scope) } };
}

export async function deniedRedirect(request: AuthorizeRequest) {
  return withParams(request.redirectUri, {
    error: "access_denied",
    error_description: "The user didn't approve the request.",
    state: request.state,
    iss: await issuer(),
  });
}

/** Records the user's approval and returns where to send them: back to the client, with a code. */
export async function approve(request: AuthorizeRequest, userId: string, grant: { permission: GrantPermission; repoIds: string[] | null }) {
  const pb = await db();
  const code = newSecret("");
  const now = Date.now();
  await pb.collection("oauth_codes").create({
    code_hash: sha256(code),
    client: request.client.id,
    user: userId,
    redirect_uri: request.redirectUri,
    code_challenge: request.codeChallenge,
    permission: grant.permission,
    all_repos: !grant.repoIds,
    repos: grant.repoIds ?? [],
    expires_at: now + CODE_MS,
  });
  await pb.collection("oauth_clients").update(request.client.id, { last_authorized_at: now });
  // Codes that were never exchanged.
  const stale = await pb.collection("oauth_codes").getFullList({ filter: filter("expires_at < {:now}", { now }), fields: "id" });
  for (const { id } of stale) await pb.collection("oauth_codes").delete(id);
  return withParams(request.redirectUri, { code, state: request.state, iss: await issuer() });
}

// ---------------------------------------------------------------------------
// Tokens

export type TokenResponse = {
  access_token: string;
  token_type: "Bearer";
  expires_in: number;
  refresh_token: string;
  scope: string;
};

/** New access and refresh tokens, with the fields that store them on a grant. */
function newTokens(permission: GrantPermission) {
  const access = newSecret(ACCESS_TOKEN_PREFIX);
  const refresh = newSecret(REFRESH_TOKEN_PREFIX);
  const now = Date.now();
  return {
    fields: {
      access_hash: sha256(access),
      access_expires_at: now + ACCESS_TOKEN_SECONDS * 1000,
      refresh_hash: sha256(refresh),
      refresh_expires_at: now + REFRESH_TOKEN_MS,
    },
    response: {
      access_token: access,
      token_type: "Bearer",
      expires_in: ACCESS_TOKEN_SECONDS,
      refresh_token: refresh,
      scope: scopeOf(permission),
    } satisfies TokenResponse,
  };
}

// Exchanges and refreshes are read-modify-write sequences on codes and grants; one lock keeps two
// requests from both using the same code or refresh token.
const withTokenLock = <T>(fn: () => Promise<T>) => withLock("oauth", fn);

export async function exchangeCode(input: { code: string; clientId: string; redirectUri: string; codeVerifier: string; resource?: string }) {
  return withTokenLock(async () => {
    const pb = await db();
    const row = input.code ? await findFirst("oauth_codes", filter("code_hash = {:hash}", { hash: sha256(input.code) })) : null;
    if (!row) throw invalidGrant("The authorization code is invalid or was already used.");
    // A code works once, whatever happens next.
    await pb.collection("oauth_codes").delete(row.id);
    if (row.expires_at < Date.now()) throw invalidGrant("The authorization code has expired.");
    if (row.client !== input.clientId) throw invalidGrant("The authorization code was issued to another client.");
    if (row.redirect_uri !== input.redirectUri) throw invalidGrant("redirect_uri doesn't match the authorization request.");
    if (!verifyPkce(input.codeVerifier, row.code_challenge)) throw invalidGrant("code_verifier doesn't match the code_challenge.");
    if (input.resource && !(await isOurResource(input.resource))) {
      throw new OAuthError("invalid_target", `Tokens can only be issued for ${await mcpResource()}.`);
    }

    // Approving an application again replaces what it was allowed before.
    const previous = await pb.collection("oauth_grants").getFullList({
      filter: filter("client = {:client} && user = {:user}", { client: row.client, user: row.user }),
      fields: "id",
    });
    for (const { id } of previous) await pb.collection("oauth_grants").delete(id);

    const tokens = newTokens(row.permission);
    await pb.collection("oauth_grants").create({
      client: row.client,
      user: row.user,
      permission: row.permission,
      all_repos: row.all_repos,
      repos: row.repos,
      ...tokens.fields,
      created_at: Date.now(),
    });
    return tokens.response;
  });
}

export async function refreshGrant(input: { refreshToken: string; clientId: string; resource?: string }) {
  return withTokenLock(async () => {
    const pb = await db();
    const hash = sha256(input.refreshToken);
    const row = input.refreshToken ? await findFirst("oauth_grants", filter("refresh_hash = {:hash}", { hash })) : null;
    if (!row) {
      // A refresh token that was already replaced: someone else may hold a copy, so end the grant.
      const reused = input.refreshToken ? await findFirst("oauth_grants", filter("previous_refresh_hash = {:hash}", { hash }), { fields: "id" }) : null;
      if (reused) await pb.collection("oauth_grants").delete(reused.id);
      throw invalidGrant("The refresh token is invalid, expired or was already used.");
    }
    if (row.client !== input.clientId) throw invalidGrant("The refresh token was issued to another client.");
    if (row.refresh_expires_at < Date.now()) {
      await pb.collection("oauth_grants").delete(row.id);
      throw invalidGrant("The refresh token has expired.");
    }
    if (input.resource && !(await isOurResource(input.resource))) {
      throw new OAuthError("invalid_target", `Tokens can only be issued for ${await mcpResource()}.`);
    }
    const tokens = newTokens(row.permission);
    await pb.collection("oauth_grants").update(row.id, { ...tokens.fields, previous_refresh_hash: hash });
    return tokens.response;
  });
}

/** Resolves an access token to its user and what it may reach, or null if unknown or expired. */
export async function authenticateAccessToken(secret: string) {
  if (!secret.startsWith(ACCESS_TOKEN_PREFIX)) return null;
  const row = await findFirst("oauth_grants", filter("access_hash = {:hash}", { hash: sha256(secret) }), { expand: "user,repos.owner" });
  const user = row?.expand?.user as RecordModel | undefined;
  const now = Date.now();
  if (!row || !user || row.access_expires_at < now) return null;
  // Agents call often; once a minute is enough to show when a grant was last used.
  if (now - row.last_used_at > 60e3) await (await db()).collection("oauth_grants").update(row.id, { last_used_at: now });
  const scope: TokenScope = { permission: row.permission, allRepos: row.all_repos, repos: coveredRepos(row) };
  return { user: toUser(user), scope };
}

/** Ends the grant an access or refresh token belongs to. Unknown tokens are ignored (RFC 7009). */
export async function revokeToken(secret: string) {
  const hash = sha256(secret);
  const row = await findFirst("oauth_grants", filter("access_hash = {:hash} || refresh_hash = {:hash}", { hash }), { fields: "id" });
  if (row) await (await db()).collection("oauth_grants").delete(row.id);
}

// ---------------------------------------------------------------------------
// Connected applications, for the settings page

export type ConnectedApp = {
  id: string;
  name: string;
  uri: string | null;
  permission: GrantPermission;
  allRepos: boolean;
  repos: { id: string; fullName: string }[];
  createdAt: number;
  lastUsedAt: number | null;
};

export async function listConnectedApps(userId: string): Promise<ConnectedApp[]> {
  const rows = await (await db()).collection("oauth_grants").getFullList({
    filter: filter("user = {:userId} && refresh_expires_at > {:now}", { userId, now: Date.now() }),
    sort: "-created_at",
    expand: "client,repos.owner",
  });
  return rows.map((r) => ({
    id: r.id,
    name: r.expand?.client?.name ?? "Unknown application",
    uri: orNull(r.expand?.client?.client_uri),
    permission: r.permission,
    allRepos: r.all_repos,
    repos: coveredRepos(r),
    createdAt: r.created_at,
    lastUsedAt: orNull(r.last_used_at),
  }));
}

export async function disconnectApp(userId: string, grantId: string) {
  const row = await findFirst("oauth_grants", filter("id = {:grantId} && user = {:userId}", { grantId, userId }), { fields: "id" });
  if (row) await (await db()).collection("oauth_grants").delete(row.id);
}

// ---------------------------------------------------------------------------
// Discovery

/** RFC 9728: tells a client that got a 401 from /mcp where to get a token. */
export async function protectedResourceMetadata() {
  return {
    resource: await mcpResource(),
    authorization_servers: [await issuer()],
    scopes_supported: ["read", "write"],
    bearer_methods_supported: ["header"],
    resource_name: "Vectrum",
  };
}

/** RFC 8414: the authorization server's endpoints and what they support. */
export async function authorizationServerMetadata() {
  const origin = await issuer();
  return {
    issuer: origin,
    authorization_endpoint: `${origin}/oauth/authorize`,
    token_endpoint: `${origin}/oauth/token`,
    registration_endpoint: `${origin}/oauth/register`,
    revocation_endpoint: `${origin}/oauth/revoke`,
    scopes_supported: ["read", "write"],
    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"],
    authorization_response_iss_parameter_supported: true,
  };
}