/**
 * Authentication and API key resolution for the LLM proxy handler.
 *
 * Extracted from handleLLMProxy to keep the main handler focused on
 * request/response orchestration. Each function is independently testable.
 */

import {
  hasArchestraTokenPrefix,
  isSupportedProvider,
  LLM_PROXY_OAUTH_SCOPE,
  providerRequiresPerUserCredential,
} from "@archestra/shared";
import type { FastifyRequest } from "fastify";
import { userHasPermission } from "@/auth";
import { type AllowedCacheKey, CacheKey, cacheManager } from "@/cache-manager";
import logger from "@/logging";
import {
  AgentModel,
  AgentTeamModel,
  LlmOauthClientModel,
  LlmProviderApiKeyModel,
  MemberModel,
  OAuthAccessTokenModel,
  OAuthClientModel,
  VirtualApiKeyModel,
} from "@/models";
import { validateExternalIdpToken } from "@/routes/mcp-gateway.utils";
import { getSecretValueForLlmProviderApiKey } from "@/secrets-manager";
import { isAppConnectorAudienceRef } from "@/services/apps/app-connector-resource";
import { type Agent, ApiError, type ResourceVisibilityScope } from "@/types";
import { resolveProviderApiKey } from "@/utils/llm-api-key-resolution";
import { isLoopbackAddress } from "@/utils/network";

// =========================================================================
// Agent Resolution
// =========================================================================

/**
 * Resolve the target agent from the request URL or fall back to the default profile.
 */
export async function resolveAgent(
  agentId: string | undefined,
): Promise<Agent> {
  if (agentId) {
    const agent = await AgentModel.findById(agentId);
    if (!agent) {
      throw new ApiError(404, `Agent with ID ${agentId} not found`);
    }
    return agent;
  }

  const defaultProfile = await AgentModel.getDefaultProfile();
  if (!defaultProfile) {
    throw new ApiError(400, "Please specify an LLMProxy ID in the URL path.");
  }
  return defaultProfile;
}

// =========================================================================
// Virtual API Key Validation
// =========================================================================

export interface VirtualKeyValidationResult {
  apiKey?: string;
  baseUrl?: string;
  /** Parent chat_api_key row ID; used by the proxy to look up per-key settings (e.g. extra headers). */
  chatApiKeyId?: string;
  virtualKeyId?: string;
  /** Scope of the resolved key; a personal key identifies its owner. */
  virtualKeyScope?: ResourceVisibilityScope;
  /** Owner of the resolved key (for cross-credential user-consistency checks). */
  virtualKeyAuthorId?: string | null;
}

export interface PassthroughVirtualKeyResult {
  /** Owner of the passthrough key — the acting Archestra user. */
  userId: string;
  passthroughVirtualKeyId: string;
}

type ResolvedVirtualApiKey = NonNullable<
  Awaited<ReturnType<typeof VirtualApiKeyModel.validateToken>>
>;

export async function validateVirtualApiKeyToken(
  tokenValue: string,
): Promise<ResolvedVirtualApiKey> {
  const resolved = await VirtualApiKeyModel.validateToken(tokenValue);
  if (!resolved) {
    throw new ApiError(401, "Invalid virtual API key");
  }

  if (
    resolved.virtualKey.expiresAt &&
    resolved.virtualKey.expiresAt < new Date()
  ) {
    throw new ApiError(401, "Virtual API key expired");
  }

  return resolved;
}

/**
 * Validate a platform-managed virtual API key.
 * Checks: token validity, expiration, and provider mapping.
 * Returns the resolved real API key and optional base URL.
 *
 * Throws ApiError on validation failure.
 */
