const models = {
  "o1-preview": "o1-preview",
  "o1-mini": "o1-preview",
  "claude-3-5-sonnet-20240620": "claude-3-5-sonnet",
  "claude-3-5-sonnet-20241022": "claude-3-5-sonnet",
  "claude-3-5-haiku-20241022": "claude-3-5-haiku",
  "claude-3-haiku-20240307": "claude-3-5-haiku",
  "gpt-4o": "gpt-4o",
  "gpt-4o-mini": "gpt-4o-mini",
  "gpt-4o-file": "gpt-4o",
  "gpt-4o-mini-file": "gpt-4o-mini",
  "claude-3-5-sonnet-file": "claude-3-5-sonnet",
  "claude-3-5-haiku-file": "claude-3-5-haiku",
  "gemini-1.5-pro-latest": "gemini-1.5-pro",
  "gemini-1.5-flash-latest": "gemini-1.5-flash",
  "genspark": "genspark",
  "deepl": "deepl", // 添加翻译模型
};

// 新增画图模型
const drawingModels = new Set([
  "dall-e-3",
  "flux",
  "flux-speed",
  "flux-pro/ultra",
  "ideogram",
  "recraft-v3",
]);

// 添加缓存相关的常量配置
const CACHE_CONFIG = {
  TTL: 30 * 60, // 缓存30分钟aa
  EXCLUDED_MODELS: ["o1-preview"], // 不缓存的模型
  MAX_CACHE_SIZE: 100, // 最大缓存条目数
};

// 简单的内存缓存实现
class MessageCache {
  constructor() {
    this.cache = new Map();
    this.keyTimestamps = new Map();
  }

  generateKey(messages, model) {
    // 生成缓存key
    const contentStr = JSON.stringify(
      messages.map((m) => ({
        role: m.role,
        content: m.content,
      }))
    );
    return `${model}:${contentStr}`;
  }

  set(key, value) {
    // 检查缓存大小
    if (this.cache.size >= CACHE_CONFIG.MAX_CACHE_SIZE) {
      // 删除最旧的缓存
      const oldestKey = [...this.keyTimestamps.entries()].sort(
        ([, a], [, b]) => a - b
      )[0][0];
      this.cache.delete(oldestKey);
      this.keyTimestamps.delete(oldestKey);
    }

    this.cache.set(key, value);
    this.keyTimestamps.set(key, Date.now());
  }

  get(key) {
    const timestamp = this.keyTimestamps.get(key);
    if (!timestamp) return null;

    // 检查是否过期
    if (Date.now() - timestamp > CACHE_CONFIG.TTL * 1000) {
      this.cache.delete(key);
      this.keyTimestamps.delete(key);
      return null;
    }

    return this.cache.get(key);
  }
}

const messageCache = new MessageCache();

class SessionPool {
  constructor() {
    this.sessions = [];
    this.currentIndex = 0;
  }
  initialize(sessionString) {
    if (!sessionString) return;
    this.sessions = sessionString
      .split(",")
      .map((s) => s.trim())
      .filter((s) => s);
    this.currentIndex = Math.floor(Math.random() * this.sessions.length);
  }
  getNext() {
    if (this.sessions.length === 0) return null;
    const session = this.sessions[this.currentIndex];
    this.currentIndex = (this.currentIndex + 1) % this.sessions.length;
    return session;
  }
  size() {
    return this.sessions.length;
  }
}

const sessionPool = new SessionPool();

const corsHeaders = {
  "Access-Control-Allow-Origin": "*",
  "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
  "Access-Control-Allow-Headers": "Content-Type, Authorization",
};

function doubleEncode(str) {
  return encodeURIComponent(encodeURIComponent(str));
}

function createErrorResponse(message, status = 500) {
  return new Response(JSON.stringify({ error: message }), {
    status,
    headers: { ...corsHeaders, "Content-Type": "application/json" },
  });
}

async function cleanupProject(session_id, project_id) {
  try {
    const headers = {
      Cookie: `session_id=${session_id}`,
      "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
      Accept: "*/*",
      Host: "www.genspark.ai",
      Connection: "keep-alive",
      "Accept-Encoding": "gzip, deflate, br", // 添加压缩支持
    };
    const response = await fetch(
      `https://www.genspark.ai/api/project/delete?project_id=${project_id}`,
      {
        method: "GET",
        headers,
      }
    );
    return response.ok;
  } catch (error) {
    console.error("Cleanup project error:", error);
    return false;
  }
}

function extractProjectId(text) {
  try {
    if (text.startsWith("data: ")) {
      const content = JSON.parse(text.slice(6));
      if (content.type === "project_start" && content.id) {
        return content.id;
      }
    }
  } catch (e) {}
  return null;
}

