adam-acs/app/services/cwmp/state_machine_manager.ts

import { EventEmitter } from 'node:events'
import type { Device } from '../../domain/device/device.js'
import type { DeviceId } from '../../domain/device/device_id.js'
import { CWMPTaskQueue, ResultStatus } from './task_queue.js'
import { MainStateMachine } from './main_state_machine.js'
import { ImmediateCommandStateMachine } from './immediate_command_state_machine.js'
import { getDeviceHashIdentifier } from '../../utils/device_utils.js'
import { CommandStatus } from '../../domain/commands/command.js'
import { ModeType } from '../../domain/shared/mode.js'
import CommandModel from '../../models/command.js'
import logger from '@adonisjs/core/services/logger'

export enum MachineStatus {
  CREATED = 'created',
  RUNNING = 'running',
  COMPLETED = 'completed',
  FAILED = 'failed',
}

export interface StateMachineInstance {
  id: string
  deviceId: DeviceId
  session: string
  stateMachine: MainStateMachine | ImmediateCommandStateMachine
  status: MachineStatus
  createdAt: Date
  completedAt?: Date
  currentState: string
  lastTaskId?: string
  events?: string[]
}

/**
 * Central state machine manager with TCP-socket-based session tracking,
 * device preemption, and CWMP task queue integration.
 */
export class StateMachineManager extends EventEmitter {
  private taskQueue: CWMPTaskQueue
  private machines: Map<string, StateMachineInstance> = new Map()
  private deviceSessions: Map<string, string> = new Map() // deviceKey -> sessionId
  private sessionMachines: Map<string, string> = new Map() // sessionId -> machineId
  private remoteToDevice: Map<string, string> = new Map() // remoteKey -> deviceKey
  private cleanupInterval: ReturnType<typeof setInterval> | null = null

  private static instance: StateMachineManager | null = null

  constructor() {
    super()
    this.taskQueue = new CWMPTaskQueue()
    this.setupTaskQueueListeners()
  }

  static getInstance(): StateMachineManager {
    if (!StateMachineManager.instance) {
      StateMachineManager.instance = new StateMachineManager()
    }
    return StateMachineManager.instance
  }

  getTaskQueue(): CWMPTaskQueue {
    return this.taskQueue
  }

  startCleanupTimer(intervalMs: number = 60000): void {
    if (this.cleanupInterval) return
    this.cleanupInterval = setInterval(() => this.cleanupCompletedMachines(), intervalMs)
  }

  stopCleanupTimer(): void {
    if (this.cleanupInterval) {
      clearInterval(this.cleanupInterval)
      this.cleanupInterval = null
    }
  }

  /**
   * Start a state machine for a device session
   */
  async startStateMachine(
    device: Device,
    session: string,
    opts?: { events?: string[] }
  ): Promise<StateMachineInstance> {
    const deviceKey = JSON.stringify(device.deviceId)
    const evts = opts?.events || []
    const lastEvent = this.normalizeEvent(evts.length > 0 ? evts[evts.length - 1] : '')

    // Check for NI commands that force MainStateMachine
    let forceMainForNI = false
    try {
      const hashId = getDeviceHashIdentifier({
        manufacturer: device.deviceId.manufacturer,
        oui: device.deviceId.oui,
        productClass: device.deviceId.productClass,
        serialNumber: device.deviceId.serialNumber,
      })
      const collection = CommandModel.getCollection()
      const niCount = await collection.countDocuments({
        deviceIdentifier: hashId,
        mode: ModeType.NI,
        status: { $in: [CommandStatus.PENDING, CommandStatus.WAITING] },
      })
      forceMainForNI = niCount > 0
    } catch (e) {
      logger.warn(`Failed to check NI commands: ${e}`)
    }

    // Handle preemption of existing machine
    const existingSessionId = this.deviceSessions.get(deviceKey)
    if (existingSessionId) {
      const existingMachineId = this.sessionMachines.get(existingSessionId)
      if (existingMachineId) {
        const existingMachine = this.machines.get(existingMachineId)
        if (existingMachine && existingMachine.status === MachineStatus.RUNNING) {
          const isImmediate = existingMachine.stateMachine instanceof ImmediateCommandStateMachine
          const onboardingEvent = ['0 BOOTSTRAP', '1 BOOT', '2 PERIODIC', '3 SCHEDULED'].includes(
            lastEvent
          )
          const shouldPreempt = isImmediate && (forceMainForNI || onboardingEvent)

          if (shouldPreempt) {
            // Rebind session if changed
            if (existingSessionId !== session) {
              this.sessionMachines.delete(existingSessionId)
              this.sessionMachines.set(session, existingMachineId)
              this.deviceSessions.set(deviceKey, session)
              existingMachine.session = session
            }
            // Replace with MainStateMachine
            const newMachine = new MainStateMachine(device, session, this.taskQueue)
            ;(existingMachine as any).stateMachine = newMachine
            existingMachine.events = evts
            await newMachine.start()
            existingMachine.currentState = newMachine.getCurrentStateName()
            if (newMachine.isCompleted()) this.completeMachine(existingMachine)
            return existingMachine
          }

          // Rebind session to existing machine
          if (existingSessionId !== session) {
            this.sessionMachines.delete(existingSessionId)
            this.sessionMachines.set(session, existingMachineId)
            this.deviceSessions.set(deviceKey, session)
            existingMachine.session = session
          }
          return existingMachine
        }
      }
    }

    // Create new machine
    const machineId = `${deviceKey}_${session}_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`

    let stateMachine: MainStateMachine | ImmediateCommandStateMachine
    if (forceMainForNI) {
      stateMachine = new MainStateMachine(device, session, this.taskQueue)
    } else if (lastEvent === '6 CONNECTION REQUEST') {
      stateMachine = new ImmediateCommandStateMachine(device, session, this.taskQueue)
    } else if (
      ['0 BOOTSTRAP', '1 BOOT', '2 PERIODIC', '3 SCHEDULED'].includes(lastEvent) ||
      !lastEvent
    ) {
      stateMachine = new MainStateMachine(device, session, this.taskQueue)
    } else {
      // Unknown event - use MainStateMachine as default
      stateMachine = new MainStateMachine(device, session, this.taskQueue)
    }

    const machineInstance: StateMachineInstance = {
      id: machineId,
      deviceId: device.deviceId,
      session,
      stateMachine,
      status: MachineStatus.CREATED,
      createdAt: new Date(),
      currentState: stateMachine.getCurrentStateName(),
      events: evts,
    }

    this.machines.set(machineId, machineInstance)
    this.deviceSessions.set(deviceKey, session)
    this.sessionMachines.set(session, machineId)

    // Start the machine
    try {
      machineInstance.status = MachineStatus.RUNNING
      machineInstance.currentState = stateMachine.getCurrentStateName()

      if (stateMachine.isCompleted()) {
        this.completeMachine(machineInstance)
      }
    } catch (error) {
      machineInstance.status = MachineStatus.FAILED
      machineInstance.completedAt = new Date()
      logger.error(`Failed to start state machine: ${error}`)
    }

    return machineInstance
  }

