import _ from "lodash";

import APIException from "@/lib/exceptions/APIException.ts";
import EX from "@/api/consts/exceptions.ts";
import util from "@/lib/util.ts";
import { getCredit, receiveCredit, request } from "./core.ts";
import logger from "@/lib/logger.ts";
import { ImagesTaskCache } from '@/api/ImagesTaskCache.ts';

const DEFAULT_ASSISTANT_ID = "513695";
export const DEFAULT_MODEL = "jimeng-3.0";
const DRAFT_VERSION = "3.0.2";
const MODEL_MAP = {
  "jimeng-4.5": "high_aes_general_v40l",
  "jimeng-4.1": "high_aes_general_v41",
  "jimeng-4.0": "high_aes_general_v40",
  "jimeng-3.1": "high_aes_general_v30l_art_fangzhou:general_v3.0_18b",
  "jimeng-3.0": "high_aes_general_v30l:general_v3.0_18b",
  "jimeng-2.1": "high_aes_general_v21_L:general_v2.1_L",
  "jimeng-2.0-pro": "high_aes_general_v20_L:general_v2.0_L",
  "jimeng-2.0": "high_aes_general_v20:general_v2.0",
  "jimeng-1.4": "high_aes_general_v14:general_v1.4",
  "jimeng-xl-pro": "text2img_xl_sft",
};

export function getModel(model: string) {
  return MODEL_MAP[model] || MODEL_MAP[DEFAULT_MODEL];
}

