NextEdu/src/shared/lib/auth-guard.ts

import type { Permission, DataScope, AuthContext, Role } from "@/shared/types/permissions"
import { db } from "@/shared/db"
import {
  classes,
  classEnrollments,
  classSubjectTeachers,
  grades,
  parentStudentRelations,
} from "@/shared/db/schema"
import { eq, inArray, or } from "drizzle-orm"
import { getSession } from "@/shared/lib/session"
import { PermissionDeniedError } from "@/shared/lib/errors"

// Re-export for backward compatibility (other modules still import from here)
export { PermissionDeniedError } from "@/shared/lib/errors"

/**
 * Get the full authentication context for the current user.
 * Throws if not authenticated.
 */
export async function getAuthContext(): Promise<AuthContext> {
  const session = await getSession()
  const userId = session?.user?.id
  if (!userId) throw new PermissionDeniedError("auth_required")

  // Prefer session data (already resolved in JWT callback)
  const roleNames = (session.user.roles ?? []) as Role[]
  const permissions = (session.user.permissions ?? []) as Permission[]

  // Resolve data scope from DB (not cached in JWT since it can change)
  const dataScope = await resolveDataScope(userId, roleNames)

  return { userId, roles: roleNames, permissions, dataScope }
}

/**
 * Assert the current user has the specified permission.
 * Returns AuthContext on success, throws PermissionDeniedError on failure.
 */
export async function requirePermission(permission: Permission): Promise<AuthContext> {
  const ctx = await getAuthContext()
  if (!ctx.permissions.includes(permission)) {
    throw new PermissionDeniedError(permission)
  }
  return ctx
}

/**
 * Check permission without throwing. Useful for conditional logic.
 */
export async function checkPermission(
  permission: Permission
): Promise<{ allowed: boolean; ctx: AuthContext }> {
  const ctx = await getAuthContext()
  return { allowed: ctx.permissions.includes(permission), ctx }
}

/**
 * Resolve the data scope for a user based on their roles.
 * Queries the DB for resource ownership information.
 */
async function resolveDataScope(userId: string, roleNames: Role[]): Promise<DataScope> {
  // Admin sees everything
  if (roleNames.includes("admin")) {
    return { type: "all" }
  }

  // Grade head / teaching head: can manage their grades
  if (roleNames.includes("grade_head") || roleNames.includes("teaching_head")) {
    const managedGrades = await db
      .select({ id: grades.id })
      .from(grades)
      .where(or(eq(grades.gradeHeadId, userId), eq(grades.teachingHeadId, userId)))

    if (managedGrades.length > 0) {
      return { type: "grade_managed", gradeIds: managedGrades.map((g) => g.id) }
    }
  }

  // Teacher: can see their own classes
  if (roleNames.includes("teacher")) {
    // Classes where user is the homeroom teacher
    const homeroomClasses = await db
      .select({ id: classes.id })
      .from(classes)
      .where(eq(classes.teacherId, userId))

    // Classes where user is a subject teacher
    const subjectClasses = await db
      .selectDistinct({ classId: classSubjectTeachers.classId, subjectId: classSubjectTeachers.subjectId })
      .from(classSubjectTeachers)
      .where(eq(classSubjectTeachers.teacherId, userId))

    const classIds = [
      ...new Set([
        ...homeroomClasses.map((c) => c.id),
        ...subjectClasses.map((c) => c.classId),
      ]),
    ]
    const subjectIds = subjectClasses
      .map((c) => c.subjectId)
      .filter((s): s is string => s !== null)

    return {
      type: "class_taught",
      classIds,
      subjectIds: subjectIds.length > 0 ? subjectIds : undefined,
    }
  }

  // Student: can see data from their enrolled classes
  // Pre-resolve classIds and gradeIds here to avoid N+1 queries in data-access layer
  if (roleNames.includes("student")) {
    const enrolledClasses = await db
      .select({ classId: classEnrollments.classId, gradeId: classes.gradeId })
      .from(classEnrollments)
      .innerJoin(classes, eq(classEnrollments.classId, classes.id))
      .where(eq(classEnrollments.studentId, userId))

    const gradeIds = [
      ...new Set(
        enrolledClasses
          .map((c) => c.gradeId)
          .filter((g): g is string => g !== null),
      ),
    ]

    return {
      type: "class_members",
      classIds: enrolledClasses.map((c) => c.classId),
      gradeIds: gradeIds.length > 0 ? gradeIds : undefined,
    }
  }

  // Parent: can see their children's data
  if (roleNames.includes("parent")) {
    const children = await db
      .select({ studentId: parentStudentRelations.studentId })
      .from(parentStudentRelations)
      .where(eq(parentStudentRelations.parentId, userId))

    const childrenIds = children.map((c) => c.studentId)

    // Pre-resolve gradeIds from children's enrolled classes
    let gradeIds: string[] | undefined
    if (childrenIds.length > 0) {
      const childrenClasses = await db
        .select({ gradeId: classes.gradeId })
        .from(classEnrollments)
        .innerJoin(classes, eq(classEnrollments.classId, classes.id))
        .where(inArray(classEnrollments.studentId, childrenIds))

      const uniqueGradeIds = [
        ...new Set(
          childrenClasses
            .map((c) => c.gradeId)
            .filter((g): g is string => g !== null),
        ),
      ]
      gradeIds = uniqueGradeIds.length > 0 ? uniqueGradeIds : undefined
    }

    return { type: "children", childrenIds, gradeIds }
  }

  // Fallback: only own data
  return { type: "owned", userId }
}

/**
 * Convenience: assert the user is authenticated (has any role).
 * Returns AuthContext on success.
 */
export async function requireAuth(): Promise<AuthContext> {
  return getAuthContext()
}