remotelink-docker/app/api/agent/register/route.ts

import { db } from '@/lib/db'
import { machines, enrollmentTokens, users } from '@/lib/db/schema'
import { eq, asc } from 'drizzle-orm'
import { NextRequest, NextResponse } from 'next/server'
import { randomBytes } from 'crypto'

// POST /api/agent/register
// Three modes:
//   1. Re-register (existing agent):  { accessKey, ... }         → updates machine
//   2. Enrollment token:              { enrollmentToken, ... }   → creates machine owned by token creator
//   3. Open enrollment (no token):    { name, ... }              → creates machine, OPEN_ENROLLMENT must be true
export async function POST(request: NextRequest) {
  try {
    const body = await request.json()
    const { accessKey, enrollmentToken, name, hostname, os, osVersion, agentVersion, ipAddress, macAddress } = body

    // ── Mode 1: existing agent re-registering ─────────────────────────────────
    if (accessKey) {
      const result = await db
        .select()
        .from(machines)
        .where(eq(machines.accessKey, accessKey))
        .limit(1)

      const machine = result[0]
      if (!machine) {
        return NextResponse.json({ error: 'Invalid access key' }, { status: 401 })
      }

      await db
        .update(machines)
        .set({
          name: name || machine.name,
          hostname: hostname || machine.hostname,
          os: os || machine.os,
          osVersion: osVersion || machine.osVersion,
          agentVersion: agentVersion || machine.agentVersion,
          ipAddress: ipAddress || machine.ipAddress,
          macAddress: macAddress || machine.macAddress,
          isOnline: true,
          lastSeen: new Date(),
          updatedAt: new Date(),
        })
        .where(eq(machines.id, machine.id))

      return NextResponse.json({ machineId: machine.id })
    }

    if (!name) {
      return NextResponse.json({ error: 'name is required' }, { status: 400 })
    }

    const newAccessKey = randomBytes(32).toString('hex')

    // ── Mode 2: enrollment token ───────────────────────────────────────────────
    if (enrollmentToken) {
      const tokenResult = await db
        .select()
        .from(enrollmentTokens)
        .where(eq(enrollmentTokens.token, enrollmentToken))
        .limit(1)

      const token = tokenResult[0]
      if (!token) return NextResponse.json({ error: 'Invalid enrollment token' }, { status: 401 })
      if (token.revokedAt) return NextResponse.json({ error: 'Enrollment token revoked' }, { status: 401 })
      if (token.expiresAt && token.expiresAt < new Date()) return NextResponse.json({ error: 'Enrollment token expired' }, { status: 401 })
      if (token.maxUses !== null && token.usedCount >= token.maxUses) return NextResponse.json({ error: 'Enrollment token use limit reached' }, { status: 401 })

      const newMachine = await db
        .insert(machines)
        .values({ userId: token.createdBy!, name, hostname, os, osVersion, agentVersion, ipAddress, macAddress, accessKey: newAccessKey, isOnline: true, lastSeen: new Date() })
        .returning({ id: machines.id })

      await db.update(enrollmentTokens).set({ usedCount: token.usedCount + 1 }).where(eq(enrollmentTokens.id, token.id))

      return NextResponse.json({ machineId: newMachine[0].id, accessKey: newAccessKey })
    }

    // ── Mode 3: open enrollment (no token required) ────────────────────────────
    const openEnrollment = process.env.OPEN_ENROLLMENT === 'true'
    if (!openEnrollment) {
      return NextResponse.json({ error: 'accessKey or enrollmentToken required' }, { status: 400 })
    }

    // Find the owner: OPEN_ENROLLMENT_USER_EMAIL env var, or first admin, or first user
    let ownerId: string | null = null
    const ownerEmail = process.env.OPEN_ENROLLMENT_USER_EMAIL
    if (ownerEmail) {
      const ownerResult = await db.select({ id: users.id }).from(users).where(eq(users.email, ownerEmail)).limit(1)
      ownerId = ownerResult[0]?.id ?? null
    }
    if (!ownerId) {
      // Fall back to first admin, then first user
      const adminResult = await db.select({ id: users.id }).from(users).where(eq(users.role, 'admin')).orderBy(asc(users.createdAt)).limit(1)
      ownerId = adminResult[0]?.id ?? null
    }
    if (!ownerId) {
      const anyUser = await db.select({ id: users.id }).from(users).orderBy(asc(users.createdAt)).limit(1)
      ownerId = anyUser[0]?.id ?? null
    }
    if (!ownerId) {
      return NextResponse.json({ error: 'No users exist on the server yet. Create an account first.' }, { status: 503 })
    }

    const newMachine = await db
      .insert(machines)
      .values({ userId: ownerId, name, hostname, os, osVersion, agentVersion, ipAddress, macAddress, accessKey: newAccessKey, isOnline: true, lastSeen: new Date() })
      .returning({ id: machines.id })

    return NextResponse.json({ machineId: newMachine[0].id, accessKey: newAccessKey })
  } catch (error) {
    console.error('[Agent Register] Error:', error)
    return NextResponse.json({ error: 'Internal server error' }, { status: 500 })
  }
}