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

import { Injectable, Logger, OnApplicationShutdown } from '@nestjs/common'
import { InjectRepository } from '@nestjs/typeorm'
import { IsNull, LessThan, Not, Repository } from 'typeorm'
import { SandboxUsagePeriod } from '../entities/sandbox-usage-period.entity'
import { OnEvent } from '@nestjs/event-emitter'
import { SandboxStateUpdatedEvent } from '../../sandbox/events/sandbox-state-updated.event'
import { SandboxState } from '../../sandbox/enums/sandbox-state.enum'
import { SandboxEvents } from './../../sandbox/constants/sandbox-events.constants'
import { Cron, CronExpression } from '@nestjs/schedule'
import { RedisLockProvider } from '../../sandbox/common/redis-lock.provider'
import { SANDBOX_WARM_POOL_UNASSIGNED_ORGANIZATION } from '../../sandbox/constants/sandbox.constants'
import { SandboxUsagePeriodArchive } from '../entities/sandbox-usage-period-archive.entity'
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 { SandboxRepository } from '../../sandbox/repositories/sandbox.repository'

@Injectable()
export class UsageService implements TrackableJobExecutions, OnApplicationShutdown {
  activeJobs = new Set<string>()
  private readonly logger = new Logger(UsageService.name)

  constructor(
    @InjectRepository(SandboxUsagePeriod)
    private sandboxUsagePeriodRepository: Repository<SandboxUsagePeriod>,
    private readonly redisLockProvider: RedisLockProvider,
    private readonly sandboxRepository: SandboxRepository,
  ) {}

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

  @OnEvent(SandboxEvents.STATE_UPDATED)
  @TrackJobExecution()
  async handleSandboxStateUpdate(event: SandboxStateUpdatedEvent) {
    await this.waitForLock(event.sandbox.id)

    try {
      switch (event.newState) {
        case SandboxState.STARTED: {
          await this.closeUsagePeriod(event.sandbox.id)
          await this.createUsagePeriod(event)
          break
        }
        case SandboxState.STOPPING:
          await this.closeUsagePeriod(event.sandbox.id)
          await this.createUsagePeriod(event, true)
          break
        case SandboxState.ERROR:
        case SandboxState.BUILD_FAILED:
        case SandboxState.ARCHIVED:
        case SandboxState.DESTROYING:
        case SandboxState.DESTROYED: {
          await this.closeUsagePeriod(event.sandbox.id)
          break
        }
      }
    } finally {
      this.releaseLock(event.sandbox.id).catch((error) => {
        this.logger.error(`Error releasing lock for sandbox ${event.sandbox.id}`, error)
      })
    }
  }

  private async createUsagePeriod(event: SandboxStateUpdatedEvent, diskOnly = false) {
    const usagePeriod = new SandboxUsagePeriod()
    usagePeriod.sandboxId = event.sandbox.id
    usagePeriod.startAt = new Date()
    usagePeriod.endAt = null
    if (!diskOnly) {
      usagePeriod.cpu = event.sandbox.cpu
      usagePeriod.gpu = event.sandbox.gpu
      usagePeriod.mem = event.sandbox.mem
    } else {
      usagePeriod.cpu = 0
      usagePeriod.gpu = 0
      usagePeriod.mem = 0
    }
    usagePeriod.disk = event.sandbox.disk
    usagePeriod.organizationId = event.sandbox.organizationId
    usagePeriod.region = event.sandbox.region

    await this.sandboxUsagePeriodRepository.save(usagePeriod)
  }

  private async closeUsagePeriod(sandboxId: string) {
    const lastUsagePeriod = await this.sandboxUsagePeriodRepository.findOne({
      where: {
        sandboxId,
        endAt: IsNull(),
      },
    })

    if (lastUsagePeriod) {
      lastUsagePeriod.endAt = new Date()
      await this.sandboxUsagePeriodRepository.save(lastUsagePeriod)
    }
  }