export async function validateVirtualApiKey(
  tokenValue: string,
  expectedProvider: string,
): Promise<VirtualKeyValidationResult> {
  const resolved = await validateVirtualApiKeyToken(tokenValue);
  if (resolved.virtualKey.keyType === "passthrough") {
    throw new ApiError(
      400,
      "Passthrough virtual keys carry no provider credential — send them in the X-Archestra-Virtual-Key header, not Authorization.",
    );
  }
  const mappedProviderKey = (
    await VirtualApiKeyModel.getProviderApiKeysForRouting(
      resolved.virtualKey.id,
    )
  ).find((mapping) => mapping.provider === expectedProvider);
  if (!mappedProviderKey) {
    throw new ApiError(
      400,
      `Virtual API key is not mapped to provider "${expectedProvider}".`,
    );
  }

  // Per-user providers (GitHub Copilot) hold an individual's token, so it may
  // only be served through the owner's OWN personal virtual key mapping to
  // their OWN personal provider key. Re-checked here at runtime (not just at
  // create/update) so a virtual key mapped before this rule existed, or one
  // whose scope/mapping changed, can never hand the token to another user.
  if (
    isSupportedProvider(expectedProvider) &&
    providerRequiresPerUserCredential(expectedProvider)
  ) {
    const parentKey = await LlmProviderApiKeyModel.findById(
      mappedProviderKey.providerApiKeyId,
    );
    if (
      resolved.virtualKey.scope !== "personal" ||
      !parentKey ||
      parentKey.scope !== "personal" ||
      parentKey.userId == null ||
      parentKey.userId !== resolved.virtualKey.authorId
    ) {
      throw new ApiError(
        403,
        `${expectedProvider} is per-user: it can only be used through your own personal virtual key linked to your own ${expectedProvider} account.`,
      );
    }
  }

  // Resolve the real provider API key from the secret.
  // If the parent key's secret was removed (orphaned row), apiKey will be
  // undefined. For providers that require keys, the upstream call will fail
  // with a clear error. For keyless providers the virtual key alone is
  // sufficient authentication.
  let apiKey: string | undefined;
  if (mappedProviderKey.secretId) {
    const secretValue = await getSecretValueForLlmProviderApiKey(
      mappedProviderKey.secretId,
    );
    if (secretValue) {
      apiKey = secretValue as string;
    } else {
      logger.warn(
        {
          virtualKeyId: resolved.virtualKey.id,
          chatApiKeyId: mappedProviderKey.providerApiKeyId,
          secretId: mappedProviderKey.secretId,
        },
        "Virtual key's parent chat API key secret could not be resolved (may be orphaned)",
      );
    }
  }

  return {
    apiKey,
    baseUrl: mappedProviderKey.baseUrl ?? undefined,
    chatApiKeyId: mappedProviderKey.providerApiKeyId,
    virtualKeyId: resolved.virtualKey.id,
    virtualKeyScope: resolved.virtualKey.scope,
    virtualKeyAuthorId: resolved.virtualKey.authorId,
  };
}

// =========================================================================
// Passthrough Virtual Key Validation
// =========================================================================

/**
 * Validate a passthrough virtual key from the X-Archestra-Virtual-Key header.
 *
 * A passthrough key carries no provider credential. It authenticates the acting
 * Archestra user; access to the target LLM proxy is then governed by the owner's
 * own agent access (the same access any authenticated user would have).
 *
 * Throws ApiError on validation failure (401 invalid/expired, 400 wrong key
 * type, 403 no proxy access).
 */
export async function validatePassthroughVirtualKey(params: {
  tokenValue: string;
  agent: Agent;
}): Promise<PassthroughVirtualKeyResult> {
  const { tokenValue, agent } = params;

  const resolved = await VirtualApiKeyModel.validateToken(tokenValue);
  if (!resolved) {
    throw new ApiError(401, "Invalid passthrough virtual key");
  }
  const { virtualKey } = resolved;

  if (virtualKey.keyType !== "passthrough") {
    throw new ApiError(
      400,
      "This is a standard virtual key. Send it in the Authorization header instead of X-Archestra-Virtual-Key.",
    );
  }
  if (virtualKey.expiresAt && virtualKey.expiresAt < new Date()) {
    throw new ApiError(401, "Passthrough virtual key expired");
  }
  if (!virtualKey.authorId) {
    throw new ApiError(401, "Passthrough virtual key has no owner");
  }

  const noAccessError = new ApiError(
    403,
    "Your passthrough virtual key does not grant access to this LLM proxy. Contact your administrator.",
  );
  if (virtualKey.organizationId !== agent.organizationId) {
    throw noAccessError;
  }

  // Proxy access follows the owner's own agent access.
  const ownerIsAgentAdmin = await userHasPermission(
    virtualKey.authorId,
    agent.organizationId,
    "agent",
    "admin",
  );
  const hasProxyAccess = await AgentTeamModel.userHasAgentAccess(
    virtualKey.authorId,
    agent.id,
    ownerIsAgentAdmin,
  );
  if (!hasProxyAccess) {
    throw noAccessError;
  }

  return {
    userId: virtualKey.authorId,
    passthroughVirtualKeyId: virtualKey.id,
  };
}