export async function generateImages(
  _model: string,
  task_id: string,
  prompt: string,
  {
    width = 1024,
    height = 1024,
    sampleStrength = 0.5,
    negativePrompt = "",
  }: {
    width?: number;
    height?: number;
    sampleStrength?: number;
    negativePrompt?: string;
  },
  refreshToken: string
) {
  const imagesTaskCache = ImagesTaskCache.getInstance();
  imagesTaskCache.startTask(task_id);

  try {
  const model = getModel(_model);
  const isV4 = _model.startsWith('jimeng-4');
  const currentVersion = isV4 ? '3.3.7' : DRAFT_VERSION;
  const resolutionType = isV4 ? '2k' : '1k';
  logger.info(`使用模型: ${_model} 映射模型: ${model} ${width}x${height} 精细度: ${sampleStrength}`);

  const { totalCredit } = await getCredit(refreshToken);
  if (totalCredit <= 0)
    await receiveCredit(refreshToken);

  const componentId = util.uuid();
  const draftContent = {
    type: "draft",
    id: util.uuid(),
    min_version: DRAFT_VERSION,
    min_features: [],
    is_from_tsn: true,
    version: currentVersion,
    main_component_id: componentId,
    component_list: [
      {
        type: "image_base_component",
        id: componentId,
        min_version: DRAFT_VERSION,
        gen_type: 1,
        generate_type: "generate",
        aigc_mode: "workbench",
        metadata: {
          type: "",
          id: util.uuid(),
          created_platform: 3,
          created_platform_version: "",
          created_time_in_ms: Date.now().toString(),
          created_did: ""
        },
        abilities: {
          type: "",
          id: util.uuid(),
          generate: {
            type: "",
            id: util.uuid(),
            core_param: {
              type: "",
              id: util.uuid(),
              model,
              prompt,
              negative_prompt: negativePrompt,
              seed: Math.floor(Math.random() * 100000000) + 2500000000,
              sample_strength: sampleStrength,
              image_ratio: 1,
              large_image_info: {
                type: "",
                id: util.uuid(),
                height,
                width,
                resolution_type: resolutionType,
              },
              intelligent_ratio: false
            },
            history_option: {
              type: "",
              id: util.uuid(),
            },
          },
          gen_option: {
            type: "",
            id: util.uuid(),
            gen_count: 1,
            generate_all: false
          }
        },
      },
    ],
  };

  const payloadData = {
    extend: {
      root_model: model,
      template_id: "",
    },
    submit_id: util.uuid(),
    metrics_extra: JSON.stringify({
      templateId: "0",
      generateCount: 1,
      promptSource: "custom",
      templateSource: "",
      lastRequestId: "",
      originRequestId: "",
      enterFrom: "use_bgimage_prompt",
      isRegenerate: false,
      isBoxSelect: false,
      isCutout: false
    }),
    draft_content: JSON.stringify(draftContent),
    http_common_info: {
      aid: Number(DEFAULT_ASSISTANT_ID),
    },
  };

  logger.info(`Jimeng Request Payload: ${JSON.stringify(payloadData)}`);

  const { aigc_data } = await request(
    "post",
    "/mweb/v1/aigc_draft/generate",
    refreshToken,
    {
      params: {
        babi_param: encodeURIComponent(
          JSON.stringify({
            scenario: "image_video_generation",
            feature_key: "aigc_to_image",
            feature_entrance: "to_image",
            feature_entrance_detail: "to_image-" + model,
          })
        ),
      },
      data: payloadData,
    }
  );
  logger.info(`Jimeng Response Data: ${JSON.stringify(aigc_data)}`);
  const historyId = aigc_data.history_record_id;
  if (!historyId)
    throw new APIException(EX.API_IMAGE_GENERATION_FAILED, "记录ID不存在");
  let status = 20, failCode = '', fail_msg = '', item_list = [];
  while (status === 20) {
    await new Promise((resolve) => setTimeout(resolve, 5000));
    const result = await request("post", "/mweb/v1/get_history_by_ids", refreshToken, {
      data: {
        history_ids: [historyId],
        image_info: {
          width: 2048,
          height: 2048,
          format: "webp",
          image_scene_list: [
            {
              scene: "smart_crop",
              width: 360,
              height: 360,
              uniq_key: "smart_crop-w:360-h:360",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 480,
              height: 480,
              uniq_key: "smart_crop-w:480-h:480",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 720,
              height: 720,
              uniq_key: "smart_crop-w:720-h:720",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 720,
              height: 480,
              uniq_key: "smart_crop-w:720-h:480",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 360,
              height: 240,
              uniq_key: "smart_crop-w:360-h:240",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 240,
              height: 320,
              uniq_key: "smart_crop-w:240-h:320",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 480,
              height: 640,
              uniq_key: "smart_crop-w:480-h:640",
              format: "webp",
            },
            {
              scene: "normal",
              width: 2400,
              height: 2400,
              uniq_key: "2400",
              format: "webp",
            },
            {
              scene: "normal",
              width: 1080,
              height: 1080,
              uniq_key: "1080",
              format: "webp",
            },
            {
              scene: "normal",
              width: 720,
              height: 720,
              uniq_key: "720",
              format: "webp",
            },
            {
              scene: "normal",
              width: 480,
              height: 480,
              uniq_key: "480",
              format: "webp",
            },
            {
              scene: "normal",
              width: 360,
              height: 360,
              uniq_key: "360",
              format: "webp",
            },
          ],
        },
        http_common_info: {
          aid: Number(DEFAULT_ASSISTANT_ID),
        },
      },
    });    
    logger.info(`1 图片生成查询 ${historyId} `);
    logger.log(`2 图片生成查询 ${historyId} `);
    if (!result[historyId])
      throw new APIException(EX.API_IMAGE_GENERATION_FAILED, "记录不存在");
    status = result[historyId].status;
    failCode = result[historyId].fail_code;
    fail_msg = result[historyId].fail_msg;
    item_list = result[historyId].item_list;
  }
  if (status === 30) {
    logger.info(`图片生成结果: failCode：${failCode} fail_msg: ${fail_msg}`);
    if (failCode == '2038' || failCode == '2041') {
      throw new APIException(EX.API_CONTENT_FILTERED);
    } else {
      throw new APIException(EX.API_IMAGE_GENERATION_FAILED);
    }
  }
  const imageUrls = item_list.map((item) => {
    let res_url = null;
    let idata = item?.image;
    let idata2 = item?.image?.large_images;
    if(!idata2 || !idata2.length){
        for(let key in idata){
            if(idata[key] && idata[key].length && idata[key].length > 0){
                if(idata[key][0].image_url){
                    res_url = idata[key][0].image_url;
                    break;
                }
            }
        }
    }else{
        res_url = idata2[0].image_url;
    }
    if(res_url) return res_url;
    return item?.common_attr?.cover_url || null;
  });
  const validImageUrls = imageUrls.filter(url => url !== null);
  if (validImageUrls.length > 0) {
    await imagesTaskCache.finishTask(task_id, -1, validImageUrls.join(",")); // Success
  } else {
    const errorMessage = "图片生成未返回有效链接";
    logger.info(`图片生成任务 ${task_id} 异常: ${errorMessage}`);
    // If no valid URLs but no explicit error thrown earlier, consider it a failure.
    // This could happen if item_list is empty or items don't have video_url.
    await imagesTaskCache.finishTask(task_id, -2, "", errorMessage); // Failure
    throw new APIException(EX.API_IMAGE_GENERATION_FAILED, errorMessage);
  }
  return validImageUrls;
}catch (error) {
  const errorMessage = error.message || "图片生成过程中发生未知错误";
  logger.info(`图片生成任务 ${task_id} 异常: ${errorMessage}`);
  await imagesTaskCache.finishTask(task_id, -2, "", errorMessage); // Failure due to exception
  throw error; // Re-throw the error to be handled by the caller
}
}

