import crypto from 'node:crypto';
import type { FastifyInstance } from 'fastify';
import { z } from 'zod';
import {
  and,
  createDb,
  eq,
  gt,
  isNull,
  mcpServers,
  oauthClients,
  oauthCodes,
  oauthTokens,
} from '@bmm/db';
import { getJWKS, signAccessToken } from '../lib/jwks.js';
import { checkDailyLimit } from '../lib/rate-limit.js';
import { requireAuth } from '../plugins/session.js';
import { config } from '../config.js';

const db = createDb();

// Access-token lifetime is short so revocation propagates within the hour.
// Refresh-token lifetime is long so legitimate clients don't have to
// re-authorize daily; rotation on every refresh limits exposure if one leaks.
const ACCESS_TOKEN_TTL_S = 3600; // 1 hour
const REFRESH_TOKEN_TTL_MS = 30 * 24 * 3600 * 1000; // 30 days

function sha256(input: string): string {
  return crypto.createHash('sha256').update(input).digest('hex');
}

// OAUTH-002: S256 only. The 'plain' PKCE method is deprecated by OAuth 2.1 and
// our AS metadata already advertises S256 exclusively — accepting plain would be
// a downgrade the metadata says we don't allow.
function pkceVerify(verifier: string, challenge: string, method: string): boolean {
  if (method !== 'S256') return false;
  const computed = crypto.createHash('sha256').update(verifier).digest('base64url');
  return computed === challenge;
}

async function resolveServerByResource(resource: string) {
  const url = new URL(resource);
  const port = url.port ? Number(url.port) : null;
  if (port !== null) {
    const [s] = await db.select().from(mcpServers).where(eq(mcpServers.hostPort, port)).limit(1);
    if (s) return s;
  }
  const slug = url.hostname.split('.')[0];
  if (slug) {
    const [s] = await db.select().from(mcpServers).where(eq(mcpServers.slug, slug)).limit(1);
    if (s) return s;
  }
  return null;
}

