import { Data, Effect, Result } from "effect";
import { jwtVerify, type JWTVerifyGetKey } from "jose";
import { JOSEError, JWKSInvalid, JWKSTimeout, JWTExpired } from "jose/errors";

export type VerifiedToken = {
  /** The WorkOS account ID (user ID). */
  accountId: string;
  /** The WorkOS organization ID, if the session has org context. */
  organizationId: string | null;
};

export class McpJwtVerificationError extends Data.TaggedError("McpJwtVerificationError")<{
  readonly cause: unknown;
  readonly reason: "expired" | "invalid" | "system";
}> {}

const classifyJwtVerificationError = (cause: unknown): McpJwtVerificationError =>
  new McpJwtVerificationError({
    cause,
    reason:
      cause instanceof JWTExpired
        ? "expired"
        : cause instanceof JWKSTimeout ||
            cause instanceof JWKSInvalid ||
            !(cause instanceof JOSEError)
          ? "system"
          : "invalid",
  });

const isExpectedJwtVerificationError = (error: McpJwtVerificationError): boolean =>
  error.reason === "expired" || error.reason === "invalid";

const withJwtVerificationSpan = <A>(
  effect: Effect.Effect<A, McpJwtVerificationError>,
): Effect.Effect<A, McpJwtVerificationError> =>
  effect.pipe(
    Effect.result,
    Effect.flatMap((outcome) =>
      Effect.gen(function* () {
        if (Result.isSuccess(outcome)) {
          yield* Effect.annotateCurrentSpan({ "mcp.auth.jwt_verify.outcome": "verified" });
          return outcome;
        }

        yield* Effect.annotateCurrentSpan({
          "mcp.auth.jwt_verify.outcome": outcome.failure.reason,
        });

        return isExpectedJwtVerificationError(outcome.failure)
          ? outcome
          : yield* Effect.fail(outcome.failure);
      }),
    ),
    Effect.withSpan("mcp.auth.jwt_verify"),
    Effect.flatMap((outcome) =>
      Result.isSuccess(outcome) ? Effect.succeed(outcome.success) : Effect.fail(outcome.failure),
    ),
  );

export const verifyMcpAccessToken = (
  token: string,
  jwks: JWTVerifyGetKey,
  options: {
    readonly issuer: string;
    readonly audience: string;
  },
) =>
  Effect.gen(function* () {
    const { payload } = yield* Effect.tryPromise({
      try: () =>
        jwtVerify(token, jwks, {
          issuer: options.issuer,
          audience: options.audience,
        }),
      catch: classifyJwtVerificationError,
    }).pipe(withJwtVerificationSpan);

    if (!payload.sub) return null;

    return {
      accountId: payload.sub,
      organizationId: (payload.org_id as string | undefined) ?? null,
    } satisfies VerifiedToken;
  });

export const verifyWorkOSMcpAccessToken = (
  token: string,
  jwks: JWTVerifyGetKey,
  options: {
    readonly issuer: string;
    readonly audience: string;
  },
) =>
  Effect.gen(function* () {
    const verified = yield* verifyMcpAccessToken(token, jwks, {
      issuer: options.issuer,
      audience: options.audience,
    });
    yield* Effect.annotateCurrentSpan({
      "mcp.auth.audience_mode": "workos_client",
    });
    return verified;
  });