import { Worker, Job } from 'bullmq';
import { getRedisConnection } from '../module/redis.ts';
import { Prompt, pbService, ai } from '../index.ts';
import type { ImageCollection } from '../services/pb.service.ts';
import { updateItemStatus } from '../routes/image-update.ts';
import { Queue } from 'bullmq';
import { notify } from '@/module/logger.ts';
import { addImageGenerateJob } from './image-creator.job.ts';

export const PERFECT_PROMPT_JOB = 'perfect-prompt';

// 状态常量
export const PerfectPromptStatus = {
  PENDING: '提示词优化中' as const,
  PLANNING: '计划中' as const,
  FAILED: '失败' as const,
};

// 最大重试次数
const MAX_RETRIES = 3;

export interface PerfectPromptJobData {
  itemId: string;
  prompt: string;
  collectionName?: string;
  data?: Record<string, any>;
}

// 优化提示词的模板
const DEFAULT_PERFECT_PROMPT = `请你将以下提示词进行完善，使其更加详细和具体，适合用于生成高质量的像素艺术图像。要求如下：
1. 只返回完善后的提示词，不要包含任何多余的内容或解释。
2. 确保提示词专注于像素艺术风格，包括但不限于像素化角色、场景和物体的描述。
3. 使用具体的细节来增强提示词的表现力，例如颜色、构图、光影效果等。
4. 避免使用与像素艺术无关的术语或描述。
5. 保持提示词的简洁性，避免过于冗长，但要确保信息量充足。
6. 如果需要颜色，需要整个图像的颜色更少的描述，而不是复杂的颜色细节, 背景默认纯蓝色。
7. 使用中文进行描述。
`;

/**
 * 单独添加优化提示词任务
 */
export async function addPerfectPromptJob(item: ImageCollection): Promise<void> {
  const connection = getRedisConnection();
  const queue = new Queue(PERFECT_PROMPT_JOB, { connection });

  const jobData: PerfectPromptJobData = {
    itemId: item.id,
    prompt: item.description || item.summary || item.title || '',
    collectionName: pbService.collectionName,
  };

  await queue.add(PERFECT_PROMPT_JOB, jobData, {
    attempts: MAX_RETRIES,
    backoff: {
      type: 'exponential',
      delay: 2000,
    },
    removeOnComplete: 100,
    removeOnFail: 100,
  });

  await updateItemStatus(item.id, PerfectPromptStatus.PENDING);
  await queue.close();
}

/**
 * 运行优化提示词 worker
 */
export async function runPerfectPromptWorker(): Promise<void> {
  const connection = getRedisConnection();
  // 获取环境变量中的 API key
  const worker = new Worker(
    PERFECT_PROMPT_JOB,
    async (job: Job<PerfectPromptJobData>) => {
      const { itemId, prompt } = job.data;
      const attemptsMade = job.attemptsMade;
      console.log(`[PerfectPrompt] Processing item: ${itemId}, attempt: ${attemptsMade + 1}/${MAX_RETRIES}`);
      try {
        if (!prompt) {
          throw new Error('Prompt is empty');
        }
        const promptTool = new Prompt({ perfectPrompt: DEFAULT_PERFECT_PROMPT });
        await ai.chat([
          {
            role: 'user',
            content: promptTool.perfect(prompt),
          },
        ]);
        const perfectText = promptTool.clearPerfectTags(ai.responseText);

        if (!perfectText) {
          throw new Error('Generated perfect prompt is empty');
        }

        console.log(`[PerfectPrompt] Perfect prompt generated for item: ${itemId}`);

        // 更新状态为已完成，并保存优化后的提示词
        await updateItemStatus(itemId, PerfectPromptStatus.PLANNING, {
          description: perfectText,
        });
        // 任务完成，把任务抛给下一个图片生成队列
        const item = await pbService.collection.getOne(itemId);
        if (item) {
          addImageGenerateJob(item)
        }

        return { success: true, perfectPrompt: perfectText };
      } catch (error: any) {
        console.error(`[PerfectPrompt] Error: ${error.message}`);

        // 重试次数用尽，标记为失败
        if (job.attemptsMade >= MAX_RETRIES - 1) {
          await updateItemStatus(itemId, PerfectPromptStatus.FAILED);
        }

        throw error;
      }
    },
    {
      connection,
      concurrency: 2,
      lockDuration: 60000 * 5, // 锁持续时间 5分钟
      stalledInterval: 30000, // 每30秒检查一次 stalled
    }
  );

  worker.on('completed', (job) => {
    console.log(`[PerfectPrompt] Job completed: ${job.id}`);
  });

  worker.on('failed', (job, err) => {
    console.error(`[PerfectPrompt] Job failed: ${job?.id}, error: ${err.message}`);
    notify.notify(`[PerfectPrompt] \nJob failed: ${job?.id}, error: ${err.message}\n Job data: ${JSON.stringify(job?.data)}`);
    if (job && job.attemptsMade >= MAX_RETRIES - 1) {
      worker.close();
      notify.notify(`[PerfectPrompt] Worker stopped after reaching max retries for item ${job.data.itemId}`);
    }
  });

  console.log('[PerfectPrompt] Worker started');
}