/**
 * All authenticated user-scoped credentials on a request must resolve to a
 * single Archestra user. Throws 401 when two differ. Unauthenticated hints
 * (e.g. the X-Archestra-User-Id header) must NOT be passed in here.
 */
export function assertConsistentUserCredentials(
  userIds: Array<string | null | undefined>,
): void {
  const distinct = new Set(userIds.filter((id): id is string => Boolean(id)));
  if (distinct.size > 1) {
    throw new ApiError(
      401,
      "Conflicting Archestra user credentials: the request's credentials identify different users.",
    );
  }
}

// =========================================================================
// LLM OAuth Access Token Validation
// =========================================================================

export type LlmOAuthAccessTokenValidationResult = {
  apiKey: string | undefined;
  baseUrl: string | undefined;
  chatApiKeyId: string | undefined;
  authMethod: "oauth_client_credentials" | "oauth_user";
  authenticatedApp?: {
    id: string;
    name: string;
    clientId: string;
  };
  userId?: string;
};

export async function validateLlmOAuthAccessToken(params: {
  tokenValue: string;
  expectedProvider: string;
  agent: Agent;
}): Promise<LlmOAuthAccessTokenValidationResult | null> {
  const accessToken = await OAuthAccessTokenModel.getByTokenHash(
    OAuthAccessTokenModel.hashTokenForLookup(params.tokenValue),
  );
  if (!accessToken) {
    return null;
  }
  if (accessToken.expiresAt < new Date()) {
    throw new ApiError(401, "Invalid LLM OAuth access token.");
  }
  if (accessToken.refreshTokenRevoked) {
    throw new ApiError(401, "Invalid LLM OAuth access token.");
  }
  if (isAppConnectorAudienceRef(accessToken.referenceId)) {
    throw new ApiError(403, "Access token is bound to an app connector.");
  }
  if (!hasLlmProxyScope(accessToken.scopes)) {
    throw new ApiError(403, "Access token is missing LLM proxy scope.");
  }
  if (accessToken.userId) {
    return validateUserLlmOAuthAccessToken({
      userId: accessToken.userId,
      clientId: accessToken.clientId,
      expectedProvider: params.expectedProvider,
      agent: params.agent,
    });
  }

  return validateClientCredentialsLlmOAuthAccessToken({
    clientId: accessToken.clientId,
    expectedProvider: params.expectedProvider,
    agent: params.agent,
  });
}

// =========================================================================
// JWKS Authentication
// =========================================================================

export interface JwksAuthResult {
  apiKey: string | undefined;
  baseUrl: string | undefined;
  /** Resolved chat_api_key row ID; used by the proxy to look up per-key settings (e.g. extra headers). */
  chatApiKeyId: string | undefined;
  userId: string | undefined;
  organizationId: string;
}

/**
 * Attempt JWKS authentication for agents with an external identity provider.
 * Returns null if no JWKS auth was attempted (no IdP configured, no bearer token,
 * virtual key token, or the bearer value is not shaped like a JWT).
 * Throws ApiError if the JWT is invalid.
 */
