auth-yes/src/core/session.ts

import type { Context } from "jsr:@hono/hono@4";
import { deleteCookie } from "jsr:@hono/hono@4/cookie";
import { sqlWrapper } from "./db.ts";
import { valkey } from "./valkey.ts";

export interface AuthenticatedUser {
  userId: string;
  sessionId: string;
  username: string;
  label?: string;
  isAgent?: boolean;
  isPaused?: boolean;
  customScopes?: string[];
}

export interface AppRecord {
  id: string;
  name: string;
  domain?: string;
  is_public?: boolean;
  bypass_paths?: string[];
  allowed_cidrs?: string[];
}

/**
 * Extracts all session_id tokens from the Cookie and Authorization headers.
 */
export function extractAllSessionIds(c: Context): string[] {
  const candidates: string[] = [];

  const authHeader = c.req.header("authorization") || "";
  if (authHeader.startsWith("Bearer ")) {
    const bearerToken = authHeader.substring(7).trim();
    if (bearerToken) candidates.push(bearerToken);
  }

  const cookieHeader = c.req.header("cookie") || "";
  if (cookieHeader) {
    const cookieMatches = [
      ...cookieHeader.matchAll(/(?:^|;\s*)session_id=([^;]+)/g),
    ]
      .map((m) => decodeURIComponent(m[1].trim()))
      .filter(Boolean);
    candidates.push(...cookieMatches);
  }

  return candidates;
}

/**
 * Resolves the authenticated user from Valkey cache or PostgreSQL sessions table.
 */
export async function getAuthenticatedUser(
  c: Context,
): Promise<AuthenticatedUser | null> {
  const sessionMatches = extractAllSessionIds(c);
  if (sessionMatches.length === 0) return null;

  for (let i = 0; i < sessionMatches.length; i++) {
    const candidateId = sessionMatches[i];

    // 1. Try Valkey cache
    try {
      const sessionDataStr = await valkey.get(candidateId);
      if (sessionDataStr) {
        const sessionData = JSON.parse(sessionDataStr);
        if (sessionData && sessionData.uuid) {
          if (i > 0) {
            deleteCookie(c, "session_id", { path: "/" });
          }

          return {
            userId: sessionData.uuid,
            sessionId: candidateId,
            username: sessionData.username || "",
            label: sessionData.label,
            isAgent: sessionData.isAgent,
            isPaused: sessionData.is_paused,
            customScopes: sessionData.customScopes,
          };
        }
      }
    } catch (_err) {
      // Valkey cache miss or connection hiccup - fallback to DB
    }

    // 2. Fallback to PostgreSQL sessions table
    try {
      const nowIso = new Date().toISOString();
      const session = await sqlWrapper.sql`
        SELECT s.user_id, s.expires_at, s.label, s.is_agent, s.is_paused, s.custom_scopes, u.username
        FROM sessions s
        JOIN users u ON s.user_id = u.id
        WHERE s.id = ${candidateId} AND s.expires_at > ${nowIso}
      `.then((res: any) => res[0]);

      if (session) {
        const username = session.username || "";
        try {
          const ttlSeconds = Math.max(
            1,
            Math.floor(
              (new Date(session.expires_at).getTime() - Date.now()) / 1000,
            ),
          );
          await valkey.setex(
            candidateId,
            ttlSeconds,
            JSON.stringify({
              uuid: session.user_id,
              username,
              label: session.label,
              isAgent: session.is_agent,
              is_paused: session.is_paused,
              customScopes: session.custom_scopes,
            }),
          );
        } catch (_e) {}

        if (i > 0) {
          deleteCookie(c, "session_id", { path: "/" });
        }

        return {
          userId: session.user_id,
          sessionId: candidateId,
          username,
          label: session.label,
          isAgent: session.is_agent,
          isPaused: session.is_paused,
          customScopes: session.custom_scopes,
        };
      }
    } catch (_err) {
      // Continue to next candidate
    }
  }

  return null;
}

/**
 * Checks if a user has a global admin role.
 */