export async function uploadImages(
  _model: string,
  prompt: string,
  {
    width = 1024,
    height = 1024,
    sampleStrength = 0.5,
    negativePrompt = "",
  }: {
    width?: number;
    height?: number;
    sampleStrength?: number;
    negativePrompt?: string;
  },
  refreshToken: string
) {
  const model = getModel(_model);
  const isV4 = _model.startsWith('jimeng-4');
  const currentVersion = isV4 ? '3.3.7' : DRAFT_VERSION;
  logger.info(`使用模型: ${_model} 映射模型: ${model} ${width}x${height} 精细度: ${sampleStrength}`);

  const { totalCredit } = await getCredit(refreshToken);
  if (totalCredit <= 0)
    await receiveCredit(refreshToken);

  const componentId = util.uuid();
  const { aigc_data } = await request(
    "post",
    "/mweb/v1/get_upload_token",
    refreshToken,
    {
      params: {
        babi_param: encodeURIComponent(
          JSON.stringify({
            scenario: "image_video_generation",
            feature_key: "aigc_to_image",
            feature_entrance: "to_image",
            feature_entrance_detail: "to_image-" + model,
          })
        ),
      },
      data: {
        extend: {
          root_model: model,
          template_id: "",
        },
        submit_id: util.uuid(),
        metrics_extra: JSON.stringify({
          templateId: "",
          generateCount: 1,
          promptSource: "custom",
          templateSource: "",
          lastRequestId: "",
          originRequestId: "",
        }),
        draft_content: JSON.stringify({
          type: "draft",
          id: util.uuid(),
          min_version: DRAFT_VERSION,
          is_from_tsn: true,
          version: currentVersion,
          main_component_id: componentId,
          component_list: [
            {
              type: "image_base_component",
              id: componentId,
              min_version: DRAFT_VERSION,
              generate_type: "generate",
              aigc_mode: "workbench",
              abilities: {
                type: "",
                id: util.uuid(),
                generate: {
                  type: "",
                  id: util.uuid(),
                  core_param: {
                    type: "",
                    id: util.uuid(),
                    model,
                    prompt,
                    negative_prompt: negativePrompt,
                    seed: Math.floor(Math.random() * 100000000) + 2500000000,
                    sample_strength: sampleStrength,
                    image_ratio: 1,
                    large_image_info: {
                      type: "",
                      id: util.uuid(),
                      height,
                      width,
                    },
                  },
                  history_option: {
                    type: "",
                    id: util.uuid(),
                  },
                },
              },
            },
          ],
        }),
        http_common_info: {
          aid: Number(DEFAULT_ASSISTANT_ID),
        },
      },
    }
  );
  const historyId = aigc_data.history_record_id;
  if (!historyId)
    throw new APIException(EX.API_IMAGE_GENERATION_FAILED, "记录ID不存在");
  let status = 20, failCode, item_list = [];
  while (status === 20) {
    await new Promise((resolve) => setTimeout(resolve, 1000));
    const result = await request("post", "/mweb/v1/get_history_by_ids", refreshToken, {
      data: {
        history_ids: [historyId],
        image_info: {
          width: 2048,
          height: 2048,
          format: "webp",
          image_scene_list: [
            {
              scene: "smart_crop",
              width: 360,
              height: 360,
              uniq_key: "smart_crop-w:360-h:360",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 480,
              height: 480,
              uniq_key: "smart_crop-w:480-h:480",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 720,
              height: 720,
              uniq_key: "smart_crop-w:720-h:720",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 720,
              height: 480,
              uniq_key: "smart_crop-w:720-h:480",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 360,
              height: 240,
              uniq_key: "smart_crop-w:360-h:240",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 240,
              height: 320,
              uniq_key: "smart_crop-w:240-h:320",
              format: "webp",
            },
            {
              scene: "smart_crop",
              width: 480,
              height: 640,
              uniq_key: "smart_crop-w:480-h:640",
              format: "webp",
            },
            {
              scene: "normal",
              width: 2400,
              height: 2400,
              uniq_key: "2400",
              format: "webp",
            },
            {
              scene: "normal",
              width: 1080,
              height: 1080,
              uniq_key: "1080",
              format: "webp",
            },
            {
              scene: "normal",
              width: 720,
              height: 720,
              uniq_key: "720",
              format: "webp",
            },
            {
              scene: "normal",
              width: 480,
              height: 480,
              uniq_key: "480",
              format: "webp",
            },
            {
              scene: "normal",
              width: 360,
              height: 360,
              uniq_key: "360",
              format: "webp",
            },
          ],
        },
        http_common_info: {
          aid: Number(DEFAULT_ASSISTANT_ID),
        },
      },
    });
    if (!result[historyId])
      throw new APIException(EX.API_IMAGE_GENERATION_FAILED, "记录不存在");
    status = result[historyId].status;
    failCode = result[historyId].fail_code;
    item_list = result[historyId].item_list;
  }
  if (status === 30) {
    if (failCode === '2038')
      throw new APIException(EX.API_CONTENT_FILTERED);
    else
      throw new APIException(EX.API_IMAGE_GENERATION_FAILED);
  }
  return item_list.map((item) => {
    let res_url = null;
    let idata = item?.image;
    let idata2 = item?.image?.large_images;
    if(!idata2 || !idata2.length){
        for(let key in idata){
            if(idata[key] && idata[key].length && idata[key].length > 0){
                if(idata[key][0].image_url){
                    res_url = idata[key][0].image_url;
                    break;
                }
            }
        }
    }else{
        res_url = idata2[0].image_url;
    }
    if(res_url) return res_url;
    return item?.common_attr?.cover_url || null;
  });
}

export default {
  generateImages,
  uploadImages,
};