import { validateJWT, type ValidationResult } from "@trigger.dev/core/v3/jwt";
import { $replica } from "~/db.server";
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
import type { AuthenticatedEnvironment } from "../apiAuth.server";

export type ValidatePublicJwtKeySuccess = {
  ok: true;
  environment: AuthenticatedEnvironment;
  claims: Record<string, unknown>;
};

export type ValidatePublicJwtKeyError = {
  ok: false;
  error: string;
};

export type ValidatePublicJwtKeyResult = ValidatePublicJwtKeySuccess | ValidatePublicJwtKeyError;

export async function validatePublicJwtKey(token: string): Promise<ValidatePublicJwtKeyResult> {
  // Get the sub claim from the token
  // Use the sub claim to find the environment
  // Validate the token against the environment.apiKey
  // Once that's done, return the environment and the claims
  const sub = extractJWTSub(token);

  if (!sub) {
    return { ok: false, error: "Invalid Public Access Token, missing subject." };
  }

  const environment = await findEnvironmentById(sub);

  if (!environment) {
    return { ok: false, error: "Invalid Public Access Token, environment not found." };
  }

  let result = await validateJWT(
    token,
    environment.parentEnvironment?.apiKey ?? environment.apiKey
  );

  // PATs are signed with the env's apiKey at mint time. If the env's apiKey
  // has since been rotated, signature verification fails against the current
  // key — fall back to any RevokedApiKey rows still in their grace window.
  // Only run this query on the failure path so the success path is unchanged.
  if (!result.ok) {
    result = await validateAgainstRevokedApiKeys(
      token,
      environment.parentEnvironment?.id ?? environment.id,
      result
    );
  }

  if (!result.ok) {
    switch (result.code) {
      case "ERR_JWT_EXPIRED": {
        return {
          ok: false,
          error:
            "Public Access Token has expired. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
        };
      }
      case "ERR_JWT_CLAIM_INVALID": {
        return {
          ok: false,
          error: `Public Access Token is invalid: ${result.error}. See https://trigger.dev/docs/frontend/overview#authentication for more information.`,
        };
      }
      default: {
        return {
          ok: false,
          error:
            "Public Access Token is invalid. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
        };
      }
    }
  }

  return {
    ok: true,
    environment,
    claims: result.payload,
  };
}

async function validateAgainstRevokedApiKeys(
  token: string,
  signingEnvironmentId: string,
  primaryResult: ValidationResult
): Promise<ValidationResult> {
  const revokedApiKeys = await $replica.revokedApiKey.findMany({
    where: {
      runtimeEnvironmentId: signingEnvironmentId,
      expiresAt: { gt: new Date() },
    },
    select: { apiKey: true },
  });

  for (const { apiKey } of revokedApiKeys) {
    const fallbackResult = await validateJWT(token, apiKey);
    if (fallbackResult.ok) {
      return fallbackResult;
    }
  }

  return primaryResult;
}

export function isPublicJWT(token: string): boolean {
  // Split the token
  const parts = token.split(".");
  if (parts.length !== 3) return false;

  try {
    // Decode the payload (second part)
    const payload = JSON.parse(decodeBase64Url(parts[1]));

    if (payload === null || typeof payload !== "object") return false;

    // Check for the pub: true claim
    return "pub" in payload && payload.pub === true;
  } catch (_error) {
    // If there's any error in decoding or parsing, it's not a valid JWT
    return false;
  }
}

export function extractJwtSigningSecretKey(environment: AuthenticatedEnvironment) {
  return environment.parentEnvironment?.apiKey ?? environment.apiKey;
}

function extractJWTSub(token: string): string | undefined {
  // Split the token
  const parts = token.split(".");
  if (parts.length !== 3) return;

  try {
    // Decode the payload (second part)
    const payload = JSON.parse(decodeBase64Url(parts[1]));

    if (payload === null || typeof payload !== "object") return;

    // Check for the pub: true claim
    return "sub" in payload && typeof payload.sub === "string" ? payload.sub : undefined;
  } catch (_error) {
    // If there's any error in decoding or parsing, it's not a valid JWT
    return;
  }
}

function decodeBase64Url(str: string): string {
  // Replace URL-safe characters and add padding
  str = str.replace(/-/g, "+").replace(/_/g, "/");
  switch (str.length % 4) {
    case 2:
      str += "==";
      break;
    case 3:
      str += "=";
      break;
  }

  // Decode using Node.js Buffer
  return Buffer.from(str, "base64").toString("utf8");
}