  /**
   * Bind a remote endpoint to a device for session rebinding
   */
  bindRemoteToDevice(remoteKey: string, deviceId: DeviceId): void {
    this.remoteToDevice.set(remoteKey, JSON.stringify(deviceId))
  }

  /**
   * Returns the hash identifier for the device bound to a remote endpoint.
   * Used by trace logging middleware to correlate non-Inform messages.
   */
  getHashIdForRemote(remoteKey: string): string | undefined {
    const deviceKeyJson = this.remoteToDevice.get(remoteKey)
    if (!deviceKeyJson) return undefined

    try {
      const deviceId = JSON.parse(deviceKeyJson)
      return getDeviceHashIdentifier({
        manufacturer: deviceId.manufacturer ?? '',
        oui: deviceId.oui,
        productClass: deviceId.productClass,
        serialNumber: deviceId.serialNumber,
      })
    } catch {
      return undefined
    }
  }

  /**
   * Ensure session is bound when TCP connection changes
   */
  ensureSessionForRemote(remoteKey: string, newSession: string): boolean {
    const deviceKey = this.remoteToDevice.get(remoteKey)
    if (!deviceKey) return false

    const existingSession = this.deviceSessions.get(deviceKey)
    if (!existingSession) return false

    const machineId = this.sessionMachines.get(existingSession)
    if (!machineId) return false

    const machine = this.machines.get(machineId)
    if (!machine) return false

    if (existingSession !== newSession) {
      this.sessionMachines.delete(existingSession)
      this.sessionMachines.set(newSession, machineId)
      this.deviceSessions.set(deviceKey, newSession)
      machine.session = newSession
    }

    return true
  }

