/*
 * Copyright 2025 Daytona Platforms Inc.
 * SPDX-License-Identifier: AGPL-3.0
 */

import { Injectable, Logger, NotFoundException, OnApplicationShutdown } from '@nestjs/common'
import { InjectRepository } from '@nestjs/typeorm'
import { Cron, CronExpression } from '@nestjs/schedule'
import { SnapshotRepository } from '../repositories/snapshot.repository'
import { In, IsNull, LessThan, Not, Repository } from 'typeorm'
import { DockerRegistryService } from '../../docker-registry/services/docker-registry.service'
import { Snapshot } from '../entities/snapshot.entity'
import { SnapshotState } from '../enums/snapshot-state.enum'
import { SnapshotRunner } from '../entities/snapshot-runner.entity'
import { Runner } from '../entities/runner.entity'
import { DockerRegistry } from '../../docker-registry/entities/docker-registry.entity'
import { RunnerState } from '../enums/runner-state.enum'
import { SnapshotRunnerState } from '../enums/snapshot-runner-state.enum'
import { v4 as uuidv4 } from 'uuid'
import { RunnerNotReadyError } from '../errors/runner-not-ready.error'
import { RedisLockProvider } from '../common/redis-lock.provider'
import { OrganizationService } from '../../organization/services/organization.service'
import { BuildInfo } from '../entities/build-info.entity'
import { fromAxiosError } from '../../common/utils/from-axios-error'
import { InjectRedis } from '@nestjs-modules/ioredis'
import { Redis } from 'ioredis'
import { RunnerService } from '../services/runner.service'
import { TrackableJobExecutions } from '../../common/interfaces/trackable-job-executions'
import { TrackJobExecution } from '../../common/decorators/track-job-execution.decorator'
import { setTimeout as sleep } from 'timers/promises'
import { LogExecution } from '../../common/decorators/log-execution.decorator'
import { WithInstrumentation } from '../../common/decorators/otel.decorator'
import { RunnerAdapterFactory, SnapshotDigestResponse } from '../runner-adapter/runnerAdapter'
import { SnapshotStateError } from '../errors/snapshot-state-error'
import { SnapshotEvents } from '../constants/snapshot-events'
import { SnapshotCreatedEvent } from '../events/snapshot-created.event'
import { SnapshotService } from '../services/snapshot.service'
import { OnAsyncEvent } from '../../common/decorators/on-async-event.decorator'
import { parseDockerImage } from '../../common/utils/docker-image.util'
import { SandboxState } from '../enums/sandbox-state.enum'
import { SandboxDesiredState } from '../enums/sandbox-desired-state.enum'
import { BackupState } from '../enums/backup-state.enum'
import { BadRequestError } from '../../exceptions/bad-request.exception'
import { SandboxRepository } from '../repositories/sandbox.repository'
import { SnapshotInfoResponse } from '@daytona/runner-api-client'
import { SnapshotActivatedEvent } from '../events/snapshot-activated.event'
import { TypedConfigService } from '../../config/typed-config.service'
import { RegionType } from '../../region/enums/region-type.enum'

/** Fisher-Yates shuffle — uniform random permutation in O(n). */
function shuffleArray<T>(array: T[]): T[] {
  const shuffled = [...array]
  for (let i = shuffled.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1))
    ;[shuffled[i], shuffled[j]] = [shuffled[j], shuffled[i]]
  }
  return shuffled
}

const SYNC_AGAIN = 'sync-again'
const DONT_SYNC_AGAIN = 'dont-sync-again'
const DEFAULT_SNAPSHOT_DEACTIVATION_TIMEOUT_MINUTES = 14 * 24 * 60 // 14 days
type SyncState = typeof SYNC_AGAIN | typeof DONT_SYNC_AGAIN
const BASE_PROPAGATION_FACTOR = 1 / 3

@Injectable()
export class SnapshotManager implements TrackableJobExecutions, OnApplicationShutdown {
  activeJobs = new Set<string>()

  private readonly logger = new Logger(SnapshotManager.name)
  //  generate a unique instance id used to ensure only one instance of the worker is handing the
  //  snapshot activation
  private readonly instanceId = uuidv4()

  constructor(
    @InjectRedis() private readonly redis: Redis,
    private readonly snapshotRepository: SnapshotRepository,
    @InjectRepository(SnapshotRunner)
    private readonly snapshotRunnerRepository: Repository<SnapshotRunner>,
    @InjectRepository(Runner)
    private readonly runnerRepository: Repository<Runner>,
    private readonly sandboxRepository: SandboxRepository,
    @InjectRepository(BuildInfo)
    private readonly buildInfoRepository: Repository<BuildInfo>,
    private readonly runnerService: RunnerService,
    private readonly dockerRegistryService: DockerRegistryService,
    private readonly runnerAdapterFactory: RunnerAdapterFactory,
    private readonly redisLockProvider: RedisLockProvider,
    private readonly organizationService: OrganizationService,
    private readonly snapshotService: SnapshotService,
    private readonly configService: TypedConfigService,
  ) {}