function processChunk(chunk, messageId, model) {
  try {
    const content = chunk.replace(/^data:\s*/, "").trim();
    if (!content || content === "[DONE]") return null;

    const parsed = JSON.parse(content);
    if (!parsed) return null;

    // 处理翻译模型的特殊返回格式
    if (model === "deepl") {
      if (parsed.field_name === "content" && parsed.field_value) {
        return JSON.stringify({
          translations: [
            {
              detected_source_language:
                parsed.session_state?.detected_source_lang || "",
              text: parsed.field_value,
            },
          ],
        });
      }
    } else if (drawingModels.has(model)) {
      // 处理画图模型的响应
      if (parsed.content) {
        return JSON.stringify({
          created: Date.now(),
          data: [
            {
              url: parsed.content,
              revised_prompt: parsed.prompt || "",
            },
          ],
        });
      }
    } else {
      if (parsed.delta) {
        return `data: ${JSON.stringify({
          id: `chatcmpl-${messageId}`,
          choices: [
            {
              index: 0,
              delta: { content: parsed.delta },
            },
          ],
          created: Math.floor(Date.now() / 1000),
          model: models[model],
          object: "chat.completion.chunk",
        })}\n\n`;
      }
      if (
        parsed.field_value &&
        parsed.field_name !== "session_state.answer_is_finished" &&
        parsed.field_name !== "content" &&
        parsed.field_name !== "session_state" &&
        !parsed.delta &&
        parsed.type !== "project_field"
      ) {
        return JSON.stringify({
          id: `chatcmpl-${messageId}`,
          object: "chat.completion",
          created: Math.floor(Date.now() / 1000),
          model,
          choices: [
            {
              index: 0,
              message: {
                role: "assistant",
                content: parsed.field_value,
              },
              finish_reason: "stop",
            },
          ],
          usage: {
            prompt_tokens: 100,
            completion_tokens: 50,
            total_tokens: parsed.field_value.length,
          },
        });
      }
    }
  } catch (e) {}
  return null;
}

async function* streamGenerator(reader) {
  let buffer = "";
  const decoder = new TextDecoder();

  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split("\n");
      buffer = lines.pop() || "";

      for (const line of lines) {
        if (line.trim()) yield line;
      }
    }
    if (buffer.trim()) yield buffer;
  } catch (e) {
    console.error("Stream error:", e);
  }
}

async function makeRequest(
  session_id,
  requestModel,
  messages,
  target_lang_code,
  source_lang_code,
  num_images = 1,
  size = "1024x1024",
  quality = "standard",
  style = "vivid"
) {
  const headers = {
    Cookie: `session_id=${session_id}`,
    "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
    "Content-Type": "application/json",
    Accept: "*/*",
    Host: "www.genspark.ai",
    Connection: "keep-alive",
  };
  let retryCount = 0;
  while (retryCount < 3) {
    try {
      let requestBody = {};
      if (requestModel === "deepl") {
        // 构建翻译模型的请求体
        requestBody = {
          type: "COPILOT_MOA_TRANSLATOR",
          current_query_string: "type=COPILOT_MOA_TRANSLATOR",
          messages,
          user_s_input: messages[messages.length - 1].content,
          action_params: {},
          extra_data: {
            models: [models[requestModel]],
            run_with_another_model: false,
            target_lang_code: target_lang_code || "zh-hans",
            source_lang_code: source_lang_code || "",
            writingContent: null,
          },
        };
      } else if (drawingModels.has(requestModel)) {
        // 构建画图模型的请求体
        requestBody = {
          type: "COPILOT_MOA_IMAGE",
          current_query_string: "type=COPILOT_MOA_IMAGE",
          messages,
          user_s_input: messages[messages.length - 1].content,
          action_params: {},
          extra_data: {
            model_configs: [],
            llm_model: "gpt-4",
            imageModelMap: {},
            writingContent: null,
          },
        };
        for (let i = 0; i < num_images; i++) {
          requestBody.extra_data.model_configs.push({
            model: requestModel === "dall-e-3" ? "dalle-3" : requestModel,
            size: size,
            quality: quality,
            style: style,
          });
        }
      } else {
        // 构建聊天模型的请求体
        requestBody = {
          type: "COPILOT_MOA_CHAT",
          current_query_string: "type=COPILOT_MOA_CHAT",
          messages,
          action_params: {},
          extra_data: {
            models: [models[requestModel] || models["claude-3-5-sonnet-20241022"]],
            run_with_another_model: false,
            writingContent: null,
          },
        };
      }

      const response = await fetch(
        "https://www.genspark.ai/api/copilot/ask",
        {
          method: "POST",
          headers,
          body: JSON.stringify(requestBody),
        }
      );
      if (!response.ok) {
        retryCount++;
        await new Promise((resolve) =>
          setTimeout(resolve, 1000 * retryCount)
        );
        continue;
      }
      return response;
    } catch (error) {
      retryCount++;
      await new Promise((resolve) =>
        setTimeout(resolve, 1000 * retryCount)
      );
      if (retryCount >= 3) throw error;
    }
  }
  throw new Error("请求失败，已重试3次");
}

