import _ from "lodash";

import Request from "@/lib/request/Request.ts";
import { generateImages } from "@/api/controllers/images.ts";
import { tokenSplit } from "@/api/controllers/core.ts";
import util from "@/lib/util.ts";
import { ImagesTaskCache } from '@/api/ImagesTaskCache.ts';
import DatabaseGenerationService from '@/lib/services/DatabaseGenerationService.js';

// 通过环境变量控制新旧方法
const USE_DATABASE_MODE = process.env.USE_DATABASE_MODE === 'true';

export default {
  prefix: "/v1/images",
  get: {
    "/query": async (request: Request) => {
      request
        .validate("query.task_id", _.isString); // 从 query 中校验
      const {
        task_id,
      } = request.query; // 从 query 中获取
      
      if (USE_DATABASE_MODE) {
        // 使用新的数据库方法
        return await DatabaseGenerationService.queryTaskResult(task_id);
      } else {
        // 使用原有的内存缓存方法
        const imagesTaskCache = ImagesTaskCache.getInstance();
        
        // 使用新的方法获取任务结果并清理缓存
        let result = imagesTaskCache.getTaskResultAndClear(task_id);
        
        if (result && typeof result.status === 'number' && result.status < 0) { 
          // 任务已完成
          return {
            created: util.unixTimestamp(),
            data: { 
              task_id, 
              url: result.url || "", 
              status: result.status,
              error_message: result.errorMessage || null
            },
          };
        } else {
          // 任务进行中或失败，不清理缓存
          result = imagesTaskCache.getTaskStatus(task_id);
          return {
            created: util.unixTimestamp(),
            data: { 
              task_id, 
              url: "", 
              status: result?.status || 0,
              error_message: result?.errorMessage || null
            },
          };
        }
      }
    },
  },

  post: {
    "/generations": async (request: Request) => {
      request
        // .validate("body.model", v => _.isUndefined(v) || _.isString(v))
        .validate("body.task_id", _.isString)
        .validate("body.prompt", _.isString)
        // .validate("body.negative_prompt", v => _.isUndefined(v) || _.isString(v))
        .validate("body.width", v => _.isUndefined(v) || _.isFinite(v))
        .validate("body.height", v => _.isUndefined(v) || _.isFinite(v))
        // .validate("body.sample_strength", v => _.isUndefined(v) || _.isFinite(v))
        .validate("body.response_format", v => _.isUndefined(v) || _.isString(v))
        .validate("headers.authorization", _.isString);
      // refresh_token切分
      const tokens = tokenSplit(request.headers.authorization);
      // 随机挑选一个refresh_token
      const token = _.sample(tokens);
      const {
        // model,
        task_id,
        prompt,
        // negative_prompt: negativePrompt,
        width,
        height,
        // sample_strength: sampleStrength,
        response_format,
      } = request.body;
      const responseFormat = _.defaultTo(response_format, "url");
      
      if (USE_DATABASE_MODE) {
        // 使用新的数据库方法
        await DatabaseGenerationService.generateImagesV2('jimeng-3.0', task_id, prompt, {
          width,
          height,
          sampleStrength: 0.5,
          negativePrompt: "",
          response_format: responseFormat
        }, token);
      } else {
        // 使用原有方法（不等待结果）
        generateImages('jimeng-3.0', task_id, prompt, {
          width,
          height,
          sampleStrength: 0.5,
          negativePrompt: "",
        }, token);
      }
      
      return {
        created: util.unixTimestamp(),
        data: 'success',
      };
    },
  },
};