export async function attemptJwksAuth(
  request: FastifyRequest,
  resolvedAgent: Agent,
  providerName: string,
): Promise<JwksAuthResult | null> {
  if (!resolvedAgent.identityProviderId) return null;

  // Read the bearer token from the RAW request headers. We cannot use
  // extractBearerToken(request) here because some provider routes (e.g.
  // OpenAI) define a headers schema with a .transform() that strips the
  // "Bearer " prefix. After Fastify applies the schema transform,
  // request.headers.authorization no longer starts with "Bearer ", causing
  // extractBearerToken to return null and silently skipping JWKS auth.
  // Reading from request.raw.headers bypasses schema transforms.
  const rawAuthHeader = request.raw.headers.authorization;
  const tokenMatch = rawAuthHeader?.match(/^Bearer\s+(.+)$/i);
  const bearerToken = tokenMatch?.[1] ?? null;
  if (!bearerToken || hasArchestraTokenPrefix(bearerToken)) return null;
  if (!isJwtLike(bearerToken)) return null;

  let jwksResult: Awaited<ReturnType<typeof validateExternalIdpToken>>;
  try {
    jwksResult = await validateExternalIdpToken(
      resolvedAgent.id,
      bearerToken,
      "llmProxy",
    );
  } catch (error) {
    // Convert any unexpected validation error to 401 (not 500)
    logger.warn(
      {
        resolvedAgentId: resolvedAgent.id,
        error: error instanceof Error ? error.message : String(error),
      },
      `[${providerName}Proxy] JWKS validation error`,
    );
    throw new ApiError(
      401,
      "JWT validation failed for the configured identity provider.",
    );
  }

  if (!jwksResult) {
    throw new ApiError(
      401,
      "Invalid JWT token for the configured identity provider.",
    );
  }

  logger.info(
    {
      resolvedAgentId: resolvedAgent.id,
      userId: jwksResult.userId,
      identityProviderId: resolvedAgent.identityProviderId,
    },
    `[${providerName}Proxy] JWKS authentication succeeded`,
  );

  let apiKey: string | undefined;
  let baseUrl: string | undefined;
  let chatApiKeyId: string | undefined;

  if (isSupportedProvider(providerName)) {
    const resolved = await resolveProviderApiKey({
      organizationId: jwksResult.organizationId,
      userId: jwksResult.userId,
      provider: providerName,
    });
    apiKey = resolved.apiKey;
    baseUrl = resolved.baseUrl ?? undefined;
    chatApiKeyId = resolved.chatApiKeyId;
  }

  return {
    apiKey,
    baseUrl,
    chatApiKeyId,
    userId: jwksResult.userId,
    organizationId: jwksResult.organizationId,
  };
}

// =========================================================================
// Keyless Provider Check
// =========================================================================

/**
 * For keyless providers (Ollama, vLLM, Vertex AI Gemini), ensure the request
 * was authenticated via a virtual API key or JWKS. Without this, anyone who
 * knows the proxy URL could call the endpoint without credentials.
 *
 * Internal requests from localhost (chat route → proxy) are allowed.
 */
export function assertAuthenticatedForKeylessProvider(
  apiKey: string | undefined,
  wasVirtualKeyResolved: boolean,
  wasJwksAuthenticated: boolean,
  requestIp: string,
): void {
  if (apiKey || wasVirtualKeyResolved || wasJwksAuthenticated) return;

  if (!isLoopbackAddress(requestIp)) {
    throw new ApiError(
      401,
      "Authentication required. Use a platform virtual API key or pass a provider API key.",
    );
  }
}

// =========================================================================
// Virtual Key Rate Limiter
// =========================================================================

const RATE_LIMIT_MAX_FAILURES = 10;
const RATE_LIMIT_WINDOW_MS = 60_000;

interface RateLimitEntry {
  count: number;
}

/**
 * Distributed rate limiter for failed virtual API key validation attempts.
 * Prevents brute-force enumeration of valid tokens by tracking failures per
 * client IP and rejecting further attempts after exceeding the threshold.
 *
 * Uses the PostgreSQL-backed CacheManager (Keyv) so rate limit state is
 * shared across all application pods. Entries expire automatically via TTL.
 */
export class VirtualKeyRateLimiter {
  private cacheManager: {
    get: <T>(key: AllowedCacheKey) => Promise<T | undefined>;
    set: <T>(
      key: AllowedCacheKey,
      value: T,
      ttl?: number,
    ) => Promise<T | undefined>;
  };