  async onApplicationShutdown() {
    //  wait for all active jobs to finish
    while (this.activeJobs.size > 0) {
      this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`)
      await sleep(1000)
    }
  }

  @Cron(CronExpression.EVERY_MINUTE, { name: 'scale-down-runner-snapshots', waitForCompletion: true })
  @TrackJobExecution()
  @LogExecution('scale-down-runner-snapshots')
  @WithInstrumentation()
  async scaleDownRunnerSnapshots() {
    const lockKey = 'scale-down-runner-snapshots-lock'
    const lockTtl = 3 * 60 // seconds (3 min)
    if (!(await this.redisLockProvider.lock(lockKey, lockTtl))) {
      return
    }

    const skip = (await this.redis.get('scale-down-runner-snapshots-skip')) || 0

    const snapshots = await this.snapshotRepository
      .createQueryBuilder('snapshot')
      .innerJoin('organization', 'org', 'org.id = snapshot.organizationId')
      .innerJoin('region_quota', 'rq', 'rq."organizationId" = org.id AND rq."regionId" = org."defaultRegionId"')
      .select(['snapshot.*', 'rq.total_cpu_quota'])
      .where('snapshot.state = :snapshotState', { snapshotState: SnapshotState.ACTIVE })
      .andWhere('org.suspended = false')
      .orderBy('snapshot.createdAt', 'ASC')
      .limit(50)
      .offset(Number(skip))
      .getRawMany()

    if (snapshots.length === 0) {
      await this.redisLockProvider.unlock(lockKey)
      await this.redis.set('scale-down-runner-snapshots-skip', 0)
      return
    }

    await this.redis.set('scale-down-runner-snapshots-skip', Number(skip) + snapshots.length)

    const results = await Promise.allSettled(
      snapshots.map(async (snapshot) => {
        const regions = await this.snapshotService.getSnapshotRegions(snapshot.id)

        const sharedRegionIds = regions
          .filter((r) => r.organizationId === null && r.regionType === RegionType.SHARED)
          .map((r) => r.id)

        return this.scaleDownSnapshotFromRunners(snapshot, sharedRegionIds)
      }),
    )

    results.forEach((result) => {
      if (result.status === 'rejected') {
        this.logger.error(`Error scaling down snapshot from runners: ${fromAxiosError(result.reason)}`)
      }
    })

    await this.redisLockProvider.unlock(lockKey)
  }

  @Cron(CronExpression.EVERY_5_SECONDS, { name: 'sync-runner-snapshots', waitForCompletion: true })
  @TrackJobExecution()
  @LogExecution('sync-runner-snapshots')
  @WithInstrumentation()
  async syncRunnerSnapshots() {
    const lockKey = 'sync-runner-snapshots-lock'
    const lockTtl = 10 * 60 // seconds (10 min)
    if (!(await this.redisLockProvider.lock(lockKey, lockTtl))) {
      return
    }

    const skip = (await this.redis.get('sync-runner-snapshots-skip')) || 0

    const snapshots = await this.snapshotRepository
      .createQueryBuilder('snapshot')
      .innerJoin('organization', 'org', 'org.id = snapshot.organizationId')
      .where('snapshot.state = :snapshotState', { snapshotState: SnapshotState.ACTIVE })
      .andWhere('org.suspended = false')
      .orderBy('snapshot.createdAt', 'ASC')
      .limit(150)
      .offset(Number(skip))
      .getMany()

    if (snapshots.length === 0) {
      await this.redisLockProvider.unlock(lockKey)
      await this.redis.set('sync-runner-snapshots-skip', 0)
      return
    }

    await this.redis.set('sync-runner-snapshots-skip', Number(skip) + snapshots.length)

    const results = await Promise.allSettled(
      snapshots.map(async (snapshot) => {
        const regions = await this.snapshotService.getSnapshotRegions(snapshot.id)

        const sharedRegionIds = regions.filter((r) => r.organizationId === null).map((r) => r.id)
        const organizationRegionIds = regions
          .filter((r) => r.organizationId === snapshot.organizationId)
          .map((r) => r.id)

        return this.propagateSnapshotToRunners(snapshot, sharedRegionIds, organizationRegionIds)
      }),
    )

    // Log all promise errors
    results.forEach((result) => {
      if (result.status === 'rejected') {
        this.logger.error(`Error propagating snapshot to runners: ${fromAxiosError(result.reason)}`)
      }
    })

    await this.redisLockProvider.unlock(lockKey)
  }

  @Cron(CronExpression.EVERY_5_SECONDS, { name: 'sync-runner-snapshot-states', waitForCompletion: true })
  @TrackJobExecution()
  @LogExecution('sync-runner-snapshot-states')
  @WithInstrumentation()
  async syncRunnerSnapshotStates() {
    //  this approach is not ideal, as if the number of runners is large, this will take a long time
    //  also, if some snapshots stuck in a "pulling" state, they will infest the queue
    //  todo: find a better approach

    const lockKey = 'sync-runner-snapshot-states-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 30))) {
      return
    }

    const runnerSnapshots = await this.snapshotRunnerRepository
      .createQueryBuilder('snapshotRunner')
      .where({
        state: In([SnapshotRunnerState.PULLING_SNAPSHOT, SnapshotRunnerState.BUILDING_SNAPSHOT]),
      })
      .orderBy('RANDOM()')
      .take(100)
      .getMany()

    await Promise.allSettled(
      runnerSnapshots.map((snapshotRunner) => {
        return this.syncRunnerSnapshotState(snapshotRunner).catch((err) => {
          if (err.code !== 'ECONNRESET') {
            if (err instanceof RunnerNotReadyError) {
              this.logger.debug(
                `Runner ${snapshotRunner.runnerId} is not ready while trying to sync snapshot runner ${snapshotRunner.id}: ${err}`,
              )
              return
            }
            this.logger.error(`Error syncing runner snapshot state ${snapshotRunner.id}: ${fromAxiosError(err)}`)
            this.snapshotRunnerRepository.update(snapshotRunner.id, {
              state: SnapshotRunnerState.ERROR,
              errorReason: fromAxiosError(err).message,
            })
          }
        })
      }),
    )

    await this.redisLockProvider.unlock(lockKey)
  }

  // REMOVING is split out to its own cron with its own lock so it doesn't compete with
  // the higher-priority PULLING/BUILDING states. We additionally wait at least 3 minutes
  // since the last update to give in-flight removals a chance to settle on the runner.
  @Cron(CronExpression.EVERY_5_SECONDS, { name: 'sync-removing-runner-snapshot-states' })
  @TrackJobExecution()
  @LogExecution('sync-removing-runner-snapshot-states')
  @WithInstrumentation()
  async syncRemovingRunnerSnapshotStates() {
    const lockKey = 'sync-removing-runner-snapshot-states-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 120))) {
      return
    }

    const threeMinutesAgo = new Date(Date.now() - 3 * 60 * 1000)
    const runnerSnapshots = await this.snapshotRunnerRepository
      .createQueryBuilder('snapshotRunner')
      .where({
        state: SnapshotRunnerState.REMOVING,
        updatedAt: LessThan(threeMinutesAgo),
      })
      .orderBy('RANDOM()')
      .take(200)
      .getMany()

    await Promise.allSettled(
      runnerSnapshots.map((snapshotRunner) => {
        return this.syncRunnerSnapshotState(snapshotRunner).catch((err) => {
          if (err.code !== 'ECONNRESET') {
            if (err instanceof RunnerNotReadyError) {
              this.logger.debug(
                `Runner ${snapshotRunner.runnerId} is not ready while trying to sync snapshot runner ${snapshotRunner.id}: ${err}`,
              )
              return
            }
            this.logger.error(`Error syncing runner snapshot state ${snapshotRunner.id}: ${fromAxiosError(err)}`)
            this.snapshotRunnerRepository.update(snapshotRunner.id, {
              state: SnapshotRunnerState.ERROR,
              errorReason: fromAxiosError(err).message,
            })
          }
        })
      }),
    )

    await this.redisLockProvider.unlock(lockKey)
  }

  async syncRunnerSnapshotState(snapshotRunner: SnapshotRunner): Promise<void> {
    const runner = await this.runnerService.findOne(snapshotRunner.runnerId)
    if (!runner) {
      //  cleanup the snapshot runner record if the runner is not found
      //  this can happen if the runner is deleted from the database without cleaning up the snapshot runners
      await this.snapshotRunnerRepository.delete(snapshotRunner.id)
      this.logger.warn(
        `Runner ${snapshotRunner.runnerId} not found while trying to process snapshot runner ${snapshotRunner.id}. Snapshot runner has been removed.`,
      )
      return
    }

    if (runner.state !== RunnerState.READY) {
      //  todo: handle timeout policy
      //  for now just remove the snapshot runner record if the runner is not ready
      await this.snapshotRunnerRepository.delete(snapshotRunner.id)

      throw new RunnerNotReadyError(`Runner ${runner.id} is not ready`)
    }

    switch (snapshotRunner.state) {
      case SnapshotRunnerState.PULLING_SNAPSHOT:
        await this.handleSnapshotRunnerStatePullingSnapshot(snapshotRunner, runner)
        break
      case SnapshotRunnerState.BUILDING_SNAPSHOT:
        await this.handleSnapshotRunnerStateBuildingSnapshot(snapshotRunner, runner)
        break
      case SnapshotRunnerState.REMOVING:
        await this.handleSnapshotRunnerStateRemoving(snapshotRunner, runner)
        break
    }
  }

  async propagateSnapshotToRunners(snapshot: Snapshot, sharedRegionIds: string[], organizationRegionIds: string[]) {
    //  todo: remove try catch block and implement error handling
    try {
      //  get all runners in the regions to propagate to
      const runners = await this.runnerRepository.find({
        where: {
          state: RunnerState.READY,
          unschedulable: Not(true),
          region: In([...sharedRegionIds, ...organizationRegionIds]),
        },
      })

      const sharedRunners = runners.filter((runner) => sharedRegionIds.includes(runner.region))
      const sharedRunnerIds = sharedRunners.map((runner) => runner.id)

      const organizationRunners = runners.filter((runner) => organizationRegionIds.includes(runner.region))
      const organizationRunnerIds = organizationRunners.map((runner) => runner.id)

      //  get all runners where the snapshot is already propagated to (or in progress)
      const sharedSnapshotRunners = await this.snapshotRunnerRepository.find({
        where: {
          snapshotRef: snapshot.ref,
          state: In([SnapshotRunnerState.READY, SnapshotRunnerState.PULLING_SNAPSHOT]),
          runnerId: In(sharedRunnerIds),
        },
      })
      const sharedSnapshotRunnersDistinctRunnersIds = new Set(
        sharedSnapshotRunners.map((snapshotRunner) => snapshotRunner.runnerId),
      )

      const organizationSnapshotRunners = await this.snapshotRunnerRepository.find({
        where: {
          snapshotRef: snapshot.ref,
          state: In([SnapshotRunnerState.READY, SnapshotRunnerState.PULLING_SNAPSHOT]),
          runnerId: In(organizationRunnerIds),
        },
      })
      const organizationSnapshotRunnersDistinctRunnersIds = new Set(
        organizationSnapshotRunners.map((snapshotRunner) => snapshotRunner.runnerId),
      )

      //  get all runners where the snapshot is not propagated to
      const unallocatedSharedRunners = sharedRunners.filter(
        (runner) => !sharedSnapshotRunnersDistinctRunnersIds.has(runner.id),
      )
      const unallocatedOrganizationRunners = organizationRunners.filter(
        (runner) => !organizationSnapshotRunnersDistinctRunnersIds.has(runner.id),
      )

      const runnersToPropagateTo: Runner[] = []

      // propagate the snapshot to all organization runners
      runnersToPropagateTo.push(...unallocatedOrganizationRunners)

      // respect the propagation limit for shared runners
      const sharedRunnersPropagateLimit = Math.max(
        0,
        Math.ceil(sharedRunners.length * BASE_PROPAGATION_FACTOR) - sharedSnapshotRunnersDistinctRunnersIds.size,
      )
      runnersToPropagateTo.push(...shuffleArray(unallocatedSharedRunners).slice(0, sharedRunnersPropagateLimit))

      if (runnersToPropagateTo.length === 0) {
        return
      }

      // regionId -> registry
      const internalRegistriesMap = new Map<string, DockerRegistry>()

      for (const regionId of [...sharedRegionIds, ...organizationRegionIds]) {
        const registry = await this.dockerRegistryService.findInternalRegistryBySnapshotRef(snapshot.ref, regionId)
        if (registry) {
          internalRegistriesMap.set(regionId, registry)
        }
      }

      const results = await Promise.allSettled(
        runnersToPropagateTo.map(async (runner) => {
          const internalRegistry = internalRegistriesMap.get(runner.region)
          if (!internalRegistry) {
            throw new Error(`No internal registry found for snapshot ${snapshot.ref} in region ${runner.region}`)
          }

          const snapshotRunner = await this.runnerService.getSnapshotRunner(runner.id, snapshot.ref)

          try {
            if (!snapshotRunner) {
              await this.runnerService.createSnapshotRunnerEntry(
                runner.id,
                snapshot.ref,
                SnapshotRunnerState.PULLING_SNAPSHOT,
              )
              await this.pullSnapshotRunner(runner, snapshot.ref, internalRegistry)
            } else if (snapshotRunner.state === SnapshotRunnerState.PULLING_SNAPSHOT) {
              await this.handleSnapshotRunnerStatePullingSnapshot(snapshotRunner, runner)
            }
          } catch (err) {
            this.logger.error(`Error propagating snapshot to runner ${runner.id}: ${fromAxiosError(err)}`)
            snapshotRunner.state = SnapshotRunnerState.ERROR
            snapshotRunner.errorReason = err.message
            await this.snapshotRunnerRepository.update(snapshotRunner.id, snapshotRunner)
          }
        }),
      )

      results.forEach((result) => {
        if (result.status === 'rejected') {
          this.logger.error(result.reason)
        }
      })
    } catch (err) {
      this.logger.error(err)
    }
  }

  async scaleDownSnapshotFromRunners(
    snapshot: Snapshot,
    sharedRegionIds: string[],
    propagationFactor: number = BASE_PROPAGATION_FACTOR,
  ): Promise<number> {
    try {
      if (sharedRegionIds.length === 0) {
        return 0
      }

      // Get all shared runners in the regions
      const sharedRunners = await this.runnerRepository.find({
        where: {
          state: RunnerState.READY,
          unschedulable: Not(true),
          region: In(sharedRegionIds),
        },
      })

      const sharedRunnerIds = sharedRunners.map((runner) => runner.id)

      if (sharedRunnerIds.length === 0) {
        return 0
      }

      // Get all snapshot runners in READY state for shared runners (excluding those created in the last hour)
      const oneHourAgo = new Date(Date.now() - 60 * 60 * 1000)
      const sharedSnapshotRunners = await this.snapshotRunnerRepository.find({
        where: {
          snapshotRef: snapshot.ref,
          state: SnapshotRunnerState.READY,
          runnerId: In(sharedRunnerIds),
          createdAt: LessThan(oneHourAgo),
        },
      })

      // Calculate the maximum allowed number of snapshot runners (same formula as propagation)
      const maxSharedSnapshotRunners = Math.ceil(propagationFactor * sharedRunners.length)

      // Only scale down if the propagated amount exceeds the limit by more than 15%
      const scaleDownThreshold = Math.ceil(maxSharedSnapshotRunners * 1.15)
      const excessCount = sharedSnapshotRunners.length - maxSharedSnapshotRunners

      if (sharedSnapshotRunners.length > scaleDownThreshold) {
        // Sort by createdAt ascending (oldest first) to remove oldest ones
        const sortedSnapshotRunners = sharedSnapshotRunners.sort(
          (a, b) => a.createdAt.getTime() - b.createdAt.getTime(),
        )

        // Take the excess ones to remove
        const snapshotRunnersToRemove = sortedSnapshotRunners.slice(0, excessCount)

        await Promise.allSettled(
          snapshotRunnersToRemove.map(async (snapshotRunner) => {
            await this.snapshotRunnerRepository.update(snapshotRunner.id, {
              state: SnapshotRunnerState.REMOVING,
            })
          }),
        )

        this.logger.log(
          `Marked ${snapshotRunnersToRemove.length} snapshot runners for removal for snapshot ${snapshot.ref}`,
        )

        return snapshotRunnersToRemove.length
      }

      return 0
    } catch (err) {
      this.logger.error(`Error scaling down snapshot ${snapshot.ref}: ${fromAxiosError(err)}`)
      return 0
    }
  }

  async pullSnapshotRunner(
    runner: Runner,
    snapshotRef: string,
    registry?: DockerRegistry,
    destinationRegistry?: DockerRegistry,
    destinationRef?: string,
  ) {
    const runnerAdapter = await this.runnerAdapterFactory.create(runner)
    // Runner returns immediately; polling for completion is handled by syncRunnerSnapshotStates cron
    await runnerAdapter.pullSnapshot(snapshotRef, registry, destinationRegistry, destinationRef)
  }

  async handleSnapshotRunnerStatePullingSnapshot(snapshotRunner: SnapshotRunner, runner: Runner) {
    const runnerAdapter = await this.runnerAdapterFactory.create(runner)
    try {
      await runnerAdapter.getSnapshotInfo(snapshotRunner.snapshotRef)
      snapshotRunner.state = SnapshotRunnerState.READY
      await this.snapshotRunnerRepository.save(snapshotRunner)
      return
    } catch (err) {
      if (err instanceof SnapshotStateError) {
        snapshotRunner.state = SnapshotRunnerState.ERROR
        snapshotRunner.errorReason = err.errorReason
        await this.snapshotRunnerRepository.save(snapshotRunner)
        return
      }
    }

    const timeoutMinutes = 60
    const timeoutMs = timeoutMinutes * 60 * 1000
    if (Date.now() - snapshotRunner.updatedAt.getTime() > timeoutMs) {
      snapshotRunner.state = SnapshotRunnerState.ERROR
      snapshotRunner.errorReason = 'Timeout while pulling snapshot to runner'
      await this.snapshotRunnerRepository.save(snapshotRunner)
      return
    }

    const retryTimeoutMinutes = 10
    const retryTimeoutMs = retryTimeoutMinutes * 60 * 1000
    if (Date.now() - snapshotRunner.createdAt.getTime() > retryTimeoutMs) {
      const internalRegistry = await this.dockerRegistryService.findInternalRegistryBySnapshotRef(
        snapshotRunner.snapshotRef,
        runner.region,
      )
      if (!internalRegistry) {
        throw new Error(
          `No internal registry found for snapshot ${snapshotRunner.snapshotRef} in region ${runner.region}`,
        )
      }
      await this.pullSnapshotRunner(runner, snapshotRunner.snapshotRef, internalRegistry)
      return
    }
  }

  async handleSnapshotRunnerStateBuildingSnapshot(snapshotRunner: SnapshotRunner, runner: Runner) {
    const runnerAdapter = await this.runnerAdapterFactory.create(runner)
    try {
      await runnerAdapter.getSnapshotInfo(snapshotRunner.snapshotRef)
      snapshotRunner.state = SnapshotRunnerState.READY
      await this.snapshotRunnerRepository.save(snapshotRunner)
      return
    } catch (err) {
      if (err instanceof SnapshotStateError) {
        snapshotRunner.state = SnapshotRunnerState.ERROR
        snapshotRunner.errorReason = err.errorReason
        await this.snapshotRunnerRepository.save(snapshotRunner)
        return
      }
    }

    const timeoutMinutes = 120
    const timeoutMs = timeoutMinutes * 60 * 1000
    if (Date.now() - snapshotRunner.updatedAt.getTime() > timeoutMs) {
      snapshotRunner.state = SnapshotRunnerState.ERROR
      snapshotRunner.errorReason = 'Timeout while building snapshot on runner'
      await this.snapshotRunnerRepository.save(snapshotRunner)
    }
  }

  // Pulls stopped sandboxes' backup snapshots to another runner to prepare for reassignment during draining
  @Cron(CronExpression.EVERY_10_SECONDS, { name: 'migrate-draining-runner-snapshots', waitForCompletion: true })
  @TrackJobExecution()
  @LogExecution('migrate-draining-runner-snapshots')
  @WithInstrumentation()
  private async handleMigrateDrainingRunnerSnapshots() {
    const lockKey = 'migrate-draining-runner-snapshots'
    const hasLock = await this.redisLockProvider.lock(lockKey, 60)
    if (!hasLock) {
      return
    }

    try {
      const drainingRunners = await this.runnerRepository.find({
        where: {
          draining: true,
          state: RunnerState.READY,
        },
      })

      this.logger.debug(`Checking ${drainingRunners.length} draining runners for snapshot migration`)

      await Promise.allSettled(
        drainingRunners.map(async (runner) => {
          try {
            const sandboxes = await this.sandboxRepository.find({
              where: {
                runnerId: runner.id,
                state: SandboxState.STOPPED,
                desiredState: SandboxDesiredState.STOPPED,
                backupState: BackupState.COMPLETED,
                backupSnapshot: Not(IsNull()),
              },
              take: 100,
            })

            this.logger.debug(
              `Found ${sandboxes.length} eligible sandboxes on draining runner ${runner.id} for snapshot migration`,
            )

            await Promise.allSettled(
              sandboxes.map(async (sandbox) => {
                const sandboxLockKey = `draining-runner-snapshot-migration:${sandbox.id}`
                const hasSandboxLock = await this.redisLockProvider.lock(sandboxLockKey, 3600)
                if (!hasSandboxLock) {
                  return
                }

                try {
                  // Get an available runner in the same region with the same class
                  const targetRunner = await this.runnerService.getRandomAvailableRunner({
                    regions: [sandbox.region],
                    sandboxClass: sandbox.class,
                    excludedRunnerIds: [runner.id],
                  })

                  // Check if snapshot runner entry already exists
                  const existingEntry = await this.runnerService.getSnapshotRunner(
                    targetRunner.id,
                    sandbox.backupSnapshot,
                  )
                  if (existingEntry) {
                    if (existingEntry.state === SnapshotRunnerState.ERROR) {
                      // Clean up the failed entry so we can retry
                      this.logger.warn(
                        `Removing ERROR snapshot runner entry ${existingEntry.id} for runner ${targetRunner.id} and snapshot ${sandbox.backupSnapshot} to allow retry`,
                      )
                      await this.snapshotRunnerRepository.delete(existingEntry.id)
                    } else {
                      this.logger.debug(
                        `Snapshot runner entry already exists for runner ${targetRunner.id} and snapshot ${sandbox.backupSnapshot} (state: ${existingEntry.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      return
                    }
                  }

                  // Find the backup registry to use as source for the pull
                  const registry = sandbox.backupRegistryId
                    ? await this.dockerRegistryService.findOne(sandbox.backupRegistryId)
                    : await this.dockerRegistryService.findInternalRegistryBySnapshotRef(
                        sandbox.backupSnapshot,
                        targetRunner.region,
                      )

                  if (!registry) {
                    this.logger.warn(
                      `No registry found for backup snapshot ${sandbox.backupSnapshot} of sandbox ${sandbox.id}`,
                    )
                    await this.redisLockProvider.unlock(sandboxLockKey)
                    return
                  }

                  // Create snapshot runner entry on the target runner
                  await this.runnerService.createSnapshotRunnerEntry(
                    targetRunner.id,
                    sandbox.backupSnapshot,
                    SnapshotRunnerState.PULLING_SNAPSHOT,
                  )
                  await this.pullSnapshotRunner(targetRunner, sandbox.backupSnapshot, registry)

                  this.logger.log(
                    `Created snapshot runner entry for sandbox ${sandbox.id} backup ${sandbox.backupSnapshot} on runner ${targetRunner.id} (migrating from draining runner ${runner.id})`,
                  )
                  await this.redisLockProvider.unlock(sandboxLockKey)
                } catch (e) {
                  if (e instanceof BadRequestError && e.message === 'No available runners') {
                    this.logger.warn(
                      `No available runners found in region ${sandbox.region} for sandbox ${sandbox.id} snapshot migration`,
                    )
                  } else {
                    this.logger.error(`Error migrating snapshot for sandbox ${sandbox.id}`, e)
                  }
                  await this.redisLockProvider.unlock(sandboxLockKey)
                }
              }),
            )
          } catch (e) {
            this.logger.error(`Error processing draining runner ${runner.id} for snapshot migration`, e)
          }
        }),
      )
    } finally {
      await this.redisLockProvider.unlock(lockKey)
    }
  }

  @Cron(CronExpression.EVERY_10_SECONDS, { name: 'check-snapshot-cleanup' })
  @TrackJobExecution()
  @LogExecution('check-snapshot-cleanup')
  @WithInstrumentation()
  async checkSnapshotCleanup() {
    const lockKey = 'check-snapshot-cleanup-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 30))) {
      return
    }

    const snapshots = await this.snapshotRepository.find({
      where: {
        state: SnapshotState.REMOVING,
      },
    })

    await Promise.all(
      snapshots.map(async (snapshot) => {
        const countActiveSnapshots = await this.snapshotRepository.count({
          where: {
            state: SnapshotState.ACTIVE,
            ref: snapshot.ref,
          },
        })

        // Only remove snapshot runners if no other snapshots depend on them
        if (countActiveSnapshots === 0) {
          await this.snapshotRunnerRepository.update(
            {
              snapshotRef: snapshot.ref,
            },
            {
              state: SnapshotRunnerState.REMOVING,
            },
          )
        }

        await this.snapshotRepository.remove(snapshot)
      }),
    )

    await this.redisLockProvider.unlock(lockKey)
  }

  @Cron(CronExpression.EVERY_10_SECONDS, { name: 'check-snapshot-state' })
  @TrackJobExecution()
  @LogExecution('check-snapshot-state')
  @WithInstrumentation()
  async checkSnapshotState() {
    //  the first time the snapshot is created it needs to be pushed to the internal registry
    //  before propagating to the runners
    //  this cron job will process the snapshot states until the snapshot is active (or error)

    //  get all snapshots
    const snapshots = await this.snapshotRepository.find({
      where: {
        state: Not(In([SnapshotState.ACTIVE, SnapshotState.ERROR, SnapshotState.BUILD_FAILED, SnapshotState.INACTIVE])),
      },
    })

    await Promise.all(
      snapshots.map(async (snapshot) => {
        this.syncSnapshotState(snapshot.id)
      }),
    )
  }

  async syncSnapshotState(snapshotId: string): Promise<void> {
    const lockKey = `sync-snapshot-state-${snapshotId}`
    if (!(await this.redisLockProvider.lock(lockKey, 720))) {
      return
    }

    const snapshot = await this.snapshotRepository.findOne({
      where: { id: snapshotId },
    })

    if (
      !snapshot ||
      [SnapshotState.ACTIVE, SnapshotState.ERROR, SnapshotState.BUILD_FAILED, SnapshotState.INACTIVE].includes(
        snapshot.state,
      )
    ) {
      await this.redisLockProvider.unlock(lockKey)
      return
    }

    let syncState = DONT_SYNC_AGAIN

    try {
      switch (snapshot.state) {
        case SnapshotState.PENDING:
          syncState = await this.handleSnapshotStatePending(snapshot)
          break
        case SnapshotState.PULLING:
        case SnapshotState.BUILDING:
          syncState = await this.handleCheckInitialRunnerSnapshot(snapshot)
          break
        case SnapshotState.REMOVING:
          syncState = await this.handleSnapshotStateRemoving(snapshot)
          break
      }
    } catch (error) {
      if (error.code === 'ECONNRESET') {
        syncState = SYNC_AGAIN
      } else {
        const message = error.message || String(error)
        await this.updateSnapshotState(snapshot, SnapshotState.ERROR, message)
      }
    }

    await this.redisLockProvider.unlock(lockKey)
    if (syncState === SYNC_AGAIN) {
      this.syncSnapshotState(snapshotId)
    }
  }

  async handleSnapshotRunnerStateRemoving(snapshotRunner: SnapshotRunner, runner: Runner) {
    if (!runner) {
      //  generally this should not happen
      //  in case the runner has been deleted from the database, delete the snapshot runner record
      const errorMessage = `Runner not found while trying to remove snapshot ${snapshotRunner.snapshotRef} from runner ${snapshotRunner.runnerId}`
      this.logger.warn(errorMessage)

      this.snapshotRunnerRepository.delete(snapshotRunner.id).catch((err) => {
        this.logger.error(fromAxiosError(err))
      })
      return
    }
    if (!snapshotRunner.snapshotRef) {
      //  this should never happen
      //  remove the snapshot runner record (it will be recreated again by the snapshot propagation job)
      this.logger.warn(`Internal snapshot name not found for snapshot runner ${snapshotRunner.id}`)
      this.snapshotRunnerRepository.delete(snapshotRunner.id).catch((err) => {
        this.logger.error(fromAxiosError(err))
      })
      return
    }

    const runnerAdapter = await this.runnerAdapterFactory.create(runner)
    const exists = await runnerAdapter.snapshotExists(snapshotRunner.snapshotRef)
    if (!exists) {
      await this.snapshotRunnerRepository.delete(snapshotRunner.id)
    } else {
      //  just in case the snapshot is still there
      runnerAdapter.removeSnapshot(snapshotRunner.snapshotRef).catch((err) => {
        //  this should not happen, and is not critical
        //  if the runner can not remove the snapshot, just delete the snapshot runner record
        this.snapshotRunnerRepository.delete(snapshotRunner.id).catch((err) => {
          this.logger.error(fromAxiosError(err))
        })
        //  and log the error for tracking
        const errorMessage = `Failed to do just in case remove snapshot ${snapshotRunner.snapshotRef} from runner ${runner.id}: ${fromAxiosError(err)}`
        this.logger.warn(errorMessage)
      })
    }
  }

  async handleSnapshotStateRemoving(snapshot: Snapshot): Promise<SyncState> {
    const snapshotRunnerItems = await this.snapshotRunnerRepository.find({
      where: {
        snapshotRef: snapshot.ref,
      },
    })

    if (snapshotRunnerItems.length === 0) {
      await this.snapshotRepository.remove(snapshot)
    }

    return DONT_SYNC_AGAIN
  }

  async handleCheckInitialRunnerSnapshot(snapshot: Snapshot): Promise<SyncState> {
    // Check for timeout - allow up to 60 minutes
    const timeoutMinutes = 60
    const timeoutMs = timeoutMinutes * 60 * 1000
    if (Date.now() - snapshot.updatedAt.getTime() > timeoutMs) {
      await this.updateSnapshotState(snapshot, SnapshotState.ERROR, 'Timeout processing snapshot on initial runner')
      return DONT_SYNC_AGAIN
    }

    // Check if the snapshot ref is already set and it is already on the runner
    const snapshotRunner = await this.snapshotRunnerRepository.findOne({
      where: {
        snapshotRef: snapshot.ref,
        runnerId: snapshot.initialRunnerId,
      },
    })

    if (snapshot.ref && snapshotRunner) {
      if (snapshotRunner.state === SnapshotRunnerState.READY && snapshot.size != null) {
        await this.updateSnapshotState(snapshot, SnapshotState.ACTIVE)
        return DONT_SYNC_AGAIN
      } else if (snapshotRunner.state === SnapshotRunnerState.ERROR) {
        await this.snapshotRunnerRepository.delete(snapshotRunner.id)
      }
    }

    const runner = await this.runnerService.findOneOrFail(snapshot.initialRunnerId)
    const runnerAdapter = await this.runnerAdapterFactory.create(runner)

    const initialImageRefOnRunner = snapshot.buildInfo ? snapshot.buildInfo.snapshotRef : snapshot.ref

    let snapshotInfoResponse: SnapshotInfoResponse
    try {
      snapshotInfoResponse = await runnerAdapter.getSnapshotInfo(initialImageRefOnRunner)
    } catch (error) {
      if (error instanceof SnapshotStateError) {
        throw error
      } else {
        return DONT_SYNC_AGAIN
      }
    }

    const internalRegistry = await this.dockerRegistryService.getAvailableInternalRegistry(runner.region)
    if (!internalRegistry) {
      throw new Error('No internal registry found for snapshot')
    }

    const digestSyncState = await this.processSnapshotDigest(
      snapshot,
      internalRegistry,
      snapshotInfoResponse.hash,
      snapshotInfoResponse.sizeGB,
      snapshotInfoResponse.entrypoint,
    )

    if (digestSyncState === DONT_SYNC_AGAIN) {
      return DONT_SYNC_AGAIN
    }

    let inspectedSnapshotDigest: SnapshotDigestResponse | undefined
    try {
      inspectedSnapshotDigest = await runnerAdapter.inspectSnapshotInRegistry(snapshot.ref, internalRegistry)
    } catch (error) {
      this.logger.error(`Failed to inspect snapshot ${snapshot.ref} in registry: ${error}`)
      return DONT_SYNC_AGAIN
    }

    if (snapshot.size == null && typeof inspectedSnapshotDigest?.sizeGB === 'number') {
      await this.processSnapshotDigest(
        snapshot,
        internalRegistry,
        snapshotInfoResponse.hash,
        inspectedSnapshotDigest.sizeGB,
        snapshotInfoResponse.entrypoint,
      )
    }

    try {
      await runnerAdapter.removeSnapshot(initialImageRefOnRunner)
    } catch (error) {
      this.logger.error(`Failed to remove snapshot ${snapshot.imageName}: ${fromAxiosError(error)}`)
    }

    // For pull snapshots, best effort cleanup the original image now that we've computed the ref from it
    // Only cleanup if there's no other snapshot in processing state using the same image
    if (!snapshot.buildInfo) {
      try {
        const anotherSnapshot = await this.snapshotRepository.findOne({
          where: {
            imageName: snapshot.imageName,
            id: Not(snapshot.id),
            state: Not(In([SnapshotState.ACTIVE, SnapshotState.INACTIVE])),
          },
        })
        if (!anotherSnapshot) {
          await runnerAdapter.removeSnapshot(snapshot.imageName)
        }
      } catch (err) {
        this.logger.error(`Failed to cleanup original image ${snapshot.imageName}: ${fromAxiosError(err)}`)
      }
    }

    if (snapshotRunner) {
      snapshotRunner.state = SnapshotRunnerState.READY
      await this.snapshotRunnerRepository.save(snapshotRunner)
    } else {
      await this.runnerService.createSnapshotRunnerEntry(runner.id, snapshot.ref, SnapshotRunnerState.READY)
    }

    if (snapshot.size == null) {
      this.logger.warn(`Snapshot ${snapshot.id} has no resolved size yet; deferring activation`)
      return DONT_SYNC_AGAIN
    }

    await this.updateSnapshotState(snapshot, SnapshotState.ACTIVE)

    // Best effort removal of old snapshot from transient registry
    const transientRegistry = await this.dockerRegistryService.findTransientRegistryBySnapshotImageName(
      snapshot.imageName,
      runner.region,
    )
    if (transientRegistry) {
      try {
        await this.dockerRegistryService.removeImage(snapshot.imageName, transientRegistry.id)
      } catch (error) {
        if (error.statusCode === 404) {
          //  image not found, just return
          return DONT_SYNC_AGAIN
        }
        this.logger.error('Failed to remove transient image:', fromAxiosError(error))
      }
    }

    return DONT_SYNC_AGAIN
  }

  async processPullOnInitialRunner(snapshot: Snapshot, runner: Runner) {
    // Check for timeout - allow up to 60 minutes
    const timeoutMinutes = 60
    const timeoutMs = timeoutMinutes * 60 * 1000
    if (Date.now() - snapshot.updatedAt.getTime() > timeoutMs) {
      await this.updateSnapshotState(
        snapshot,
        SnapshotState.ERROR,
        'Timeout processing snapshot pull on initial runner',
      )
      return DONT_SYNC_AGAIN
    }

    let sourceRegistry = await this.dockerRegistryService.findSourceRegistryBySnapshotImageName(
      snapshot.imageName,
      runner.region,
      snapshot.organizationId,
    )
    if (!sourceRegistry) {
      sourceRegistry = await this.dockerRegistryService.getDefaultDockerHubRegistry()
    }
    const destinationRegistry = await this.dockerRegistryService.getAvailableInternalRegistry(runner.region)

    // Fire pull request (runner returns 202 immediately)
    // Post-processing (digest, cleanup) is handled by handleCheckInitialRunnerSnapshot on the next poll cycle
    try {
      await this.pullSnapshotRunner(
        runner,
        snapshot.imageName,
        sourceRegistry,
        destinationRegistry ?? undefined,
        snapshot.ref ? snapshot.ref : undefined,
      )
    } catch (err) {
      // Validation errors are still returned synchronously
      await this.updateSnapshotState(snapshot, SnapshotState.ERROR, err.message)
      throw err
    }
  }

  async processBuildOnRunner(snapshot: Snapshot, runner: Runner) {
    try {
      const registry = await this.dockerRegistryService.getAvailableInternalRegistry(runner.region)

      const sourceRegistries = await this.dockerRegistryService.getSourceRegistriesForDockerfile(
        snapshot.buildInfo.dockerfileContent,
        snapshot.organizationId,
      )

      const runnerAdapter = await this.runnerAdapterFactory.create(runner)

      registry.url = registry.url.replace(/^(https?:\/\/)/, '')
      // Runner returns immediately; polling for completion is handled by handleCheckInitialRunnerSnapshot
      await runnerAdapter.buildSnapshot(
        snapshot.buildInfo,
        snapshot.organizationId,
        sourceRegistries.length > 0 ? sourceRegistries : undefined,
        registry ?? undefined,
        true,
      )
    } catch (err) {
      this.logger.error(`Error building snapshot ${snapshot.name}: ${fromAxiosError(err)}`)
      await this.updateSnapshotState(snapshot, SnapshotState.BUILD_FAILED, fromAxiosError(err).message)
    }
  }

  async handleSnapshotStatePending(snapshot: Snapshot): Promise<SyncState> {
    let initialRunner: Runner | undefined = undefined

    if (!snapshot.initialRunnerId) {
      // TODO: get only runners where the base snapshot is available (extract from buildInfo)
      const excludedRunnerIds = snapshot.buildInfo
        ? await this.runnerService.getRunnersWithMultipleSnapshotsBuilding()
        : await this.runnerService.getRunnersWithMultipleSnapshotsPulling()

      try {
        const regions = await this.snapshotService.getSnapshotRegions(snapshot.id)
        if (!regions.length) {
          throw new Error('No regions found for snapshot')
        }

        const availabilityThreshold =
          this.configService.getOrThrow('runnerScore.thresholds.availability') +
          this.configService.getOrThrow('runnerScore.thresholds.initialRunnerScoreAddon')

        initialRunner = await this.runnerService.getRandomAvailableRunner({
          regions: regions.map((region) => region.id),
          excludedRunnerIds: excludedRunnerIds,
          availabilityScoreThreshold: availabilityThreshold,
        })
      } catch (error) {
        this.logger.warn(`Failed to get initial runner: ${fromAxiosError(error)}`)
      }

      if (!initialRunner) {
        // No runners available, retry later
        return DONT_SYNC_AGAIN
      }

      const updateData: Partial<Snapshot> = {
        initialRunnerId: initialRunner.id,
      }

      await this.snapshotRepository.update(snapshot.id, { updateData, entity: snapshot })
    } else {
      initialRunner = await this.runnerService.findOneOrFail(snapshot.initialRunnerId)
    }

    if (snapshot.buildInfo) {
      await this.updateSnapshotState(snapshot, SnapshotState.BUILDING)
      await this.runnerService.createSnapshotRunnerEntry(
        initialRunner.id,
        snapshot.buildInfo.snapshotRef,
        SnapshotRunnerState.BUILDING_SNAPSHOT,
      )
      await this.processBuildOnRunner(snapshot, initialRunner)
    } else {
      if (!snapshot.ref) {
        const runnerAdapter = await this.runnerAdapterFactory.create(initialRunner)
        const registry = await this.dockerRegistryService.findRegistryByImageName(
          snapshot.imageName,
          initialRunner.region,
          snapshot.organizationId,
        )

        const image = parseDockerImage(snapshot.imageName)
        if (registry && !image.registry) {
          image.registry = registry.url.replace(/^(https?:\/\/)/, '')
        }
        const imageName = image.getFullName()

        const internalRegistry = await this.dockerRegistryService.getAvailableInternalRegistry(initialRunner.region)
        if (!internalRegistry) {
          throw new Error('No internal registry found for snapshot')
        }

        const snapshotDigestResponse = await runnerAdapter.inspectSnapshotInRegistry(imageName, registry)
        const digestSyncState = await this.processSnapshotDigest(
          snapshot,
          internalRegistry,
          snapshotDigestResponse.hash,
          snapshotDigestResponse.sizeGB,
        )

        if (digestSyncState === DONT_SYNC_AGAIN) {
          return DONT_SYNC_AGAIN
        }
      }

      await this.updateSnapshotState(snapshot, SnapshotState.PULLING)
      await this.runnerService.createSnapshotRunnerEntry(
        initialRunner.id,
        snapshot.ref,
        SnapshotRunnerState.PULLING_SNAPSHOT,
      )
      await this.processPullOnInitialRunner(snapshot, initialRunner)
    }

    return SYNC_AGAIN
  }

  private async updateSnapshotState(snapshot: Snapshot, state: SnapshotState, errorReason?: string, size?: number) {
    const updateData: Partial<Snapshot> = {
      state,
    }

    if (state === SnapshotState.ACTIVE) {
      updateData.lastUsedAt = new Date()
    }

    if (errorReason !== undefined) {
      updateData.errorReason = errorReason
    }

    if (size !== undefined) {
      updateData.size = size
    }

    await this.snapshotRepository.update(snapshot.id, { updateData, entity: snapshot })
  }

  @Cron(CronExpression.EVERY_MINUTE, { name: 'cleanup-old-buildinfo-snapshot-runners' })
  @TrackJobExecution()
  @LogExecution('cleanup-old-buildinfo-snapshot-runners')
  @WithInstrumentation()
  async cleanupOldBuildInfoSnapshotRunners() {
    const lockKey = 'cleanup-old-buildinfo-snapshots-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 300))) {
      return
    }

    try {
      const stalenessDays = this.configService.getOrThrow('buildInfoSnapshotRunnerStalenessDays')
      const stalenessInterval = `interval '${stalenessDays} days'`

      const staleEntries = await this.snapshotRunnerRepository
        .createQueryBuilder('sr')
        .select('sr.id')
        .innerJoin(BuildInfo, 'bi', 'sr."snapshotRef" = bi."snapshotRef"')
        .innerJoin('runner', 'r', 'r.id = sr."runnerId"::uuid')
        .where('sr.state = :readyState', { readyState: SnapshotRunnerState.READY })
        .andWhere(`bi.lastUsedAt < now() - ${stalenessInterval}`)
        .andWhere(`sr.updatedAt < now() - ${stalenessInterval}`)
        .andWhere("sr.snapshotRef LIKE 'daytona-%'")
        .limit(500)
        .getMany()

      if (staleEntries.length === 0) {
        return
      }

      const ids = staleEntries.map((sr) => sr.id)
      const result = await this.snapshotRunnerRepository.update(
        { id: In(ids) },
        { state: SnapshotRunnerState.REMOVING },
      )

      if (result.affected > 0) {
        this.logger.debug(`Marked ${result.affected} SnapshotRunners for removal due to unused BuildInfo`)
      }
    } catch (error) {
      this.logger.error(`Failed to mark old BuildInfo SnapshotRunners for removal: ${fromAxiosError(error)}`)
    } finally {
      await this.redisLockProvider.unlock(lockKey)
    }
  }

  @Cron(CronExpression.EVERY_10_MINUTES, { name: 'deactivate-old-snapshots' })
  @TrackJobExecution()
  @LogExecution('deactivate-old-snapshots')
  @WithInstrumentation()
  async deactivateOldSnapshots() {
    const lockKey = 'deactivate-old-snapshots-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 300))) {
      return
    }

    try {
      const cutoff = `NOW() - INTERVAL '1 minute' * COALESCE(org."snapshot_deactivation_timeout_minutes", ${DEFAULT_SNAPSHOT_DEACTIVATION_TIMEOUT_MINUTES})`

      const oldSnapshots = await this.snapshotRepository
        .createQueryBuilder('snapshot')
        .leftJoin('organization', 'org', `org."id" = snapshot."organizationId"`)
        .where('snapshot.general = false')
        .andWhere('snapshot.state = :snapshotState', { snapshotState: SnapshotState.ACTIVE })
        .andWhere(`(snapshot."lastUsedAt" IS NULL OR snapshot."lastUsedAt" < ${cutoff})`)
        .andWhere(`snapshot."createdAt" < ${cutoff}`)
        .andWhere(
          `NOT EXISTS (
            SELECT 1 FROM snapshot s
            WHERE s."ref" = snapshot."ref"
            AND s.state = :activeState
            AND (s."lastUsedAt" >= ${cutoff} OR s."createdAt" >= ${cutoff})
          )`,
          {
            activeState: SnapshotState.ACTIVE,
          },
        )
        .take(100)
        .getMany()

      if (oldSnapshots.length === 0) {
        return
      }

      // Deactivate the snapshots
      const settledResults = await Promise.allSettled(
        oldSnapshots.map((snapshot) =>
          this.snapshotRepository.update(snapshot.id, {
            updateData: { state: SnapshotState.INACTIVE },
            entity: snapshot,
          }),
        ),
      )

      const deactivatedSnapshots: Snapshot[] = []
      for (const [i, result] of settledResults.entries()) {
        if (result.status === 'fulfilled') {
          deactivatedSnapshots.push(oldSnapshots[i])
        } else {
          this.logger.warn(`Failed to deactivate snapshot ${oldSnapshots[i].id}: ${result.reason}`)
        }
      }

      // Get internal names of deactivated snapshots
      const refs = deactivatedSnapshots.map((snapshot) => snapshot.ref).filter((name) => name) // Filter out null/undefined values

      if (refs.length > 0) {
        // Set associated SnapshotRunner records to REMOVING state
        const result = await this.snapshotRunnerRepository.update(
          { snapshotRef: In(refs) },
          { state: SnapshotRunnerState.REMOVING },
        )

        this.logger.debug(
          `Deactivated ${deactivatedSnapshots.length} snapshots and marked ${result.affected} SnapshotRunners for removal`,
        )
      }
    } catch (error) {
      this.logger.error(`Failed to deactivate old snapshots: ${fromAxiosError(error)}`)
    } finally {
      await this.redisLockProvider.unlock(lockKey)
    }
  }

  @Cron(CronExpression.EVERY_MINUTE, { name: 'cleanup-inactive-snapshots-from-runners' })
  @TrackJobExecution()
  @LogExecution('cleanup-inactive-snapshots-from-runners')
  @WithInstrumentation()
  async cleanupInactiveSnapshotsFromRunners() {
    const lockKey = 'cleanup-inactive-snapshots-from-runners-lock'
    if (!(await this.redisLockProvider.lock(lockKey, 300))) {
      return
    }

    try {
      // Only fetch inactive snapshots that have associated snapshot runner entries
      const queryResult = await this.snapshotRepository
        .createQueryBuilder('snapshot')
        .select('snapshot."ref"')
        .where('snapshot.state = :inactiveState', { inactiveState: SnapshotState.INACTIVE })
        .andWhere('snapshot."ref" IS NOT NULL')
        .andWhereExists(
          this.snapshotRunnerRepository
            .createQueryBuilder('snapshot_runner')
            .select('1')
            .where('snapshot_runner."snapshotRef" = snapshot."ref"')
            .andWhere('snapshot_runner.state != :snapshotRunnerState', {
              snapshotRunnerState: SnapshotRunnerState.REMOVING,
            }),
        )
        .andWhere(
          () => {
            const query = this.snapshotRepository
              .createQueryBuilder('s')
              .select('1')
              .where('s."ref" = snapshot."ref"')
              .andWhere('s.state = :activeState')
            return `NOT EXISTS (${query.getQuery()})`
          },
          {
            activeState: SnapshotState.ACTIVE,
          },
        )
        .orderBy('snapshot."updatedAt"', 'DESC')
        .take(100)
        .getRawMany()

      const inactiveSnapshotRefs = queryResult.map((result) => result.ref)

      if (inactiveSnapshotRefs.length > 0) {
        // Set associated SnapshotRunner records to REMOVING state
        const result = await this.snapshotRunnerRepository.update(
          { snapshotRef: In(inactiveSnapshotRefs) },
          { state: SnapshotRunnerState.REMOVING },
        )

        this.logger.debug(`Marked ${result.affected} SnapshotRunners for removal`)
      }
    } catch (error) {
      this.logger.error(`Failed to cleanup inactive snapshots from runners: ${fromAxiosError(error)}`)
    } finally {
      await this.redisLockProvider.unlock(lockKey)
    }
  }

  private async processSnapshotDigest(
    snapshot: Snapshot,
    internalRegistry: DockerRegistry,
    hash: string,
    sizeGB?: number,
    entrypoint?: string[] | string,
  ) {
    const updateData: Partial<Snapshot> = {}

    if (!snapshot.ref) {
      const sanitizedUrl = internalRegistry.url.replace(/^https?:\/\//, '')
      updateData.ref = `${sanitizedUrl}/${internalRegistry.project || 'daytona'}/daytona-${hash}:daytona`
    }

    if (snapshot.size == null && typeof sizeGB === 'number') {
      const organization = await this.organizationService.findOne(snapshot.organizationId)
      if (!organization) {
        throw new NotFoundException(`Organization with ID ${snapshot.organizationId} not found`)
      }

      const MAX_SIZE_GB = organization.maxSnapshotSize

      updateData.size = sizeGB

      if (sizeGB > MAX_SIZE_GB) {
        await this.updateSnapshotState(
          snapshot,
          SnapshotState.ERROR,
          `Snapshot size (${sizeGB.toFixed(2)}GB) exceeds maximum allowed size of ${MAX_SIZE_GB}GB`,
          sizeGB,
        )
        return DONT_SYNC_AGAIN
      }
    }

    // If entrypoint is not explicitly set, set it from snapshotInfoResponse
    if (!snapshot.entrypoint) {
      if (entrypoint && entrypoint.length > 0) {
        updateData.entrypoint = Array.isArray(entrypoint) ? entrypoint : [entrypoint]
      }
    }

    if (Object.keys(updateData).length > 0) {
      await this.snapshotRepository.update(snapshot.id, { updateData, entity: snapshot })
    }
  }

  @OnAsyncEvent({
    event: SnapshotEvents.CREATED,
  })
  private async handleSnapshotCreatedEvent(event: SnapshotCreatedEvent) {
    await this.syncSnapshotState(event.snapshot.id)
  }

  @OnAsyncEvent({
    event: SnapshotEvents.ACTIVATED,
  })
  private async handleSnapshotActivatedEvent(event: SnapshotActivatedEvent) {
    await this.syncSnapshotState(event.snapshot.id)
  }
}
