src/services/implementations/taskQueue.service.ts

import { ITaskQueueService } from '../interfaces/ITaskQueueService';
import { ITaskRepository } from '../../repositories/interfaces/ITaskRepository';
import { IOffsetIndexService } from '../interfaces/IOffsetIndexService';
import { IWorkerService } from '../interfaces/IWorkerService';
import { ISchedulerService } from '../interfaces/ISchedulerService';
import { ITask, CreateTaskInput, TaskFilter, TaskStatus } from '../../types/task.types';
import { TaskModel } from '../../models/task.model';
import { CONFIG } from '../../config/app.config';
import { logger } from '../../utils/logger';

export class TaskQueueService implements ITaskQueueService {
  private taskRepo: ITaskRepository;
  private offsetIndexService: IOffsetIndexService;
  private workerService: IWorkerService;
  private schedulerService: ISchedulerService;

  constructor(
    taskRepo: ITaskRepository,
    offsetIndexService: IOffsetIndexService,
    workerService: IWorkerService,
    schedulerService: ISchedulerService
  ) {
    this.taskRepo = taskRepo;
    this.offsetIndexService = offsetIndexService;
    this.workerService = workerService;
    this.schedulerService = schedulerService;
  }

  public async createTask(input: CreateTaskInput): Promise<ITask> {
    const taskId = `task-${Date.now()}-${Math.random().toString(36).substring(2, 7)}`;
    const offsets = await this.offsetIndexService.getOffsetsForRange(input.startRow, input.endRow);

    const taskModel = new TaskModel({
      taskId,
      startRow: input.startRow,
      endRow: input.endRow,
      startOffset: offsets.startOffset,
      endOffset: offsets.endOffset,
      status: 'pending',
      retries: 0,
      progress: 0,
      workerId: null,
      leaseUntil: null,
      completedAt: null,
    });

    const saved = await this.taskRepo.save(taskModel);
    logger.info(`Created task [${saved.taskId}] rows ${saved.startRow} -> ${saved.endRow} (byte offsets: ${saved.startOffset} -> ${saved.endOffset})`);
    return saved;
  }

  public async generateBatchTasks(totalRows: number, defaultChunkSize: number = CONFIG.SCHEDULER.DEFAULT_BATCH_SIZE): Promise<number> {
    const existingCount = await this.taskRepo.count();
    if (existingCount > 0) {
      logger.info(`Task queue already contains ${existingCount} tasks. Skipping auto batch generation.`);
      return existingCount;
    }

    const tasksToCreate: ITask[] = [];
    let start = 0;
    const nowIso = new Date().toISOString();

    logger.info(`Generating task batch for dataset of ${totalRows} total rows (chunkSize=${defaultChunkSize})...`);

    while (start < totalRows) {
      const end = Math.min(start + defaultChunkSize, totalRows);
      const taskId = `task-${start}-${end}`;

      const task: ITask = {
        taskId,
        startRow: start,
        endRow: end,
        startOffset: 0, // Lazy fetched on assignment
        endOffset: 0,
        workerId: null,
        status: 'pending',
        retries: 0,
        progress: 0,
        rowsProcessed: 0,
        leaseUntil: null,
        createdAt: nowIso,
        updatedAt: nowIso,
        completedAt: null,
      };

      tasksToCreate.push(task);
      start = end;
    }

    const createdCount = await this.taskRepo.saveBatch(tasksToCreate);
    logger.info(`Successfully generated and queued ${createdCount} dataset tasks.`);
    return createdCount;
  }

  public async getTask(taskId: string): Promise<ITask | null> {
    return this.taskRepo.findById(taskId);
  }

  public async getTasks(filter?: TaskFilter): Promise<ITask[]> {
    return this.taskRepo.findAll(filter);
  }