async function searchModel(session_id, messages) {
  const headers = {
    Cookie: `session_id=${session_id}`,
    "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
    "Content-Type": "application/json",
    Accept: "*/*",
    Host: "www.genspark.ai",
    Connection: "keep-alive",
  };
  let retryCount = 0;
  while (retryCount < 3) {
    try {
      const query = doubleEncode(messages[messages.length - 1].content);
      const response = await fetch(
        `https://www.genspark.ai/api/search/stream?query=${query}`,
        {
          method: "POST",
          headers,
        }
      );
      if (!response.ok) {
        retryCount++;
        await new Promise((resolve) =>
          setTimeout(resolve, 1000 * retryCount)
        );
        continue;
      }
      return response;
    } catch (error) {
      retryCount++;
      await new Promise((resolve) =>
        setTimeout(resolve, 1000 * retryCount)
      );
      if (retryCount >= 3) throw error;
    }
  }
  throw new Error("请求失败，已重试3次");
}

async function pollTaskId(taskId, session_id, maxRetries, retryInterval) {
  let retries = 0;
  while (retries < maxRetries) {
    try {
      const statusResponse = await fetch(
        `https://www.genspark.ai/api/spark/image_generation_task_status?task_id=${taskId}`,
        {
          headers: {
            Cookie: `session_id=${session_id}`,
            "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
            Accept: "*/*",
            Host: "www.genspark.ai",
            Connection: "keep-alive",
          },
        }
      );
      if (!statusResponse.ok) {
        throw new Error(`Status check failed: ${statusResponse.status}`);
      }
      const data = await statusResponse.json();
      if (data.data.status === "SUCCESS") {
        const urls = data.data.image_urls_nowatermark;
        if (urls?.length) {
          return urls;
        }
      } else if (data.data.status === "FAILED") {
        throw new Error("Image generation failed");
      }
    } catch (error) {
      console.error(`Poll error for task ${taskId} (retry ${retries}):`, error);
    }
    await new Promise((r) => setTimeout(r, retryInterval));
    retries++;
  }
  throw new Error(`Timeout waiting for task: ${taskId}`);
}

function arrayBufferToBase64(buffer) {
  let binary = '';
  const bytes = new Uint8Array(buffer);
  const len = bytes.byteLength;
  const chunkSize = 8192; // 分块大小，可以根据需要调整

  for (let i = 0; i < len; i += chunkSize) {
    const chunk = bytes.subarray(i, i + chunkSize);
    binary += String.fromCharCode.apply(null, chunk);
  }
  return btoa(binary);
}

// 添加处理消息的函数
async function processMessages(messages, env) {
  // 对每个消息进行处理
  for (let message of messages) {
    // 检查 message.content 是否为数组
    if (Array.isArray(message.content)) {
      // 处理 content 数组中的每个条目
      for (let i = 0; i < message.content.length; i++) {
        let entry = message.content[i];
        if (entry.type === 'image_url') {
          // 处理 image_url
          let url = entry.image_url?.url;
          if (!url) {
            throw new Error('image_url.url is required');
          }

          let base64ImageData;
          if (url.startsWith('data:image')) {
            // 已经是 base64 编码
            base64ImageData = url;
          } else {
            // 从 URL 获取图片并编码为 base64
            base64ImageData = await fetchImageAsDataURL(url, env);
          }
          // 用 base64 数据 URL 替换
          entry.image_url.url = base64ImageData;
        }
      }
    }
    // 否则，message.content 是字符串，无需处理
  }
}

// 添加获取图片并编码为 base64 的函数
async function fetchImageAsDataURL(url, env) {
  // 检查是否在 KV 缓存中
  let cacheKey = `image:${url}`;
  let cachedData = await env.kv.get(cacheKey);
  if (cachedData) {
    return cachedData;
  }
  // 获取图片
  let response = await fetch(url);
  if (!response.ok) {
    throw new Error(`Failed to fetch image from ${url}`);
  }
  let arrayBuffer = await response.arrayBuffer();
  let base64Data = arrayBufferToBase64(arrayBuffer);
  // 获取内容类型
  let contentType = response.headers.get('Content-Type') || 'image/jpeg';
  // 构建数据 URL
  let dataURL = `data:${contentType};base64,${base64Data}`;
  // 存储到 KV 缓存
  await env.kv.put(cacheKey, dataURL, { expirationTtl: CACHE_CONFIG.TTL });
  return dataURL;
}