  constructor(cacheManager: {
    get: <T>(key: AllowedCacheKey) => Promise<T | undefined>;
    set: <T>(
      key: AllowedCacheKey,
      value: T,
      ttl?: number,
    ) => Promise<T | undefined>;
  }) {
    this.cacheManager = cacheManager;
  }

  async check(ip: string): Promise<void> {
    const entry = await this.cacheManager.get<RateLimitEntry>(this.key(ip));
    if (!entry) return;

    if (entry.count >= RATE_LIMIT_MAX_FAILURES) {
      throw new ApiError(
        429,
        "Too many failed virtual API key attempts. Please try again later.",
      );
    }
  }

  async recordFailure(ip: string): Promise<void> {
    const entry = await this.cacheManager.get<RateLimitEntry>(this.key(ip));
    const newCount = (entry?.count ?? 0) + 1;
    await this.cacheManager.set<RateLimitEntry>(
      this.key(ip),
      { count: newCount },
      RATE_LIMIT_WINDOW_MS,
    );
  }

  private key(ip: string): AllowedCacheKey {
    return `${CacheKey.VirtualKeyRateLimit}-${ip}`;
  }
}

export const virtualKeyRateLimiter = new VirtualKeyRateLimiter(cacheManager);

async function validateClientCredentialsLlmOAuthAccessToken(params: {
  clientId: string;
  expectedProvider: string;
  agent: Agent;
}): Promise<LlmOAuthAccessTokenValidationResult> {
  const oauthClient = await LlmOauthClientModel.findByClientId(params.clientId);
  if (!oauthClient) {
    throw new ApiError(401, "LLM OAuth client is no longer available.");
  }
  if (oauthClient.disabled) {
    throw new ApiError(401, "LLM OAuth client is disabled.");
  }
  if (oauthClient.organizationId !== params.agent.organizationId) {
    throw new ApiError(403, "LLM OAuth client cannot access this LLM Proxy.");
  }
  if (!oauthClient.allowedLlmProxyIds.includes(params.agent.id)) {
    throw new ApiError(403, "LLM OAuth client cannot access this LLM Proxy.");
  }
  const mappedProviderKey = oauthClient.providerApiKeys.find(
    (mapping) => mapping.provider === params.expectedProvider,
  );
  if (!mappedProviderKey) {
    throw new ApiError(
      400,
      `LLM OAuth client is not mapped to provider "${params.expectedProvider}".`,
    );
  }

  // OAuth client credentials are a service-to-service credential with no acting
  // user. Per-user providers (GitHub Copilot) are an individual's token, so
  // they can never be served this way — there's no user to attribute, and the
  // mapped key would be one person's token for every caller.
  if (
    isSupportedProvider(params.expectedProvider) &&
    providerRequiresPerUserCredential(params.expectedProvider)
  ) {
    throw new ApiError(
      400,
      `${params.expectedProvider} is per-user and cannot be used via OAuth client credentials; each user must connect their own account.`,
    );
  }

  const providerApiKey = await LlmProviderApiKeyModel.findById(
    mappedProviderKey.providerApiKeyId,
  );
  if (!providerApiKey) {
    throw new ApiError(
      500,
      "LLM OAuth client references a missing provider API key.",
    );
  }
  return resolveOAuthProviderApiKey({
    chatApiKeyId: providerApiKey.id,
    secretId: providerApiKey.secretId,
    baseUrl: providerApiKey.inferenceBaseUrl ?? providerApiKey.baseUrl,
    actualProvider: providerApiKey.provider,
    expectedProvider: params.expectedProvider,
    authMethod: "oauth_client_credentials",
    authenticatedApp: {
      id: oauthClient.id,
      name: oauthClient.name,
      clientId: oauthClient.clientId,
    },
  });
}

