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

import { db } from '@/lib/db'
import { machines, enrollmentTokens } from '@/lib/db/schema'
import { eq, and, isNull, or, gt } from 'drizzle-orm'
import { NextRequest, NextResponse } from 'next/server'
import { randomBytes } from 'crypto'

// POST /api/agent/register
// Two modes:
//   1. First-time: { enrollmentToken, name, hostname, os, osVersion, agentVersion, ipAddress }
//      → creates machine, returns { machineId, accessKey }
//   2. Re-register: { accessKey, name, hostname, os, osVersion, agentVersion, ipAddress }
//      → updates existing machine, returns { machineId }
export async function POST(request: NextRequest) {
  try {
    const body = await request.json()
    const { accessKey, enrollmentToken, name, hostname, os, osVersion, agentVersion, ipAddress } = body

    // ── Mode 2: 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,
          isOnline: true,
          lastSeen: new Date(),
          updatedAt: new Date(),
        })
        .where(eq(machines.id, machine.id))

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

    // ── Mode 1: first-time registration with enrollment token ─────────────────
    if (!enrollmentToken) {
      return NextResponse.json({ error: 'accessKey or enrollmentToken required' }, { status: 400 })
    }

    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 })
    }

    // Check revoked
    if (token.revokedAt) {
      return NextResponse.json({ error: 'Enrollment token has been revoked' }, { status: 401 })
    }

    // Check expiry
    if (token.expiresAt && token.expiresAt < new Date()) {
      return NextResponse.json({ error: 'Enrollment token has expired' }, { status: 401 })
    }

    // Check max uses
    if (token.maxUses !== null && token.usedCount >= token.maxUses) {
      return NextResponse.json({ error: 'Enrollment token has reached its use limit' }, { status: 401 })
    }

    if (!name) {
      return NextResponse.json({ error: 'name is required for first-time registration' }, { status: 400 })
    }

    // Generate a secure access key for this machine
    const newAccessKey = randomBytes(32).toString('hex')

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

    // Increment use count
    await db
      .update(enrollmentTokens)
      .set({ usedCount: token.usedCount + 1 })
      .where(eq(enrollmentTokens.id, token.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 })
  }
}