// 提取最后的用户输入文本
function extractLastUserInput(messages) {
  // 从后向前遍历 messages
  for (let i = messages.length - 1; i >= 0; i--) {
    let message = messages[i];
    if (message.role === 'user') {
      if (Array.isArray(message.content)) {
        // 处理 content 数组
        for (let j = message.content.length - 1; j >= 0; j--) {
          let entry = message.content[j];
          if (entry.type === 'text' && entry.text) {
            return entry.text;
          }
        }
      } else if (typeof message.content === 'string') {
        // content 是字符串
        return message.content;
      }
    }
  }
  return '';
}

export default {
  async fetch(request, env, ctx) {
    if (request.method === "OPTIONS") {
      return new Response(null, { headers: corsHeaders });
    }

    const url = new URL(request.url);
    if (url.pathname === "/health") {
      return new Response(JSON.stringify({ status: "ok" }), {
        headers: { ...corsHeaders, "Content-Type": "application/json" },
      });
    }

    if (url.pathname === "/v1/models" && request.method === "GET") {
      return new Response(
        JSON.stringify({
          object: "list",
          data: Object.keys(models).map((model) => ({
            id: model,
            object: "model",
            created: 1706745938,
            owned_by: "genspark",
          })),
        }),
        { headers: { ...corsHeaders, "Content-Type": "application/json" } }
      );
    }

    // 添加对 /translate 接口的支持
    if (
      (url.pathname === "/v1/chat/completions" ||
        url.pathname === "/translate") &&
      request.method === "POST"
    ) {
      try {
        let requestData = await request.json();
        let {
          messages,
          prompt,
          stream = false,
          model = "deepl",
          target_lang_code,
          source_lang_code,
        } = requestData;

        // 如果是 /translate 接口，调整请求数据
        if (url.pathname === "/translate") {
          const { text, source_lang, target_lang } = requestData;
          messages = [{ role: "user", content: text }];
          model = "deepl";
          target_lang_code = target_lang || "zh-hans";
          source_lang_code = source_lang || "";
        }

        // 如果没有 messages，且有 prompt，创建 messages
        if (!messages && prompt) {
          messages = [{ role: "user", content: prompt }];
        }

        // 提取 num_images, size, response_format
        const {
          num_images = 1,
          size = "1024x1024",
          response_format = "url",
          quality = "standard",
          style = "vivid",
        } = requestData;

        const authHeader = request.headers.get("authorization");

        if (sessionPool.size() === 0) {
          const sessionString = authHeader?.replace("Bearer ", "");
          sessionPool.initialize(sessionString);
          if (sessionPool.size() === 0) {
            return createErrorResponse("未提供有效的 session_id", 401);
          }
        }
        const session_id = sessionPool.getNext();
        if (!session_id) {
          return createErrorResponse("无可用的会话ID", 500);
        }

        // 添加缓存处理
        const cacheKey = messageCache.generateKey(messages, model);
        let cachedResponse = null;
        if (!CACHE_CONFIG.EXCLUDED_MODELS.includes(model)) {
          cachedResponse = messageCache.get(cacheKey);
        }
        if (cachedResponse) {
          return new Response(cachedResponse, {
            headers: { ...corsHeaders, "Content-Type": "application/json" },
          });
        }

        // 处理 messages，处理图片
        await processMessages(messages, env);
        const user_s_input = extractLastUserInput(messages);

        if (drawingModels.has(model)) {
          // 处理画图请求
          // 构建请求体
          const requestBody = {
            type: "COPILOT_MOA_IMAGE",
            current_query_string: "type=COPILOT_MOA_IMAGE",
            messages,
            user_s_input: user_s_input,
            action_params: {},
            extra_data: {
              model_configs: [],
              llm_model: "gpt-4",
              imageModelMap: {},
              writingContent: null,
            },
          };
          for (let i = 0; i < num_images; i++) {
            requestBody.extra_data.model_configs.push({
              model: model === "dall-e-3" ? "dalle-3" : model,
              size: size,
              quality: quality,
              style: style,
            });
          }
          const headers = {
            Cookie: `session_id=${session_id}`,
            "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
            "Content-Type": "application/json",
            Accept: "*/*",
            Host: "www.genspark.ai",
            Connection: "keep-alive",
          };
          if (stream === true || stream === "true") {
            // Handle streaming response
            const { readable, writable } = new TransformStream();
            const writer = writable.getWriter();

            ctx.waitUntil(
              (async () => {
                try {
                  const response = await fetch(
                    "https://www.genspark.ai/api/copilot/ask",
                    {
                      method: "POST",
                      headers,
                      body: JSON.stringify(requestBody),
                    }
                  );

                  if (!response.ok) {
                    await writer.write(
                      new TextEncoder().encode(
                        `data: ${JSON.stringify({
                          error: "上游服务请求失败",
                        })}\n\n`
                      )
                    );
                    await writer.close();
                    return;
                  }

                  const reader = response.body.getReader();
                  let taskIds = [];

                  // Read the response from upstream and extract task IDs
                  let responseText = "";
                  for await (const line of streamGenerator(reader)) {
                    responseText += line + "\n";
                    if (line.includes("task_id")) {
                      try {
                        const contentLine = line.replace("data: ", "");
                        const { content: innerContent } = JSON.parse(contentLine);
                        if (innerContent) {
                          const { generated_images } = JSON.parse(innerContent);
                          if (generated_images) {
                            for (const img of generated_images) {
                              if (img.task_id) {
                                taskIds.push(img.task_id);
                              }
                            }
                          }
                        }
                      } catch (e) {}
                    }
                  }

                  if (!taskIds.length) {
                    await writer.write(
                      new TextEncoder().encode(
                        `data: ${JSON.stringify({
                          error: "获取任务ID失败",
                        })}\n\n`
                      )
                    );
                    await writer.close();
                    return;
                  }

                  // Now we can poll the task IDs and send images to the client
                  const maxRetries = 30;
                  const retryInterval = 2000;

                  for (const taskId of taskIds) {
                    let retries = 0;
                    let imageSent = false;
                    while (retries < maxRetries && !imageSent) {
                      try {
                        const statusResponse = await fetch(
                          `https://www.genspark.ai/api/spark/image_generation_task_status?task_id=${taskId}`,
                          {
                            headers,
                          }
                        );

                        if (!statusResponse.ok) {
                          throw new Error(
                            `Status check failed: ${statusResponse.status}`
                          );
                        }

                        const data = await statusResponse.json();
                        if (data.data.status === "SUCCESS") {
                          const urls = data.data.image_urls_nowatermark;
                          if (urls?.length) {
                            // Send the image URLs to the client
                            for (const url of urls) {
                              const data = {
                                id: `chatcmpl-${crypto.randomUUID()}`,
                                object: "chat.completion.chunk",
                                created: Math.floor(Date.now() / 1000),
                                model: model,
                                choices: [
                                  {
                                    delta: { content: `![Image](${url})\n` },
                                    index: 0,
                                    finish_reason: null,
                                  },
                                ],
                              };
                              await writer.write(
                                new TextEncoder().encode(
                                  `data: ${JSON.stringify(data)}\n\n`
                                )
                              );
                            }
                            imageSent = true;
                            break;
                          }
                        } else if (data.data.status === "FAILED") {
                          throw new Error("Image generation failed");
                        }
                        // No else, we just wait and retry
                      } catch (error) {
                        console.error(
                          `Poll error for task ${taskId} (retry ${retries}):`,
                          error
                        );
                      }
                      await new Promise((r) => setTimeout(r, retryInterval));
                      retries++;
                    }
                    if (!imageSent) {
                      const errorData = {
                        id: `chatcmpl-${crypto.randomUUID()}`,
                        object: "chat.completion.chunk",
                        created: Math.floor(Date.now() / 1000),
                        model: model,
                        choices: [
                          {
                            delta: {
                              content: `![Image](Image generation failed for task ${taskId})\n`,
                            },
                            index: 0,
                            finish_reason: null,
                          },
                        ],
                      };
                      await writer.write(
                        new TextEncoder().encode(
                          `data: ${JSON.stringify(errorData)}\n\n`
                        )
                      );
                    }
                  }

                  // Send the completion signal
                  await writer.write(new TextEncoder().encode("data: [DONE]\n\n"));
                } catch (error) {
                  console.error("Stream error:", error);
                  await writer.abort(error);
                } finally {
                  await writer.close();
                }
              })()
            );

            return new Response(readable, {
              headers: {
                ...corsHeaders,
                "Content-Type": "text/event-stream",
                "Cache-Control": "no-cache",
                Connection: "keep-alive",
              },
            });
          } else {
            // 非流式处理
            // 发送请求获取任务ID
            const response = await fetch(
              "https://www.genspark.ai/api/copilot/ask",
              {
                method: "POST",
                headers,
                body: JSON.stringify(requestBody),
              }
            );
            if (!response.ok) {
              return createErrorResponse("上游服务请求失败", response.status);
            }
            // 解析响应获取 task_ids
            const text = await response.text();
            const taskIds = [];
            const lines = text.split("\n");
            for (const line of lines) {
              if (line.includes("task_id")) {
                try {
                  const content = line.replace("data: ", "");
                  if (!content) continue;
                  const { content: innerContent } = JSON.parse(content);
                  if (!innerContent) continue;
                  const { generated_images } = JSON.parse(innerContent);
                  if (!generated_images) continue;
                  for (const img of generated_images) {
                    if (img.task_id) {
                      taskIds.push(img.task_id);
                    }
                  }
                } catch (e) {}
              }
            }
            if (!taskIds.length) {
              return createErrorResponse("获取任务ID失败");
            }
            // 并行轮询任务结果
            const maxRetries = 30;
            const retryInterval = 2000;

            const taskPromises = taskIds.map((taskId) =>
              pollTaskId(taskId, session_id, maxRetries, retryInterval)
            );
            let imageUrls = [];
            try {
              const results = await Promise.all(taskPromises);
              results.forEach((urls) => {
                imageUrls.push(...urls);
              });
            } catch (error) {
              console.error("Polling error:", error);
              return createErrorResponse(error.message);
            }

            // 处理 response_format
            if (response_format === "b64_json") {
              // 使用 KV 缓存图像数据
              const imageDatas = await Promise.all(
                imageUrls.map(async (url) => {
                  const cachedData = await env.kv.get(url);
                  if (cachedData) {
                    return { b64_json: cachedData };
                  } else {
                    const imageResponse = await fetch(url);
                    const imageArrayBuffer = await imageResponse.arrayBuffer();
                    const base64Data = arrayBufferToBase64(imageArrayBuffer);
                    // 存储到 KV
                    await env.kv.put(url, base64Data, {
                      expirationTtl: CACHE_CONFIG.TTL,
                    });
                    return { b64_json: base64Data };
                  }
                })
              );

              // 构建聊天回复
              const imageMarkdowns = imageDatas
                .map(
                  (imgData, idx) =>
                    `![Image](data:image/png;base64,${imgData.b64_json})`
                )
                .join("\n");

              const messageId = crypto.randomUUID();

              const result = {
                id: `chatcmpl-${messageId}`,
                object: "chat.completion",
                created: Math.floor(Date.now() / 1000),
                model: model,
                choices: [
                  {
                    index: 0,
                    message: {
                      role: "assistant",
                      content: imageMarkdowns,
                    },
                    finish_reason: "stop",
                  },
                ],
                usage: {
                  prompt_tokens: 0,
                  completion_tokens: 0,
                  total_tokens: 0,
                },
              };

              // 缓存结果
              if (!CACHE_CONFIG.EXCLUDED_MODELS.includes(model)) {
                messageCache.set(cacheKey, JSON.stringify(result));
              }

              return new Response(JSON.stringify(result), {
                headers: {
                  ...corsHeaders,
                  "Content-Type": "application/json",
                },
              });
            } else {
              // 默认返回 URL，构建聊天回复
              const imageMarkdowns = imageUrls
                .map((url) => `![Image](${url})`)
                .join("\n");

              const messageId = crypto.randomUUID();

              const result = {
                id: `chatcmpl-${messageId}`,
                object: "chat.completion",
                created: Math.floor(Date.now() / 1000),
                model: model,
                choices: [
                  {
                    index: 0,
                    message: {
                      role: "assistant",
                      content: imageMarkdowns,
                    },
                    finish_reason: "stop",
                  },
                ],
                usage: {
                  prompt_tokens: 0,
                  completion_tokens: 0,
                  total_tokens: 0,
                },
              };

              // 缓存结果
              if (!CACHE_CONFIG.EXCLUDED_MODELS.includes(model)) {
                messageCache.set(cacheKey, JSON.stringify(result));
              }

              return new Response(JSON.stringify(result), {
                headers: {
                  ...corsHeaders,
                  "Content-Type": "application/json",
                },
              });
            }
          }
        }

        // 以下是处理非画图模型的逻辑

        let response;
        if (model === "genspark") {
          response = await searchModel(session_id, messages);
        } else {
          response = await makeRequest(
            session_id,
            model,
            messages,
            target_lang_code,
            source_lang_code
          );
        }

        if (!response.ok) {
          return createErrorResponse("上游服务请求失败", response.status);
        }

        const messageId = crypto.randomUUID();
        let project_id = null;

        if (stream === true || stream === "true") {
          const { readable, writable } = new TransformStream();
          const writer = writable.getWriter();
          ctx.waitUntil(
            (async () => {
              try {
                const reader = response.body.getReader();
                for await (const line of streamGenerator(reader)) {
                  if (!project_id) {
                    project_id = extractProjectId(line);
                  }
                  const processed = processChunk(line, messageId, model);
                  if (processed) {
                    await writer.write(new TextEncoder().encode(processed));
                  }
                }
                await writer.write(new TextEncoder().encode("data: [DONE]\n\n"));
              } catch (error) {
                console.error("Stream error:", error);
              } finally {
                await writer.close();
                if (project_id) {
                  await cleanupProject(session_id, project_id);
                }
              }
            })()
          );
          return new Response(readable, {
            headers: {
              ...corsHeaders,
              "Content-Type": "text/event-stream",
              "Cache-Control": "no-cache",
              Connection: "keep-alive",
            },
          });
        }

        const chunks = [];
        const reader = response.body.getReader();
        try {
          while (true) {
            const { done, value } = await reader.read();
            if (done) break;
            const text = new TextDecoder().decode(value);
            chunks.push(text);
            if (!project_id) {
              const lines = text.split("\n");
              for (const line of lines) {
                project_id = extractProjectId(line);
                if (project_id) break;
              }
            }
          }
        } finally {
          reader.releaseLock();
          if (project_id) {
            ctx.waitUntil(cleanupProject(session_id, project_id));
          }
        }

        const fullResponse = chunks.join("");
        for (const line of fullResponse.split("\n")) {
          if (line.startsWith("data: ")) {
            try {
              const content = JSON.parse(line.slice(6));
              if (model === "deepl") {
                // 处理翻译模型的响应
                if (content.field_name === "content" && content.field_value) {
                  const result = JSON.stringify({
                    code: 200,
                    data: {
                      translations: [
                        {
                          detected_source_language:
                            content.session_state?.detected_source_lang || "",
                          text: content.field_value,
                        },
                      ],
                    },
                    msg: "ok",
                  });
                  // 缓存结果
                  if (!CACHE_CONFIG.EXCLUDED_MODELS.includes(model)) {
                    messageCache.set(cacheKey, result);
                  }
                  return new Response(result, {
                    headers: { ...corsHeaders, "Content-Type": "application/json" },
                  });
                }
              } else {
                // 处理其他模型的响应
                if (
                  content.field_value &&
                  content.field_name !== "session_state.answer_is_finished" &&
                  content.field_name !== "content" &&
                  content.field_name !== "session_state" &&
                  !content.delta &&
                  content.type !== "project_field"
                ) {
                  const result = JSON.stringify({
                    id: `chatcmpl-${messageId}`,
                    object: "chat.completion",
                    created: Math.floor(Date.now() / 1000),
                    model,
                    choices: [
                      {
                        index: 0,
                        message: {
                          role: "assistant",
                          content: content.field_value,
                        },
                        finish_reason: "stop",
                      },
                    ],
                    usage: {
                      prompt_tokens: 0,
                      completion_tokens: 0,
                      total_tokens: content.field_value.length,
                    },
                  });
                  // 缓存结果
                  if (!CACHE_CONFIG.EXCLUDED_MODELS.includes(model)) {
                    messageCache.set(cacheKey, result);
                  }
                  return new Response(result, {
                    headers: { ...corsHeaders, "Content-Type": "application/json" },
                  });
                }
              }
            } catch (e) {}
          }
        }
        return createErrorResponse("无效的响应数据");
      } catch (error) {
        console.error("Request processing error:", error);
        return createErrorResponse("请求处理失败");
      }
    }

    // 添加新的画图路由
    if (
      url.pathname === "/v1/images/generations" &&
      request.method === "POST"
    ) {
      try {
        const requestData = await request.json();
        // 参数验证
        if (!requestData.prompt && !requestData.messages) {
          return new Response(
            JSON.stringify({ error: "Prompt is required" }),
            {
              status: 400,
              headers: { ...corsHeaders, "Content-Type": "application/json" },
            }
          );
        }
        // 获取随机session_id
        const authHeader = request.headers.get("authorization");
        if (sessionPool.size() === 0) {
          const sessionString = authHeader?.replace("Bearer ", "");
          sessionPool.initialize(sessionString);
          if (sessionPool.size() === 0) {
            return createErrorResponse("未提供有效的 session_id", 401);
          }
        }
        const session_id = sessionPool.getNext();
        if (!session_id) {
          return createErrorResponse("无可用的会话ID", 500);
        }
        // 提取 response_format
        const {
          response_format = "url",
          num_images = 1,
          size = "1024x1024",
          quality = "standard",
          style = "auto",
          image,
          model = "dall-e-3",
          aspect_ratio=auto,
          hd=true,
          reflection_enabled=false,
        } = requestData;
        // 构建请求体
        let messages = [];

        if (requestData.messages) {
          messages = requestData.messages;
          // 处理 messages，处理图片
          await processMessages(messages, env);
        } else {
          // 构建带有 prompt 和 image 的 messages
          if (image) {
            // image 可以是 base64 编码的图像数据
            let base64ImageData = image;
            if (!base64ImageData.startsWith('data:image')) {
              // 假设是 base64 编码的数据，包装为数据 URL
              base64ImageData = `data:image/jpeg;base64,${image}`;
            }
            messages.push({
              role: 'user',
              content: [
                { type: 'image_url', image_url: { url: base64ImageData } },
                { type: 'text', text: requestData.prompt },
              ],
            });
          } else {
            // 没有图片，仅使用提示词
            messages.push({ role: 'user', content: requestData.prompt });
          }
        }

        await processMessages(messages, env);
        const user_s_input = extractLastUserInput(messages);

        const requestBody = {
          type: "COPILOT_MOA_IMAGE",
          current_query_string: "type=COPILOT_MOA_IMAGE",
          messages: messages,
          user_s_input: user_s_input,
          action_params: {},
          extra_data: {
            model_configs: [],
            llm_model: "gpt-4o",
            imageModelMap: {},
            writingContent: null,
          },
        };
        for (let i = 0; i < num_images; i++) {
          requestBody.extra_data.model_configs.push({
            model: model === "dall-e-3" ? "dalle-3" : model,
            aspect_ratio,
            use_personalized_models:false,
            fashion_profile_id:null,
            hd,
            reflection_enabled,
            style: style
          });
        }
        // 发送请求获取任务ID
        const response = await fetch(
          "https://www.genspark.ai/api/copilot/ask",
          {
            method: "POST",
            headers: {
              Cookie: `session_id=${session_id}`,
              "User-Agent": "Apifox/1.0.0 (https://apifox.com)",
              "Content-Type": "application/json",
              Accept: "*/*",
              Host: "www.genspark.ai",
              Connection: "keep-alive",
            },
            body: JSON.stringify(requestBody),
          }
        );
        if (!response.ok) {
          return createErrorResponse("上游服务请求失败", response.status);
        }
        // 解析响应获取task_id
        const text = await response.text();
        const taskIds = [];
        const lines = text.split("\n");
        for (const line of lines) {
          if (line.includes("task_id")) {
            try {
              const content = line.replace("data: ", "");
              if (!content) continue;
              const { content: innerContent } = JSON.parse(content);
              if (!innerContent) continue;
              const { generated_images } = JSON.parse(innerContent);
              if (!generated_images) continue;
              for (const img of generated_images) {
                if (img.task_id) {
                  taskIds.push(img.task_id);
                }
              }
            } catch (e) {}
          }
        }
        if (!taskIds.length) {
          return createErrorResponse("获取任务ID失败");
        }
        // 并行轮询任务结果
        const maxRetries = 300;
        const retryInterval = 2000;

        const taskPromises = taskIds.map((taskId) =>
          pollTaskId(taskId, session_id, maxRetries, retryInterval)
        );
        let imageUrls = [];
        try {
          const results = await Promise.all(taskPromises);
          results.forEach((urls) => {
            imageUrls.push(...urls);
          });
        } catch (error) {
          console.error("Polling error:", error);
          return createErrorResponse(error.message);
        }

        if (response_format === "b64_json") {
          // 使用 KV 缓存图像数据
          const imageDatas = await Promise.all(
            imageUrls.map(async (url) => {
              const cachedData = await env.kv.get(url);
              if (cachedData) {
                return { b64_json: cachedData };
              } else {
                const imageResponse = await fetch(url);
                const imageArrayBuffer = await imageResponse.arrayBuffer();
                const base64Data = arrayBufferToBase64(imageArrayBuffer);
                // 存储到 KV
                await env.kv.put(url, base64Data, { expirationTtl: CACHE_CONFIG.TTL });
                return { b64_json: base64Data };
              }
            })
          );

          return new Response(
            JSON.stringify({
              created: Date.now(),
              data: imageDatas,
            }),
            { headers: { ...corsHeaders, "Content-Type": "application/json" } }
          );
        } else {
          // 默认返回 URL
          return new Response(
            JSON.stringify({
              created: Date.now(),
              data: imageUrls.map((url) => ({
                url,
                revised_prompt: requestData.prompt,
              })),
            }),
            { headers: { ...corsHeaders, "Content-Type": "application/json" } }
          );
        }
      } catch (error) {
        console.error("Request processing error:", error);
        return createErrorResponse("请求处理失败");
      }
    }


    return new Response("Not Found", {
      status: 404,
      headers: corsHeaders,
    });
  },
};