  @Cron(CronExpression.EVERY_MINUTE, { name: 'close-and-reopen-usage-periods' })
  @TrackJobExecution()
  @LogExecution('close-and-reopen-usage-periods')
  @WithInstrumentation()
  async closeAndReopenUsagePeriods() {
    if (!(await this.redisLockProvider.lock('close-and-reopen-usage-periods', 60))) {
      return
    }

    const usagePeriods = await this.sandboxUsagePeriodRepository.find({
      where: {
        endAt: IsNull(),
        // 1 day ago
        startAt: LessThan(new Date(Date.now() - 1000 * 60 * 60 * 24)),
        organizationId: Not(SANDBOX_WARM_POOL_UNASSIGNED_ORGANIZATION),
      },
      order: {
        startAt: 'ASC',
      },
      take: 100,
    })

    for (const usagePeriod of usagePeriods) {
      if (!(await this.aquireLock(usagePeriod.sandboxId))) {
        continue
      }

      // validate that the usage period should remain active just in case
      try {
        const sandbox = await this.sandboxRepository.findOne({
          where: {
            id: usagePeriod.sandboxId,
          },
        })

        await this.sandboxUsagePeriodRepository.manager.transaction(async (transactionalEntityManager) => {
          // Close usage period
          const closeTime = new Date()
          usagePeriod.endAt = closeTime
          await transactionalEntityManager.save(usagePeriod)

          if (
            sandbox &&
            (sandbox.state === SandboxState.STARTED ||
              sandbox.state === SandboxState.STOPPED ||
              sandbox.state === SandboxState.STOPPING)
          ) {
            // Create new usage period
            const newUsagePeriod = SandboxUsagePeriod.fromUsagePeriod(usagePeriod)
            newUsagePeriod.startAt = closeTime
            newUsagePeriod.endAt = null
            await transactionalEntityManager.save(newUsagePeriod)
          }
        })
      } catch (error) {
        this.logger.error(`Error closing and reopening usage period ${usagePeriod.sandboxId}`, error)
      } finally {
        await this.releaseLock(usagePeriod.sandboxId)
      }
    }

    await this.redisLockProvider.unlock('close-and-reopen-usage-periods')
  }

  @Cron(CronExpression.EVERY_5_SECONDS, { name: 'archive-usage-periods' })
  @TrackJobExecution()
  @LogExecution('archive-usage-periods')
  @WithInstrumentation()
  async archiveUsagePeriods() {
    const lockKey = 'archive-usage-periods'
    if (!(await this.redisLockProvider.lock(lockKey, 60))) {
      return
    }

    await this.sandboxUsagePeriodRepository.manager.transaction(async (transactionalEntityManager) => {
      const usagePeriods = await transactionalEntityManager.find(SandboxUsagePeriod, {
        where: {
          endAt: Not(IsNull()),
        },
        order: {
          startAt: 'ASC',
        },
        take: 5000,
      })

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

      this.logger.debug(`Found ${usagePeriods.length} usage periods to archive`)

      await transactionalEntityManager.delete(
        SandboxUsagePeriod,
        usagePeriods.map((usagePeriod) => usagePeriod.id),
      )
      await transactionalEntityManager.save(usagePeriods.map(SandboxUsagePeriodArchive.fromUsagePeriod))
    })

    await this.redisLockProvider.unlock(lockKey)
  }

  private async waitForLock(sandboxId: string) {
    while (!(await this.aquireLock(sandboxId))) {
      await new Promise((resolve) => setTimeout(resolve, 500))
    }
  }

  private async aquireLock(sandboxId: string): Promise<boolean> {
    return await this.redisLockProvider.lock(`usage-period-${sandboxId}`, 60)
  }

  private async releaseLock(sandboxId: string) {
    await this.redisLockProvider.unlock(`usage-period-${sandboxId}`)
  }
}