export async function oauthRoutes(app: FastifyInstance): Promise<void> {
  // Authorization Server Metadata (RFC 8414) — control-plane wide
  app.get('/oauth/.well-known/oauth-authorization-server', async (_req, reply) => {
    const base = `${config.CONTROL_PLANE_PUBLIC_URL}`;
    return reply.send({
      issuer: `${base}/oauth`,
      authorization_endpoint: `${base}/oauth/authorize`,
      token_endpoint: `${base}/oauth/token`,
      registration_endpoint: `${base}/oauth/register`,
      jwks_uri: `${base}/oauth/jwks`,
      response_types_supported: ['code'],
      grant_types_supported: ['authorization_code', 'refresh_token'],
      code_challenge_methods_supported: ['S256'],
      token_endpoint_auth_methods_supported: [
        'client_secret_basic',
        'client_secret_post',
        'none',
      ],
      scopes_supported: ['mcp:read', 'mcp:write'],
    });
  });

  app.get('/oauth/jwks', async (_req, reply) => {
    reply.header('cache-control', 'public, max-age=300');
    return reply.send(await getJWKS());
  });

  // RFC 7591 Dynamic Client Registration — rate-limited per-IP to prevent
  // DB-row spam. 20/day per visitor IP is well above legitimate MCP-client
  // bootstrap rates (each AI client registers once per resource server, ever).
  // (Zb-002.)
  app.post('/oauth/register', async (req, reply) => {
    const rl = await checkDailyLimit('oauth_register', req.ip, 20);
    if (!rl.ok) {
      return reply.code(429).send({
        error: 'rate_limited',
        detail: 'Too many client registrations from this IP. Try again tomorrow.',
      });
    }
    const Body = z.object({
      client_name: z.string().min(1).max(128).optional(),
      redirect_uris: z.array(z.string().url()).min(1).max(10),
      grant_types: z.array(z.string()).optional(),
      response_types: z.array(z.string()).optional(),
      token_endpoint_auth_method: z.string().optional(),
      resource: z.string().url().optional(),
    });
    const parsed = Body.safeParse(req.body);
    if (!parsed.success) return reply.code(400).send({ error: 'invalid_request' });

    let serverId: string | null = null;
    if (parsed.data.resource) {
      const server = await resolveServerByResource(parsed.data.resource);
      if (!server) return reply.code(400).send({ error: 'invalid_resource' });
      serverId = server.id;
    } else {
      return reply.code(400).send({ error: 'resource_required' });
    }

    const clientId = `bmm_${crypto.randomBytes(12).toString('hex')}`;
    const isPublic = parsed.data.token_endpoint_auth_method === 'none';
    let clientSecret: string | null = null;
    let clientSecretHash: string | null = null;
    if (!isPublic) {
      clientSecret = crypto.randomBytes(32).toString('base64url');
      clientSecretHash = sha256(clientSecret);
    }

    await db.insert(oauthClients).values({
      serverId,
      clientId,
      clientSecretHash,
      redirectUris: parsed.data.redirect_uris,
      metadata: { ...parsed.data },
    });

    return reply.code(201).send({
      client_id: clientId,
      ...(clientSecret ? { client_secret: clientSecret } : {}),
      redirect_uris: parsed.data.redirect_uris,
      grant_types: parsed.data.grant_types ?? ['authorization_code', 'refresh_token'],
      response_types: ['code'],
      token_endpoint_auth_method: isPublic ? 'none' : 'client_secret_basic',
    });
  });

  // /oauth/authorize — user-consent step. Requires a logged-in dashboard user.
  app.get('/oauth/authorize', { preHandler: requireAuth }, async (req, reply) => {
    const user = req.user!;
    const Query = z.object({
      response_type: z.literal('code'),
      client_id: z.string(),
      redirect_uri: z.string().url(),
      code_challenge: z.string(),
      code_challenge_method: z.literal('S256').default('S256'),
      state: z.string().optional(),
      scope: z.string().optional(),
      resource: z.string().url(),
    });
    const parsed = Query.safeParse(req.query);
    if (!parsed.success) return reply.code(400).send({ error: 'invalid_request' });

    const [client] = await db
      .select()
      .from(oauthClients)
      .where(eq(oauthClients.clientId, parsed.data.client_id))
      .limit(1);
    if (!client) return reply.code(400).send({ error: 'unknown_client' });

    const redirectOk = (client.redirectUris as string[]).includes(parsed.data.redirect_uri);
    if (!redirectOk) return reply.code(400).send({ error: 'invalid_redirect_uri' });

    const server = await resolveServerByResource(parsed.data.resource);
    if (!server || server.id !== client.serverId) {
      return reply.code(400).send({ error: 'invalid_resource' });
    }
    if (server.orgId !== user.orgId) {
      return reply.code(403).send({ error: 'forbidden_resource' });
    }

    const code = crypto.randomBytes(24).toString('base64url');
    await db.insert(oauthCodes).values({
      clientDbId: client.id,
      code,
      codeChallenge: parsed.data.code_challenge,
      codeChallengeMethod: parsed.data.code_challenge_method,
      redirectUri: parsed.data.redirect_uri,
      scope: parsed.data.scope,
      resource: parsed.data.resource,
      userId: user.userId,
      expiresAt: new Date(Date.now() + 5 * 60_000),
    });

    const url = new URL(parsed.data.redirect_uri);
    url.searchParams.set('code', code);
    if (parsed.data.state) url.searchParams.set('state', parsed.data.state);
    return reply.redirect(url.toString());
  });

  app.post('/oauth/token', async (req, reply) => {
    const Body = z.object({
      grant_type: z.enum(['authorization_code', 'refresh_token']),
      code: z.string().optional(),
      redirect_uri: z.string().url().optional(),
      client_id: z.string().optional(),
      client_secret: z.string().optional(),
      code_verifier: z.string().optional(),
      resource: z.string().url().optional(),
      refresh_token: z.string().optional(),
    });
    const body = (req.body ?? {}) as Record<string, string>;
    const parsed = Body.safeParse(body);
    if (!parsed.success) return reply.code(400).send({ error: 'invalid_request' });

    if (parsed.data.grant_type === 'authorization_code') {
      const { code, code_verifier, client_id, client_secret, redirect_uri, resource } = parsed.data;
      if (!code || !code_verifier || !client_id || !redirect_uri || !resource) {
        return reply.code(400).send({ error: 'invalid_request' });
      }
      const [row] = await db
        .select({ code: oauthCodes, client: oauthClients })
        .from(oauthCodes)
        .innerJoin(oauthClients, eq(oauthClients.id, oauthCodes.clientDbId))
        .where(and(eq(oauthCodes.code, code), gt(oauthCodes.expiresAt, new Date())))
        .limit(1);
      if (!row || row.code.consumedAt) return reply.code(400).send({ error: 'invalid_grant' });
      if (row.client.clientId !== client_id) return reply.code(400).send({ error: 'invalid_client' });
      if (row.code.redirectUri !== redirect_uri) return reply.code(400).send({ error: 'invalid_redirect_uri' });
      if (row.code.resource !== resource) return reply.code(400).send({ error: 'invalid_resource' });
      if (!pkceVerify(code_verifier, row.code.codeChallenge, row.code.codeChallengeMethod)) {
        return reply.code(400).send({ error: 'invalid_grant' });
      }
      if (row.client.clientSecretHash) {
        if (!client_secret || sha256(client_secret) !== row.client.clientSecretHash) {
          return reply.code(401).send({ error: 'invalid_client' });
        }
      }
      // OAUTH-001: consume the code atomically. The UPDATE only succeeds while
      // consumed_at is still NULL, so two requests racing the same code can
      // never both mint a token — the loser gets zero rows back.
      const consumed = await db
        .update(oauthCodes)
        .set({ consumedAt: new Date() })
        .where(and(eq(oauthCodes.id, row.code.id), isNull(oauthCodes.consumedAt)))
        .returning({ id: oauthCodes.id });
      if (consumed.length === 0) {
        return reply.code(400).send({ error: 'invalid_grant' });
      }

      const subject = row.code.userId ?? row.client.clientId;
      const accessToken = await signAccessToken({
        subject,
        audience: resource,
        issuer: `${config.CONTROL_PLANE_PUBLIC_URL}/oauth`,
        scope: row.code.scope ?? '',
        ttlSeconds: ACCESS_TOKEN_TTL_S,
      });
      const refreshToken = crypto.randomBytes(32).toString('base64url');
      await db.insert(oauthTokens).values({
        clientDbId: row.client.id,
        accessTokenHash: sha256(accessToken),
        refreshTokenHash: sha256(refreshToken),
        scope: row.code.scope ?? null,
        resource,
        subject,
        // expiresAt is the REFRESH-token lifetime — 30 days. Access-token
        // expiry lives inside the JWT's `exp` claim (1h, set above).
        expiresAt: new Date(Date.now() + REFRESH_TOKEN_TTL_MS),
      });

      return reply.send({
        access_token: accessToken,
        token_type: 'Bearer',
        expires_in: ACCESS_TOKEN_TTL_S,
        refresh_token: refreshToken,
        scope: row.code.scope ?? '',
      });
    }

    // ─── grant_type: refresh_token ─────────────────────────────────────
    // OAuth 2.1 with rotation: every successful refresh issues a NEW refresh
    // token and atomically invalidates the old one. If a stolen refresh token
    // gets used after the legitimate client refreshed, the second use sees
    // invalid_grant — that's how rotation surfaces token theft.
    if (parsed.data.grant_type === 'refresh_token') {
      const { refresh_token, client_id, client_secret, resource: requestedResource } =
        parsed.data;
      if (!refresh_token || !client_id) {
        return reply.code(400).send({ error: 'invalid_request' });
      }

      const refreshHash = sha256(refresh_token);
      const [row] = await db
        .select({ token: oauthTokens, client: oauthClients })
        .from(oauthTokens)
        .innerJoin(oauthClients, eq(oauthClients.id, oauthTokens.clientDbId))
        .where(
          and(
            eq(oauthTokens.refreshTokenHash, refreshHash),
            gt(oauthTokens.expiresAt, new Date()),
          ),
        )
        .limit(1);
      if (!row) return reply.code(400).send({ error: 'invalid_grant' });
      if (row.client.clientId !== client_id) {
        return reply.code(401).send({ error: 'invalid_client' });
      }
      if (row.client.clientSecretHash) {
        if (!client_secret || sha256(client_secret) !== row.client.clientSecretHash) {
          return reply.code(401).send({ error: 'invalid_client' });
        }
      }
      // RFC 8707: requested resource must equal the stored one — refreshes
      // don't allow audience changes (would be a downgrade/escalation vector).
      if (requestedResource && requestedResource !== row.token.resource) {
        return reply.code(400).send({ error: 'invalid_resource' });
      }

      const subject = row.token.subject ?? row.client.clientId;
      const newAccessToken = await signAccessToken({
        subject,
        audience: row.token.resource ?? '',
        issuer: `${config.CONTROL_PLANE_PUBLIC_URL}/oauth`,
        scope: row.token.scope ?? '',
        ttlSeconds: ACCESS_TOKEN_TTL_S,
      });
      const newRefreshToken = crypto.randomBytes(32).toString('base64url');
      const newRefreshHash = sha256(newRefreshToken);

      // Atomic rotation: UPDATE only succeeds if the row still has the OLD
      // refresh-hash. Two parallel refreshes with the same token can't both
      // win — the loser sees zero rows and gets invalid_grant.
      const rotated = await db
        .update(oauthTokens)
        .set({
          accessTokenHash: sha256(newAccessToken),
          refreshTokenHash: newRefreshHash,
          expiresAt: new Date(Date.now() + REFRESH_TOKEN_TTL_MS),
        })
        .where(
          and(eq(oauthTokens.id, row.token.id), eq(oauthTokens.refreshTokenHash, refreshHash)),
        )
        .returning({ id: oauthTokens.id });
      if (rotated.length === 0) {
        return reply.code(400).send({ error: 'invalid_grant' });
      }

      return reply.send({
        access_token: newAccessToken,
        token_type: 'Bearer',
        expires_in: ACCESS_TOKEN_TTL_S,
        refresh_token: newRefreshToken,
        scope: row.token.scope ?? '',
      });
    }

    return reply.code(400).send({ error: 'unsupported_grant_type' });
  });

  // Resource discovery proxy — generated servers ask the AS for their own metadata
  // but during dev we expose a control-plane endpoint so the dashboard can show URLs.
  app.get('/oauth/resource-metadata', async (req, reply) => {
    const Query = z.object({ resource: z.string().url() });
    const parsed = Query.safeParse(req.query);
    if (!parsed.success) return reply.code(400).send({ error: 'invalid_request' });
    const server = await resolveServerByResource(parsed.data.resource);
    if (!server) return reply.code(404).send({ error: 'not_found' });
    const base = config.CONTROL_PLANE_PUBLIC_URL;
    return reply.send({
      resource: parsed.data.resource,
      authorization_servers: [`${base}/oauth`],
      bearer_methods_supported: ['header'],
      scopes_supported: ['mcp:read', 'mcp:write'],
    });
  });
}

