import {
  importPKCS8,
  SignJWT,
  decodeJwt,
  createRemoteJWKSet,
  jwtVerify,
} from "jose";
import { BaseProvider, OAuthTokens } from "../../types";
import {
  ConfigError,
  ProviderGetUserError,
  TokenError,
} from "../../utils/errors";
import { parseQuerystring } from "../../utils/helpers";
import { logger } from "../../utils/logger";
import { Apple } from "./types";

function replaceEscapeCharacters(input: string): string {
  return input.replace(/\\n/g, "\n");
}

export async function convertPrivateKeyToClientSecret({
  privateKey,
  keyIdentifier,
  teamId,
  clientId,
  expAfter,
}: Apple.ConvertPrivateKeyToClientSecretOptions): Promise<string> {
  const now = Math.floor(Date.now() / 1000);
  const exp = now + expAfter;

  privateKey = replaceEscapeCharacters(privateKey);

  const payload = {
    iss: teamId,
    iat: now,
    exp: now + expAfter,
    aud: "https://appleid.apple.com",
    sub: clientId,
  };

  const key = await importPKCS8(privateKey, "ES256");

  const clientSecret = await new SignJWT(payload)
    .setProtectedHeader({ alg: "ES256", kid: keyIdentifier, typ: "JWT" })
    .setIssuedAt(now)
    .setExpirationTime(exp)
    .sign(key);

  return clientSecret;
}

export async function getTokensFromCode(
  code: string,
  { clientId, clientSecret, redirectUrl }: BaseProvider.TokensFromCodeOptions
): Promise<OAuthTokens> {
  const params = {
    grant_type: "authorization_code",
    code,
    client_id: clientId,
    client_secret: clientSecret,
    redirect_uri: redirectUrl,
  };

  const response = await fetch("https://appleid.apple.com/auth/token", {
    method: "POST",
    headers: {
      "content-type": "application/x-www-form-urlencoded",
    },
    body: new URLSearchParams(params).toString(),
  });

  const result: any = await response.json();
  logger.log(`[tokens], ${JSON.stringify(result)}`, "info");

  if (result.error) {
    throw new TokenError({
      message: result.error_description,
    });
  }
  return result as OAuthTokens;
}

export async function getUser(token: string): Promise<Apple.UserResponse> {
  try {
    const data = decodeJwt(token) as Apple.UserResponse;
    logger.log(`[provider user data], ${JSON.stringify(data)}`, "info");
    return data;
  } catch (e: any) {
    logger.log(`[error], ${JSON.stringify(e.stack)}`, "error");
    throw new ProviderGetUserError({
      message: "There was an error fetching the user",
    });
  }
}

export async function verifyIdToken(
  idToken: string,
  clientId?: string,
  userData?: Apple.UserData
): Promise<Apple.TokenVerificationResponse> {
  try {
    // Apple's JWKS (JSON Web Key Set) endpoint
    const JWKS = createRemoteJWKSet(
      new URL("https://appleid.apple.com/auth/keys")
    );

    // Prepare verification options
    const verificationOptions = {
      issuer: "https://appleid.apple.com",
      ...(clientId && { audience: clientId }), // Only include audience check if clientId is provided
    };

    // Verify the token's signature and decode its payload
    const { payload } = await jwtVerify(idToken, JWKS, verificationOptions);

    logger.log(`[token verification data], ${JSON.stringify(payload)}`, "info");

    // Validate required fields
    if (!payload.sub || !payload.iss || !payload.exp || !payload.iat) {
      throw new Error("Missing required JWT claims");
    }

    // Additional validations before creating the verified payload
    if (Date.now() >= payload.exp! * 1000) {
      throw new Error("Token has expired");
    }

    if (payload.iss !== "https://appleid.apple.com") {
      throw new Error("Invalid token issuer");
    }

    // Only verify audience if clientId is provided
    if (clientId) {
      const tokenAudience = Array.isArray(payload.aud)
        ? payload.aud
        : [payload.aud];
      if (!tokenAudience.includes(clientId)) {
        throw new Error("Invalid audience");
      }
    }

    // Create verified payload with proper type checking
    const verifiedPayload: Apple.TokenVerificationResponse = {
      // Standard JWT claims
      iss: payload.iss as string,
      sub: payload.sub,
      aud: Array.isArray(payload.aud)
        ? payload.aud[0]
        : (payload.aud as string),
      iat: payload.iat,
      exp: payload.exp,

      // Apple specific claims
      email: payload.email as string,
      email_verified: Boolean(payload.email_verified),
      is_private_email: Boolean(payload.is_private_email),

      // Optional claims
      nonce: payload.nonce as string,
      nonce_supported: Boolean(payload.nonce_supported),

      // Additional user information from initial sign-in if available
      given_name: userData?.name?.given_name,
      family_name: userData?.name?.family_name,
      middle_name: userData?.name?.middle_name,
      real_user_status: userData?.real_user_status,
    };

    return verifiedPayload;
  } catch (e: any) {
    logger.log(`[error], ${JSON.stringify(e.stack)}`, "error");

    // Specific error handling based on error type
    let errorMessage: string;

    switch (e.code) {
      case "ERR_JWS_SIGNATURE_VERIFICATION_FAILED":
        errorMessage = "Invalid token signature";
        break;
      case "ERR_JWT_EXPIRED":
        errorMessage = "Token has expired";
        break;
      case "ERR_JWS_INVALID":
        errorMessage = "Invalid token format";
        break;
      case "ERR_JWS_VERIFICATION_FAILED":
        errorMessage = "Token verification failed";
        break;
      default:
        errorMessage = e.message || "There was an error verifying the ID token";
    }

    throw new ProviderGetUserError({
      message: errorMessage,
    });
  }
}

export default async function callback({
  options,
  request,
}: BaseProvider.CallbackOptions): Promise<Apple.CallbackResponse> {
  const { query }: any = parseQuerystring(request);
  logger.setEnabled(options?.isLogEnabled || false);
  logger.log(`[code], ${JSON.stringify(query.code)}`, "info");
  if (!query.code) {
    throw new ConfigError({
      message: "No code is passed!",
    });
  }
  const tokens = await getTokensFromCode(query.code, options);
  const accessToken = tokens.access_token;
  logger.log(`[access_token], ${JSON.stringify(accessToken)}`, "info");
  const providerUser = await getUser(accessToken);
  return {
    user: providerUser,
    tokens,
  };
}