export async function isGlobalAdmin(userId: string): Promise<boolean> {
  try {
    // Check 1: User has an explicit 'admin' grant
    const adminGrant = await sqlWrapper.sql`
      SELECT g.id
      FROM grants g
      LEFT JOIN apps a ON g.app_id = a.id
      WHERE g.user_id = ${userId}
        AND g.role = 'admin'
        AND (
          a.spiffe_id = 'spiffe://system.local/auth-yes-management'
          OR a.name = 'Auth-Yes Management Console'
          OR g.app_id IS NULL
        )
    `.then((res: any) => res[0]);

    if (adminGrant) return true;

    // Check 2: First registered user in system fallback
    const firstUser = await sqlWrapper.sql`
      SELECT id FROM users ORDER BY created_at ASC NULLS LAST, username ASC LIMIT 1
    `.then((res: any) => res[0]);

    if (firstUser && firstUser.id === userId) {
      return true;
    }
  } catch (err) {
    console.error("[Auth API] isGlobalAdmin error:", err);
  }

  return false;
}

/**
 * Evaluates if the current user's capabilities satisfy the required scope.
 */
export function hasScope(
  auth: AuthenticatedUser,
  requiredScope: string,
): boolean {
  if (!auth.isAgent) return true;
  if (!Array.isArray(auth.customScopes)) return false;
  return auth.customScopes.includes("*") ||
    auth.customScopes.includes(requiredScope);
}

/**
 * Helper to check if the session itself is authorized as an admin.
 */
export async function isSessionAdmin(
  auth: AuthenticatedUser,
): Promise<boolean> {
  const globalAdmin = await isGlobalAdmin(auth.userId);
  if (!globalAdmin) return false;
  if (!auth.isAgent) return true;
  return Array.isArray(auth.customScopes) && auth.customScopes.includes("*");
}

/**
 * Computes root cookie domain from RP_ID or host.
 */
export function getCookieDomain(customRpId?: string): string {
  const rpId = customRpId || Deno.env.get("RP_ID");
  if (rpId) {
    if (rpId.includes("localhost") || rpId.includes("127.0.0.1")) {
      return rpId;
    }
    return `.${rpId}`;
  }
  return "";
}

/**
 * Hono Middleware: Blocks access if the session is a delegated agent session.
 */
export async function requirePrimarySession(
  c: Context,
  next: () => Promise<void>,
) {
  const auth = await getAuthenticatedUser(c);
  if (!auth) return c.json({ error: "Unauthorized" }, 401);
  if (auth.isAgent) {
    return c.json({ error: "Forbidden: Primary session required" }, 403);
  }
  await next();
}

/**
 * Hono Middleware Factory: Requires a specific scope.
 */
export function requireScope(scope: string) {
  return async (c: Context, next: () => Promise<void>) => {
    const auth = await getAuthenticatedUser(c);
    if (!auth) return c.json({ error: "Unauthorized" }, 401);
    if (!hasScope(auth, scope)) {
      return c.json({ error: "Forbidden: Insufficient scopes" }, 403);
    }
    await next();
  };
}

/**
 * Hono Middleware: Blocks access unless the session is an admin session.
 */
export async function requireAdmin(c: Context, next: () => Promise<void>) {
  const auth = await getAuthenticatedUser(c);
  if (!auth) return c.json({ error: "Unauthorized" }, 401);
  const isAdmin = await isSessionAdmin(auth);
  if (!isAdmin) {
    return c.json({ error: "Forbidden: Global admin access required" }, 403);
  }
  await next();
}

/**
 * Resolves application record by domain with Valkey caching.
 */
export async function getAppByHost(host: string): Promise<AppRecord | null> {
  const cacheKey = `auth:app_by_host:${host}`;
  try {
    const cached = await valkey.get(cacheKey);
    if (cached) return JSON.parse(cached);
  } catch (_e) {}

  try {
    const app = await sqlWrapper.sql`
      SELECT id, name, domain, is_public, bypass_paths, allowed_cidrs
      FROM apps WHERE domain = ${host} LIMIT 1
    `.then((res: any) => res[0]);

    if (app) {
      await valkey.setex(cacheKey, 300, JSON.stringify(app)).catch(() => {});
      return app;
    }
  } catch (_e) {}
  return null;
}

/**
 * Resolves role grant for a user on a given app.
 */
export async function getUserGrant(
  userId: string,
  appId: string,
): Promise<string | null> {
  try {
    const grant = await sqlWrapper.sql`
      SELECT role FROM user_grants WHERE user_id = ${userId} AND app_id = ${appId} LIMIT 1
    `.then((res: any) => res[0]);
    return grant?.role || null;
  } catch (_e) {
    return null;
  }
}