import { Kevisual } from "@kevisual/ai";
import { useConfig } from '@kevisual/use-config'
import path from "node:path";
import fs from "node:fs";
import { createStorage } from "unstorage";
import { Prompt } from "../src/index.ts";
import fsDriver from "unstorage/drivers/fs";

const promptPath = path.join(process.cwd(), "./docs/prompts-01.json");

const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

const storage = createStorage({
  driver: fsDriver({ base: 'storage' }),
});
const config = useConfig();
const kevisual = new Kevisual({
  token: config.KEVISUAL_NEW_API_KEY || "",
  model: "qwen-turbo"
});

if (!fs.existsSync(promptPath)) {
  console.error("Prompt file not found:", promptPath);
  process.exit(1);
}
const prompts = JSON.parse(fs.readFileSync(promptPath, "utf-8")) as { key: string, value: string, id: string }[];

const currentUsage = (await storage.getItem("usage.json") as any) || { prompt_tokens: 0, total_tokens: 0, completion_tokens: 0 };

console.log("Current usage:", currentUsage);
async function generatePerfectPrompts() {
  for (const promptData of prompts) {
    const id = promptData.id + '.json';
    const has = await storage.get(id);
    if (has) {
      continue;
    }
    try {
      const perfectPrompt = `请你将以下提示词进行完善，使其更加详细和具体，适合用于生成高质量的像素艺术图像。要求如下：
1. 只返回完善后的提示词，不要包含任何多余的内容或解释。
2. 确保提示词专注于像素艺术风格，包括但不限于像素化角色、场景和物体的描述。
3. 使用具体的细节来增强提示词的表现力，例如颜色、构图、光影效果等。
4. 避免使用与像素艺术无关的术语或描述。
5. 保持提示词的简洁性，避免过于冗长，但要确保信息量充足。
6. 如果需要颜色，需要整个图像的颜色更少的描述，而不是复杂的颜色细节, 背景默认纯蓝色。
7. 使用中文进行描述。
`;
      const prompt = new Prompt({ perfectPrompt });
      const result = await kevisual.chat([
        {
          role: "user",
          content: prompt.perfect(promptData.value),
        }
      ])
      const text = prompt.clearPerfectTags(kevisual.responseText);
      await storage.setItem(id, {
        value: promptData.value,
        id: promptData.id,
        perfect: text,
      });

      await sleep(2000); // Avoid rate limits
      const _usage = kevisual.getChatUsage()
      console.log("Generated perfect prompt for:", promptData.id, _usage?.total_tokens);
      currentUsage.total_tokens += _usage?.total_tokens || 0;
      currentUsage.prompt_tokens += _usage?.prompt_tokens || 0;
      currentUsage.completion_tokens += _usage?.completion_tokens || 0;
      // Update usage
      await storage.setItem("usage.json", {
        prompt_tokens: currentUsage.prompt_tokens,
        total_tokens: currentUsage.total_tokens,
        completion_tokens: currentUsage.completion_tokens,
      });
      console.log('优化的提示词', text);
    } catch (error) {
      console.error("Error generating perfect prompt for:", promptData, error);
      // 如果是超时错误，等待一段时间后继续
      if ((error as any)?.message?.includes("timeout")) {
        console.log("Timeout occurred, waiting for 30 seconds before retrying...");
        await sleep(10000);
        continue;
      } else {
        process.exit(1);
      }
    }
  }

}

generatePerfectPrompts();

//  bun test/generate-perfect.ts
// pm2 start --name "generate-perfect" --interpreter=bun -- test/generate-perfect.ts