import { createMiddleware } from "hono/factory";
import { createRemoteJWKSet, jwtVerify, type JWTPayload } from "jose";
import type { AuthUser } from "./upsert-user.js";
import { IdentityConflictError, upsertUserFromIdentity } from "./upsert-user.js";
import type { ApiEnv, UserRole } from "../env.js";
import { isUserRole } from "../env.js";
import type { Db } from "../db/client.js";

export type AppVariables = {
  user: AuthUser;
};

export type AppBindings = {
  Variables: AppVariables;
};

function roleFromClaims(claims: Record<string, unknown>, fallback: UserRole): UserRole {
  const raw =
    (typeof claims["custom:role"] === "string" && claims["custom:role"]) ||
    (typeof claims.role === "string" && claims.role) ||
    fallback;
  return isUserRole(raw) ? raw : fallback;
}

function audienceMatches(payload: JWTPayload, expected: string): boolean {
  const claims = payload as JWTPayload & { token_use?: string; client_id?: string };
  if (claims.token_use === "access") {
    return claims.client_id === expected;
  }

  if (typeof payload.aud === "string") {
    return payload.aud === expected;
  }
  if (Array.isArray(payload.aud)) {
    return payload.aud.includes(expected);
  }
  return false;
}

function identityFromPayload(payload: JWTPayload): {
  sub: string;
  email: string;
  name: string;
} | null {
  const sub = typeof payload.sub === "string" ? payload.sub : null;
  const email = typeof payload.email === "string" ? payload.email : null;
  const name =
    (typeof payload.name === "string" && payload.name) ||
    (typeof payload["cognito:username"] === "string" && payload["cognito:username"]) ||
    email;

  if (!sub || !email || !name) {
    return null;
  }
  return { sub, email, name };
}

export function createAuthMiddleware(env: ApiEnv, handle: Db) {
  const jwks =
    env.cognitoIssuer.length > 0
      ? createRemoteJWKSet(new URL(`${env.cognitoIssuer}/.well-known/jwks.json`))
      : null;

  return createMiddleware<AppBindings>(async (c, next) => {
    if (env.devAuthBypass) {
      try {
        const user = await upsertUserFromIdentity(handle, {
          cognitoSub: env.devAuthSub,
          email: env.devAuthEmail,
          name: env.devAuthName,
          role: env.devAuthRole,
        });
        c.set("user", user);
      } catch (error) {
        if (error instanceof IdentityConflictError) {
          return c.json({ error: error.message }, 409);
        }
        throw error;
      }
      await next();
      return;
    }

    const header = c.req.header("authorization");
    if (!header?.startsWith("Bearer ")) {
      return c.json({ error: "Missing or invalid Authorization header." }, 401);
    }

    if (!jwks) {
      return c.json({ error: "JWT verification is not configured." }, 401);
    }

    const token = header.slice("Bearer ".length);
    let payload: JWTPayload;
    try {
      ({ payload } = await jwtVerify(token, jwks, {
        issuer: env.cognitoIssuer,
      }));
    } catch {
      return c.json({ error: "Invalid or expired token." }, 401);
    }

    if (!audienceMatches(payload, env.cognitoAudience)) {
      return c.json({ error: "Token audience does not match this API." }, 401);
    }

    const identity = identityFromPayload(payload);
    if (!identity) {
      return c.json(
        {
          error:
            "Token is missing required identity claims. Use a Cognito ID token or an access token that includes email.",
        },
        401,
      );
    }

    let user: AuthUser;
    try {
      user = await upsertUserFromIdentity(handle, {
        cognitoSub: identity.sub,
        email: identity.email,
        name: identity.name,
        role: roleFromClaims(payload as Record<string, unknown>, "viewer"),
      });
    } catch (error) {
      if (error instanceof IdentityConflictError) {
        return c.json({ error: error.message }, 409);
      }
      throw error;
    }

    c.set("user", user);
    await next();
  });
}