async function validateUserLlmOAuthAccessToken(params: {
  userId: string;
  clientId: string;
  expectedProvider: string;
  agent: Agent;
}): Promise<LlmOAuthAccessTokenValidationResult> {
  const member = await MemberModel.getFirstMembershipForUser(params.userId);
  if (!member || member.organizationId !== params.agent.organizationId) {
    throw new ApiError(401, "OAuth user is no longer available.");
  }

  // Access has two additive sources: the user's own RBAC, or an admin-controlled
  // grant on the authorization_code LLM OAuth client that minted this token — its
  // allowedLlmProxyIds may grant access to proxies the user could not otherwise
  // reach (e.g. a proxy reachable only through a specific pre-registered app).
  const hasAgentAccess = await AgentTeamModel.userHasAgentAccess(
    params.userId,
    params.agent.id,
    false,
  );
  const hasClientGrant = hasAgentAccess
    ? false
    : await llmOauthClientGrantsProxyAccess({
        clientId: params.clientId,
        proxyId: params.agent.id,
        organizationId: member.organizationId,
      });
  if (!hasAgentAccess && !hasClientGrant) {
    throw new ApiError(403, "OAuth user cannot access this LLM Proxy.");
  }
  if (!isSupportedProvider(params.expectedProvider)) {
    throw new ApiError(
      400,
      `OAuth user access is not supported for provider "${params.expectedProvider}".`,
    );
  }

  const resolved = await resolveProviderApiKey({
    organizationId: member.organizationId,
    userId: params.userId,
    provider: params.expectedProvider,
  });
  const oauthClient = await OAuthClientModel.findByClientId(params.clientId);

  return {
    apiKey: resolved.apiKey,
    baseUrl: resolved.baseUrl ?? undefined,
    chatApiKeyId: resolved.chatApiKeyId,
    authMethod: "oauth_user",
    authenticatedApp: oauthClient
      ? {
          id: oauthClient.id,
          name: oauthClient.name ?? oauthClient.clientId,
          clientId: oauthClient.clientId,
        }
      : undefined,
    userId: params.userId,
  };
}

/**
 * Whether the authorization_code LLM OAuth client that minted a user-bound token
 * grants access to an LLM proxy beyond the user's own RBAC. Additive,
 * admin-controlled grant (see the MCP gateway equivalent). Disabled or deleted
 * clients grant nothing.
 */
async function llmOauthClientGrantsProxyAccess(params: {
  clientId: string;
  proxyId: string;
  organizationId: string;
}): Promise<boolean> {
  const oauthClient = await LlmOauthClientModel.findByClientId(params.clientId);
  if (!oauthClient || oauthClient.disabled) {
    return false;
  }
  if (oauthClient.grantType !== "authorization_code") {
    return false;
  }
  if (oauthClient.organizationId !== params.organizationId) {
    return false;
  }
  return oauthClient.allowedLlmProxyIds.includes(params.proxyId);
}

async function resolveOAuthProviderApiKey(params: {
  chatApiKeyId: string;
  secretId: string | null;
  baseUrl: string | null;
  actualProvider: string;
  expectedProvider: string;
  authMethod: "oauth_client_credentials" | "oauth_user";
  authenticatedApp?: {
    id: string;
    name: string;
    clientId: string;
  };
}): Promise<LlmOAuthAccessTokenValidationResult> {
  if (params.actualProvider !== params.expectedProvider) {
    throw new ApiError(
      400,
      `LLM OAuth client provider key is for provider "${params.actualProvider}", but request is for "${params.expectedProvider}"`,
    );
  }

  const apiKey = params.secretId
    ? await getSecretValueForLlmProviderApiKey(params.secretId)
    : undefined;
  return {
    apiKey,
    baseUrl: params.baseUrl ?? undefined,
    chatApiKeyId: params.chatApiKeyId,
    authMethod: params.authMethod,
    authenticatedApp: params.authenticatedApp,
  };
}

function hasLlmProxyScope(scopes: string[] | null | undefined): boolean {
  return scopes?.some((scope) => scope === LLM_PROXY_OAUTH_SCOPE) ?? false;
}

function isJwtLike(token: string): boolean {
  const parts = token.split(".");
  if (parts.length !== 3) return false;
  return parts.every((part) => /^[A-Za-z0-9_-]+$/.test(part));
}