  public async leaseNextTask(workerId: string, requestedChunkSize?: number): Promise<ITask | null> {
    const worker = await this.workerService.getWorker(workerId);
    if (!worker) {
      throw new Error(`Worker [${workerId}] not found.`);
    }

    const pendingTask = await this.taskRepo.findNextPending();
    if (!pendingTask) {
      return null;
    }

    // Determine target chunk size dynamically
    const idealChunkSize = requestedChunkSize || this.schedulerService.calculateTaskChunkSize(worker);
    let targetTask = pendingTask;

    // Split task if pending task is larger than dynamic allocation target
    const currentTaskSize = pendingTask.endRow - pendingTask.startRow;
    if (currentTaskSize > idealChunkSize * 1.5) {
      const originalEndRow = pendingTask.endRow;
      const allocatedEndRow = pendingTask.startRow + idealChunkSize;

      // Mutate current task to allocated range
      pendingTask.endRow = allocatedEndRow;
      
      // Create leftover slice task back in queue
      const leftoverTaskId = `task-${allocatedEndRow}-${originalEndRow}`;
      const leftoverTask: ITask = {
        taskId: leftoverTaskId,
        startRow: allocatedEndRow,
        endRow: originalEndRow,
        workerId: null,
        status: 'pending',
        retries: 0,
        progress: 0,
        rowsProcessed: 0,
        leaseUntil: null,
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        completedAt: null,
      };
      await this.taskRepo.save(leftoverTask);
    }

    // Resolve byte offsets from Index Service
    const offsets = await this.offsetIndexService.getOffsetsForRange(targetTask.startRow, targetTask.endRow);
    
    const leaseDuration = CONFIG.LEASE.DURATION_SECONDS * 1000;
    const leaseUntil = new Date(Date.now() + leaseDuration).toISOString();

    targetTask.workerId = workerId;
    targetTask.status = 'leased';
    targetTask.startOffset = offsets.startOffset;
    targetTask.endOffset = offsets.endOffset;
    targetTask.leaseUntil = leaseUntil;
    targetTask.updatedAt = new Date().toISOString();

    const savedTask = await this.taskRepo.save(targetTask);

    // Update worker status
    await this.workerService.setWorkerStatus(workerId, 'busy', savedTask.taskId);

    logger.info(`Leased task [${savedTask.taskId}] (rows ${savedTask.startRow} -> ${savedTask.endRow}) to worker [${workerId}] until ${leaseUntil}`);
    return savedTask;
  }

  public async releaseTask(taskId: string, reason?: string): Promise<ITask | null> {
    const task = await this.taskRepo.findById(taskId);
    if (!task) return null;

    logger.warn(`Releasing task [${taskId}] back to pending queue. Reason: ${reason || 'Unspecified'}`);
    task.status = 'pending';
    task.workerId = null;
    task.leaseUntil = null;
    task.progress = 0;
    task.retries += 1;
    task.errorMessage = reason || null;

    if (task.retries >= CONFIG.LEASE.MAX_RETRIES) {
      task.status = 'failed';
      logger.error(`Task [${taskId}] exceeded max retries (${CONFIG.LEASE.MAX_RETRIES}). Marked as permanently failed.`);
    }

    return this.taskRepo.save(task);
  }

  public async completeTask(taskId: string, rowsProcessed?: number): Promise<ITask | null> {
    const task = await this.taskRepo.findById(taskId);
    if (!task) return null;

    const nowIso = new Date().toISOString();
    task.status = 'completed';
    task.progress = 100;
    task.rowsProcessed = rowsProcessed || (task.endRow - task.startRow);
    task.completedAt = nowIso;
    task.leaseUntil = null;

    const saved = await this.taskRepo.save(task);

    if (task.workerId) {
      await this.workerService.updateWorkerStats(task.workerId, 1, 0, task.rowsProcessed);
      await this.workerService.setWorkerStatus(task.workerId, 'idle', null);
    }

    logger.info(`Task [${taskId}] completed successfully by worker [${task.workerId}]`);
    return saved;
  }

  public async failTask(taskId: string, error: string, canRetry: boolean = true): Promise<ITask | null> {
    const task = await this.taskRepo.findById(taskId);
    if (!task) return null;

    const workerId = task.workerId;
    task.errorMessage = error;

    if (canRetry && task.retries < CONFIG.LEASE.MAX_RETRIES) {
      return this.releaseTask(taskId, `Failed: ${error}`);
    } else {
      task.status = 'failed';
      task.leaseUntil = null;
      task.completedAt = new Date().toISOString();
      const saved = await this.taskRepo.save(task);

      if (workerId) {
        await this.workerService.updateWorkerStats(workerId, 0, 1, 0);
        await this.workerService.setWorkerStatus(workerId, 'idle', null);
      }

      logger.error(`Task [${taskId}] marked failed. Error: ${error}`);
      return saved;
    }
  }

  public async updateTaskProgress(taskId: string, progress: number, rowsProcessed?: number): Promise<ITask | null> {
    const task = await this.taskRepo.findById(taskId);
    if (!task || task.status !== 'leased') return null;

    task.progress = Math.min(Math.max(progress, 0), 100);
    if (rowsProcessed !== undefined) {
      task.rowsProcessed = rowsProcessed;
    }

    return this.taskRepo.save(task);
  }

  public async getTaskCounts(): Promise<Record<TaskStatus, number>> {
    const [pending, leased, completed, failed] = await Promise.all([
      this.taskRepo.count('pending'),
      this.taskRepo.count('leased'),
      this.taskRepo.count('completed'),
      this.taskRepo.count('failed'),
    ]);

    return { pending, leased, completed, failed };
  }

  public async resetAllTasks(): Promise<number> {
    return this.taskRepo.resetAll();
  }
}