  /**
   * Continue when CPE sends empty body
   */
  async continue(session: string): Promise<string> {
    const machineId = this.sessionMachines.get(session)
    if (!machineId) return ''

    const machineInstance = this.machines.get(machineId)
    if (!machineInstance || machineInstance.status !== MachineStatus.RUNNING) return ''

    try {
      // If a task is in-flight and CPE sends empty body, it means the CPE
      // cannot/will not respond to our request. Treat as task failure.
      if (machineInstance.lastTaskId) {
        const failedTaskId = machineInstance.lastTaskId
        machineInstance.lastTaskId = undefined

        logger.warn(
          { session, taskId: failedTaskId },
          'CPE sent empty body while task in-flight — treating as task failure'
        )

        this.taskQueue.failTask(failedTaskId, 'CPE sent empty body (no response)')
        try {
          await this.taskQueue.executeTaskCallback(failedTaskId, ResultStatus.Error, {
            FaultCode: 9000,
            FaultString: 'CPE did not respond (empty body)',
          })
        } catch (cbErr) {
          logger.error(`Task callback error on empty-body failure: ${cbErr}`)
        }

        // Let the state machine handle the failure and advance
        const response = await machineInstance.stateMachine.handleResponse(false, {
          FaultCode: 9000,
          FaultString: 'CPE did not respond (empty body)',
        })
        machineInstance.currentState = machineInstance.stateMachine.getCurrentStateName()

        if (machineInstance.stateMachine.isCompleted()) {
          this.completeMachine(machineInstance)
          return ''
        }

        return response
      }

      // Serve queued tasks
      const pendingTask = this.taskQueue.getTaskForDevice(machineInstance.deviceId)
      if (pendingTask) {
        const builder = CWMPTaskQueue.task_map[pendingTask.task]
        if (typeof builder === 'function') {
          const envelope = builder(pendingTask.props)
          if (pendingTask.id) machineInstance.lastTaskId = pendingTask.id
          return envelope
        }
      }

      // Let the state machine continue
      const response = await machineInstance.stateMachine.handleResponse(true, { isFirst: false })
      machineInstance.currentState = machineInstance.stateMachine.getCurrentStateName()

      if (machineInstance.stateMachine.isCompleted()) {
        this.completeMachine(machineInstance)
        return ''
      }

      return response
    } catch (error) {
      logger.error(`Error continuing session ${session}: ${error}`)
      this.failMachine(machineInstance, String(error))
      return ''
    }
  }

  /**
   * Handle a CWMP response message
   */
  async handleCWMPResponse(session: string, success: boolean, message: any): Promise<string> {
    const machineId = this.sessionMachines.get(session)
    if (!machineId) return ''

    const machineInstance = this.machines.get(machineId)
    if (!machineInstance || machineInstance.status !== MachineStatus.RUNNING) return ''

    try {
      const responseTaskId = machineInstance.lastTaskId
      const response = await machineInstance.stateMachine.handleResponse(success, message)

      // Execute task callback
      if (responseTaskId) {
        if (success) {
          this.taskQueue.completeTask(responseTaskId, message)
        } else {
          this.taskQueue.failTask(responseTaskId, JSON.stringify(message))
        }
        try {
          await this.taskQueue.executeTaskCallback(
            responseTaskId,
            success ? ResultStatus.Success : ResultStatus.Error,
            message
          )
        } catch (cbErr) {
          logger.error(`Task callback error: ${cbErr}`)
        }
        if (machineInstance.lastTaskId === responseTaskId) {
          machineInstance.lastTaskId = undefined
        }
      }

      machineInstance.currentState = machineInstance.stateMachine.getCurrentStateName()

      if (machineInstance.stateMachine.isCompleted()) {
        this.completeMachine(machineInstance)
        return ''
      }

      return response
    } catch (error) {
      logger.error(`Error handling CWMP response: ${error}`)
      this.failMachine(machineInstance, String(error))
      return ''
    }
  }

  hasSession(session: string): boolean {
    return this.sessionMachines.has(session)
  }

  /**
   * Enqueue an immediate task for a device (e.g., after Inform)
   */
  async enqueueImmediate(device: Device, type: string, payload: any): Promise<void> {
    this.taskQueue.addCWMPTask(device.deviceId, type, payload)
  }

  private completeMachine(machine: StateMachineInstance): void {
    machine.status = MachineStatus.COMPLETED
    machine.completedAt = new Date()
    const deviceKey = JSON.stringify(machine.deviceId)
    this.deviceSessions.delete(deviceKey)
    this.sessionMachines.delete(machine.session)
  }

  private failMachine(machine: StateMachineInstance, error: string): void {
    machine.status = MachineStatus.FAILED
    machine.completedAt = new Date()
    const deviceKey = JSON.stringify(machine.deviceId)
    this.deviceSessions.delete(deviceKey)
    this.sessionMachines.delete(machine.session)
    logger.error(`State machine failed: ${error}`)
  }

  private normalizeEvent(evt: string): string {
    return (evt || '').toString().trim().replace(/\s+/g, ' ').toUpperCase()
  }

  cleanupCompletedMachines(olderThanHours: number = 1): void {
    const cutoff = new Date(Date.now() - olderThanHours * 60 * 60 * 1000)
    for (const [id, machine] of this.machines.entries()) {
      if (
        (machine.status === MachineStatus.COMPLETED || machine.status === MachineStatus.FAILED) &&
        machine.completedAt &&
        machine.completedAt < cutoff
      ) {
        this.machines.delete(id)
      }
    }
  }

  async shutdown(): Promise<void> {
    this.stopCleanupTimer()
    this.cleanupCompletedMachines(0)
  }

  private setupTaskQueueListeners(): void {
    this.taskQueue.on('taskSent', (task: any) => {
      // Find the machine for this device and set lastTaskId
      const deviceKey = JSON.stringify(task.deviceId)
      const sessionId = this.deviceSessions.get(deviceKey)
      if (sessionId) {
        const machineId = this.sessionMachines.get(sessionId)
        if (machineId) {
          const machine = this.machines.get(machineId)
          if (machine) machine.lastTaskId = task.id
        }
      }
    })
  }
}