main.ts

import { serve } from "https://deno.land/std@0.208.0/http/server.ts";
import { decode } from "https://deno.land/std@0.208.0/encoding/base64.ts";

// --- 常量定义 ---
const MAX_DOCUMENT_SIZE_MB = 20; // 设置最大文档大小限制（单位：MB）
const MAX_DOCUMENT_SIZE_BYTES = MAX_DOCUMENT_SIZE_MB * 1024 * 1024;
const MODELS_CACHE_DURATION = 60000; // 1分钟模型缓存

interface OpenAIMessage {
  role: "system" | "user" | "assistant";
  content: string | Array<{
    type: string;
    text?: string;
    image_url?: { url: string };
    document?: { url: string; type: string }; // 支持多种文档类型
  }>;
}

interface OpenAIRequest {
  model: string;
  messages: OpenAIMessage[];
  max_tokens?: number;
  temperature?: number;
  stream?: boolean;
}

interface OpenAITTSRequest {
  model: string;
  input: string;
  voice: 'Zephyr' | 'Puck' | 'Charon' | 'Kore' | 'Fenrir' | 'Leda' | string;
}

class GoogleAIService {
  public apiKeys: string[];
  public currentKeyIndex = 0;
  public cachedModels: any[] = [];
  public modelsLastFetch = 0;

  constructor() {
    this.apiKeys = [];
    this.apiKeys = Deno.env.get(`GOOGLE_AI_KEYS`).split(',').map(s => s.trim());
    if (this.apiKeys.length === 0) {
      throw new Error("No Google AI API keys found in environment variables (e.g., GOOGLE_AI_KEYS)");
    }
  }

  private getNextApiKey(): string {
    const key = this.apiKeys[this.currentKeyIndex];
    console.log(key)
    this.currentKeyIndex = (this.currentKeyIndex + 1) % this.apiKeys.length;
    return key;
  }

  async fetchOfficialModels(): Promise<any[]> {
    const now = Date.now();
    if (this.cachedModels.length > 0 && (now - this.modelsLastFetch) < MODELS_CACHE_DURATION) {
      return this.cachedModels;
    }

    const apiKey = this.getNextApiKey();
    try {
      const response = await fetch(
        `https://generativelanguage.googleapis.com/v1beta/models?key=${apiKey}`,
        { method: "GET", headers: { "Content-Type": "application/json" } }
      );

      if (!response.ok) {
        console.warn(`Failed to fetch models from Google AI: ${response.status}. Using fallback models.`);
        return this.getFallbackModels();
      }

      const data = await response.json();
      if (data.models && Array.isArray(data.models)) {
        this.cachedModels = data.models.filter((model: any) =>
          model.supportedGenerationMethods?.includes('generateContent')
        );
        this.modelsLastFetch = now;
        this.cachedModels.push({
        "id": "gemini-2.0-flash-search",
        "name": "gemini-2.0-flash-search",
        "object": "model",
        "created": now,
        "owned_by": "google",
        "description": "Gemini 2.0 Flash with GoogleSearch",
        "maxTokens": 1048576
        })
        this.cachedModels.push({
        "id": "gemini-2.5-flash-search",
        "name": "gemini-2.5-flash-search",
        "object": "model",
        "created": now,
        "owned_by": "google",
        "description": "Gemini 2.5 Flash with GoogleSearch",
        "maxTokens": 1048576
        })
        this.cachedModels.push({
        "id": "gemini-2.5-pro-search",
        "name": "gemini-2.5-pro-search",
        "object": "model",
        "created": now,
        "owned_by": "google",
        "description": "Gemini 2.5 Pro with GoogleSearch",
        "maxTokens": 1048576
        })
        console.log(`Fetched ${this.cachedModels.length} models from Google AI`);
        return this.cachedModels;
      }
      return this.getFallbackModels();
    } catch (error) {
      console.warn("Error fetching models from Google AI:", error.message, ". Using fallback models.");
      return this.getFallbackModels();
    }
  }

  private getFallbackModels(): any[] {
    return [
      { name: "models/gemini-1.5-pro", displayName: "Gemini 1.5 Pro", description: "Mid-size multimodal model that supports up to 1 million tokens, images, and documents (PDF, TXT, MD)", supportedGenerationMethods: ["generateContent"], maxTokens: 1000000, supportsDocuments: true },
      { name: "models/gemini-1.5-flash", displayName: "Gemini 1.5 Flash", description: "Fast and versatile multimodal model for diverse tasks, supports images and documents (PDF, TXT, MD)", supportedGenerationMethods: ["generateContent"], maxTokens: 1000000, supportsDocuments: true },
      { name: "models/gemini-2.0-flash-preview-image-generation", displayName: "Gemini 2.0 Flash Image Generation", description: "Advanced model for generating and editing high-quality images with text and image outputs", supportedGenerationMethods: ["generateContent"], maxTokens: 100000, capabilities: ["text", "image_generation", "image_editing"] },
      { name: "models/gemini-2.5-flash-preview-tts", displayName: "Gemini 2.5 Flash TTS", description: "Advanced model for generating high-quality speech from text.", supportedGenerationMethods: ["generateContent"] },
    ];
  }

  public isVisionModel = (modelName: string): boolean => modelName.toLowerCase().includes('vision') || modelName.toLowerCase().includes('pro');
  public isImageGenerationModel = (modelName: string): boolean => modelName.includes('image') || modelName === 'gemini-2.0-flash-preview-image-generation' || modelName === 'gemini-2.5-flash-image-preview';
  public isImageEditingModel = (modelName: string): boolean => modelName.includes('image') || modelName === 'gemini-2.0-flash-preview-image-generation' || modelName === 'gemini-2.5-flash-image-preview';
  public isDocumentModel = (modelName: string): boolean => modelName.toLowerCase().includes('gemini-1.5') || modelName.toLowerCase().includes('pro') || modelName.toLowerCase().includes('flash');
  public isTTSModel = (modelName: string): boolean => modelName.toLowerCase().includes('tts');

  async generateSpeech(text: string, modelName: string, voiceName: string): Promise<string> {
    const apiKey = this.getNextApiKey();
    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;

    console.log(`Generating speech with model: ${fullModelName}, voice: ${voiceName}`);

    const requestBody = {
      contents: [{
        parts: [{ "text": text }]
      }],
      generationConfig: {
        responseModalities: ["AUDIO"],
        speechConfig: {
          voiceConfig: {
            prebuiltVoiceConfig: {
              voiceName: voiceName
            }
          }
        }
      },
      model: fullModelName,
    };

    const response = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
      {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(requestBody),
      }
    );

    if (!response.ok) {
      const errorBody = await response.json().catch(() => response.text());
      const errorMessage = errorBody?.error?.message || JSON.stringify(errorBody);
      console.error(`Google TTS API Error: ${response.status} - ${errorMessage}`);
      throw new Error(`Google TTS API request failed with status ${response.status}: ${errorMessage}`);
    }

    const data = await response.json();
    const audioData = data.candidates?.[0]?.content?.parts?.[0]?.inlineData?.data;
    
    if (!audioData) {
      console.error("Invalid TTS response from Google AI:", JSON.stringify(data));
      throw new Error("No audio data received from Google AI TTS service.");
    }

    return audioData;
  }

  private getDocumentType(url: string): string {
    const lowerUrl = url.toLowerCase();
    if (lowerUrl.startsWith('data:application/pdf') || lowerUrl.includes('.pdf')) return 'pdf';
    if (lowerUrl.startsWith('data:text/plain') || lowerUrl.includes('.txt')) return 'txt';
    if (lowerUrl.startsWith('data:text/markdown') || lowerUrl.includes('.md')) return 'md';
    if (lowerUrl.startsWith('data:application/msword') || lowerUrl.includes('.doc')) return 'doc';
    if (lowerUrl.startsWith('data:application/vnd.openxmlformats-officedocument.wordprocessingml.document') || lowerUrl.includes('.docx')) return 'docx';
    return 'unknown';
  }

  /**
   * [关键改进] 提取并验证文档数据，增加大小检查和更稳健的解析
   */
  private extractDocumentData(documentUrl: string): { mimeType: string; data: string; text?: string; docType: string } {
    const docType = this.getDocumentType(documentUrl);

    if (!documentUrl.startsWith("data:")) {
      if (documentUrl.startsWith("http")) {
        throw new Error("Document URL downloads are not supported. Please provide base64 encoded data URLs.");
      }
      // 如果不是data url或http url，则假定为纯base64数据，但这是一种不推荐的格式
      // 为了健壮性，我们强制要求使用标准的 data URL
      throw new Error("Document must be provided as a standard base64 data URL (e.g., 'data:application/pdf;base64,...').");
    }

    const parts = documentUrl.split(",");
    if (parts.length !== 2) {
        throw new Error("Invalid data URL format for document. Expected 'data:[mime];base64,[data]'.");
    }
    const [mimeInfo, base64Data] = parts;

    // **改进1: 检查文件大小**
    // Base64 字符串的长度约是原始数据的 4/3。
    const approxSizeInBytes = base64Data.length * 0.75;
    if (approxSizeInBytes > MAX_DOCUMENT_SIZE_BYTES) {
        throw new Error(`Document size (${(approxSizeInBytes / 1024 / 1024).toFixed(2)}MB) exceeds the ${MAX_DOCUMENT_SIZE_MB}MB limit.`);
    }

    const mimeType = mimeInfo.split(":")[1]?.split(";")[0] || 'application/octet-stream';

    if (docType === 'txt' || docType === 'md') {
      try {
        const textContent = atob(base64Data);
        return { mimeType, data: base64Data, text: textContent, docType };
      } catch (error) {
        console.error(`Failed to decode base64 content for ${docType}:`, error);
        throw new Error(`Invalid base64 encoding for ${docType} document.`);
      }
    }
    
    // 自动识别PDF的MIME类型
    const finalMimeType = docType === 'pdf' ? 'application/pdf' : mimeType;
    return { mimeType: finalMimeType, data: base64Data, docType };
  }
  
  private extractImageData(imageUrl: string): { mimeType: string; data: string } {
    if (imageUrl.startsWith("data:image/")) {
      const [mimeInfo, base64Data] = imageUrl.split(",");
      const mimeType = mimeInfo.split(":")[1].split(";")[0];
      return { mimeType, data: base64Data };
    } else if (imageUrl.startsWith("http")) {
      throw new Error("URL images are not supported yet. Please provide base64 encoded images.");
    } else {
      return { mimeType: "image/jpeg", data: imageUrl };
    }
  }

  async generateContentWithDocument(messages: OpenAIMessage[], modelName: string, maxTokens?: number): Promise<string> {
    const apiKey = this.getNextApiKey();
    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
    const documentModel = this.isDocumentModel(fullModelName) ? fullModelName : 'models/gemini-1.5-pro-latest';

    console.log(`Processing document with model: ${documentModel}`);

    let contents;
    try {
      contents = messages.map(msg => {
        if (typeof msg.content === "string") {
          return { role: msg.role === "assistant" ? "model" : "user", parts: [{ text: msg.content }] };
        }

        const messageParts = msg.content.map(part => {
          if (part.type === "text") return { text: part.text };

          if (part.type === "image_url" && part.image_url) {
            const { mimeType, data } = this.extractImageData(part.image_url.url);
            return { inlineData: { mimeType, data } };
          }

          if (part.type === "document" && part.document) {
            const docData = this.extractDocumentData(part.document.url);
            console.log(`Processing document: ${docData.docType}, mime: ${docData.mimeType}, size: ${(docData.data.length * 0.75 / 1024).toFixed(2)} KB`);

            if (docData.docType === 'txt' || docData.docType === 'md') {
              const prefix = docData.docType === 'md' ? 'Markdown document content:\n' : 'Text document content:\n';
              return { text: `${prefix}${docData.text}` };
            }
            if (docData.docType === 'pdf') {
              return { inlineData: { mimeType: docData.mimeType, data: docData.data } };
            }
            return { text: `[Document type '${docData.docType}' is not supported for direct processing. Please convert to PDF, TXT, or MD.]` };
          }
          return { text: "" };
        });
        return { role: msg.role === "assistant" ? "model" : "user", parts: messageParts.filter(p => p.text || p.inlineData) };
      });
    } catch (error) {
      throw error;
    }

    const requestBody = {
      contents,
      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
    };

    const response = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/${documentModel}:generateContent?key=${apiKey}`,
      {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(requestBody),
      }
    );

    if (!response.ok) {
      const errorBody = await response.json().catch(() => response.text());
      const errorMessage = errorBody?.error?.message || JSON.stringify(errorBody);
      console.error(`Google API Error: ${response.status} - ${errorMessage}`);
      throw new Error(`Google API request failed with status ${response.status}: ${errorMessage}`);
    }

    const data = await response.json();
    const promptFeedback = data.promptFeedback;
    if (promptFeedback && promptFeedback.blockReason) {
      const reason = promptFeedback.blockReason;
      const safetyRatings = promptFeedback.safetyRatings?.map((r: any) => `${r.category}: ${r.probability}`).join(', ') || 'N/A';
      throw new Error(`Request blocked by Google API. Reason: ${reason}. Safety Ratings: [${safetyRatings}]`);
    }

    if (!data.candidates || data.candidates.length === 0) {
      throw new Error("No response generated for document content. The content might be empty or unreadable.");
    }

    const candidate = data.candidates[0];
    if (candidate.finishReason === "SAFETY") {
        throw new Error("Response blocked due to safety filters. Check content for sensitive topics.");
    }
    if (candidate.finishReason === "RECITATION") {
        throw new Error("Response blocked due to recitation policy. The model's output was too similar to a copyrighted source.");
    }

    return candidate.content?.parts[0]?.text || "Document processed, but no text response was generated.";
  }
  
  // The rest of the original methods from the user's code
  async generateContent(messages: OpenAIMessage[], modelName: string, maxTokens?: number, enableSearch: boolean = false): Promise<string> {
    const hasDocument = messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "document"));
    if (hasDocument) {
      return await this.generateContentWithDocument(messages, modelName, maxTokens);
    }

    const apiKey = this.getNextApiKey();
    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;

    const contents = messages.map(msg => {
      if (typeof msg.content === "string") {
        return { role: msg.role === "assistant" ? "model" : "user", parts: [{ text: msg.content }] };
      } else {
        const messageParts = msg.content.map(part => {
          if (part.type === "text") {
            return { text: part.text };
          } else if (part.type === "image_url" && part.image_url) {
            const imageData = part.image_url.url;
            if (imageData.startsWith("data:image/")) {
              const { mimeType, data } = this.extractImageData(imageData);
              return { inlineData: { mimeType, data } };
            } else {
              return { fileData: { mimeType: "image/jpeg", fileUri: imageData } };
            }
          }
          return { text: "" };
        });
        return { role: msg.role === "assistant" ? "model" : "user", parts: messageParts };
      }
    });

    const requestBody: any = {
      contents,
      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
    };
    if (enableSearch) {
      requestBody.tools = [{ googleSearchRetrieval: {} }];
    }

    const response = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
    );

    if (!response.ok) {
      const errorText = await response.text();
      throw new Error(`Google AI API error: ${response.status} - ${errorText}`);
    }
    const data = await response.json();
    if (!data.candidates || data.candidates.length === 0) {
      throw new Error("No response generated from Google AI");
    }
    const candidate = data.candidates[0];
    if (candidate.finishReason === "SAFETY") {
      throw new Error("Response blocked due to safety filters");
    }
    return candidate.content?.parts[0]?.text || "No response generated";
  }

  async generateOrEditImageWithGemini(prompt: string, modelName: string = "gemini-2.0-flash-preview-image-generation", inputImage?: { mimeType: string; data: string }): Promise<{ text?: string; imageBase64?: string; imageUrl?: string }> {
    const apiKey = this.getNextApiKey();
    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
    const requestParts: any[] = [{ text: prompt }];

    if (inputImage) {
      requestParts.push({ inline_data: { mime_type: inputImage.mimeType, data: inputImage.data } });
      console.log(`Editing image with model: ${fullModelName}`);
    } else {
      console.log(`Generating image with model: ${fullModelName}`);
    }

    const requestBody = {
      contents: [{ parts: requestParts }],
      generationConfig: { responseModalities: ["TEXT", "IMAGE"], temperature: 0.7 }
    };

    const response = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
    );

    if (!response.ok) {
      const errorText = await response.text();
      throw new Error(`Image ${inputImage ? 'editing' : 'generation'} failed: ${response.status} - ${errorText}`);
    }
    const data = await response.json();
    if (!data.candidates || data.candidates.length === 0) {
      throw new Error(`No ${inputImage ? 'edited' : 'generated'} image returned`);
    }

    const candidate = data.candidates[0];
    if (candidate.finishReason === "SAFETY") {
      throw new Error(`Image ${inputImage ? 'editing' : 'generation'} blocked due to safety filters`);
    }

    const responseParts = candidate.content?.parts || [];
    let textResponse = "";
    let imageBase64 = "";

    for (const part of responseParts) {
      if (part.text) textResponse += part.text;
      if (part.inlineData?.data) imageBase64 = part.inlineData.data;
      if (part.inline_data?.data) imageBase64 = part.inline_data.data;
    }

    const result: { text?: string; imageBase64?: string; imageUrl?: string } = {};
    if (textResponse) result.text = textResponse;
    if (imageBase64) {
      result.imageBase64 = imageBase64;
      result.imageUrl = `data:image/png;base64,${imageBase64}`;
    }
    return result;
  }
    
  async generateContentWithGrounding(messages: OpenAIMessage[], modelName: string, maxTokens?: number): Promise<string> {
    const apiKey = this.getNextApiKey();
    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
    const contents = messages.map(msg => ({ role: msg.role === 'assistant' ? 'model' : 'user', parts: [{ text: typeof msg.content === 'string' ? msg.content : '' }] }));

    const requestBody = {
      contents,
      tools: [{ googleSearch: {} }],
      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
    };
    const response = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
    );
    if (!response.ok) {
      console.warn(`Google Search API failed: ${response.status}, trying alternative.`);
      return await this.generateContentWithSearchPrompt(messages, modelName, maxTokens);
    }
    const data = await response.json();
    if (!data.candidates || data.candidates.length === 0) {
      return await this.generateContentWithSearchPrompt(messages, modelName, maxTokens);
    }
    
    const candidate = data.candidates[0];
    if (candidate.finishReason === "SAFETY") {
      throw new Error("Response blocked due to safety filters");
    }
    return candidate.content?.parts[0]?.text || "No response generated";
  }

  async generateContentWithSearchPrompt(messages: OpenAIMessage[], modelName: string, maxTokens?: number): Promise<string> {
    const enhancedMessages = [...messages];
    const lastMessage = enhancedMessages[enhancedMessages.length - 1];
    if (typeof lastMessage.content === "string") {
      lastMessage.content = `Please provide the most current and accurate information available about: ${lastMessage.content}.`;
    }
    return await this.generateContent(enhancedMessages, modelName, maxTokens, false);
  }

  async generateOrEditImage(prompt: string, modelName: string, inputImages?: any[]): Promise<string> {
    if (this.isImageGenerationModel(modelName)) {
      try {
        let inputImage: { mimeType: string; data: string } | undefined;
        if (inputImages && inputImages.length > 0) {
          inputImage = this.extractImageData(inputImages[0].url);
        }
        const result = await this.generateOrEditImageWithGemini(prompt, modelName, inputImage);
        let response = "";
        if (result.text) response += result.text + "\\\\n\\\\n";
        if (result.imageUrl) response += `![image](${result.imageUrl})`;
        return response || `Image processing complete.`;
      } catch (error) {
        return `Image processing failed: ${error.message}`;
      }
    }
    return `Model ${modelName} does not support image generation. Use a model like gemini-2.0-flash-preview-image-generation.`;
  }
}

class OpenAICompatibleServer {
  public googleAI: GoogleAIService;
  private authKey: string;

  constructor() {
    this.googleAI = new GoogleAIService();
    this.authKey = Deno.env.get("AUTH_KEY") || "";
  }

  private _writeString(view: DataView, offset: number, str: string) {
    for (let i = 0; i < str.length; i++) {
      view.setUint8(offset + i, str.charCodeAt(i));
    }
  }
  
  private _createWavFile(pcmData: Uint8Array): Uint8Array {
    const numChannels = 1;
    const sampleRate = 24000;
    const bitsPerSample = 16;
    const dataSize = pcmData.length;
    const headerSize = 44;
    const buffer = new ArrayBuffer(headerSize + dataSize);
    const view = new DataView(buffer);
    
    this._writeString(view, 0, "RIFF");
    view.setUint32(4, 36 + dataSize, true);
    this._writeString(view, 8, "WAVE");
    this._writeString(view, 12, "fmt ");
    view.setUint32(16, 16, true);
    view.setUint16(20, 1, true);
    view.setUint16(22, numChannels, true);
    view.setUint32(24, sampleRate, true);
    view.setUint32(28, sampleRate * numChannels * (bitsPerSample / 8), true);
    view.setUint16(32, numChannels * (bitsPerSample / 8), true);
    view.setUint16(34, bitsPerSample, true);
    this._writeString(view, 36, "data");
    view.setUint32(40, dataSize, true);

    const wavBytes = new Uint8Array(buffer);
    wavBytes.set(pcmData, headerSize);
    return wavBytes;
  }

  private authenticate(request: Request): boolean {
    if (!this.authKey) return true;
    const authHeader = request.headers.get("Authorization");
    return authHeader ? authHeader.replace("Bearer ", "") === this.authKey : false;
  }

  private async handleAudioSpeech(request: Request): Promise<Response> {
      try {
        const body: OpenAITTSRequest = await request.json();
        const modelMap: { [key: string]: string } = { 'tts-1': 'gemini-2.5-flash-preview-tts', 'tts-1-hd': 'gemini-2.5-flash-preview-tts' };
        const geminiModel = modelMap[body.model] || (this.googleAI.isTTSModel(body.model) ? body.model : 'gemini-2.5-flash-preview-tts');
        const voiceMap: { [key: string]: string } = { 'alloy': 'Krew', 'echo': 'Kore', 'fable': 'Chiron', 'onyx': 'Calypso', 'nova': 'Cria', 'shimmer': 'Estrella' };
        const geminiVoice = voiceMap[body.voice] || 'Kore';
  
        if (!body.input) throw new Error("The 'input' field is required for TTS requests.");
  
        const audioBase64 = await this.googleAI.generateSpeech(body.input, geminiModel, geminiVoice);
        const pcmBytes = decode(audioBase64);
        const wavBytes = this._createWavFile(pcmBytes);
  
        return new Response(wavBytes, { headers: { "Content-Type": "audio/wav" } });
      } catch (error) {
        console.error("Error in audio speech generation:", error.message);
        const status = error.message.includes("required") ? 400 : 500;
        return new Response(JSON.stringify({ error: { message: error.message, type: status === 400 ? "invalid_request_error" : "api_error", code: "tts_failed" } }), { status, headers: { "Content-Type": "application/json" } });
      }
  }

  private isDocumentContent(url?: string): boolean {
    if (!url) return false;
    const lowerUrl = url.toLowerCase();
    return lowerUrl.includes('.pdf') || lowerUrl.startsWith('data:application/pdf') ||
           lowerUrl.includes('.txt') || lowerUrl.startsWith('data:text/plain') ||
           lowerUrl.includes('.md') || lowerUrl.startsWith('data:text/markdown');
  }

  private async handleChatCompletions(request: Request): Promise<Response> {
    try {
      const body: OpenAIRequest = await request.json();
      const requestedModel = body.model || "gemini-1.5-pro";
      const stream = body.stream || false;
      const maxTokens = body.max_tokens || 1048576;
      console.log(`Request for model: ${requestedModel}, stream: ${stream}, max_tokens: ${maxTokens}`);
      const lastMessage = body.messages[body.messages.length - 1];
      const content = typeof lastMessage.content === "string"
        ? lastMessage.content
        : (Array.isArray(lastMessage.content) ? lastMessage.content.map(p => p.text || "").join(" ") : "");
      if (content == 'ping'){
        const responsePayload = {
          id: `chatcmpl-${Date.now()}`, object: "chat.completion", created: Math.floor(Date.now() / 1000), model: requestedModel,
          choices: [{ index: 0, message: { role: "assistant", content: "pong" }, finish_reason: "stop" }],
          usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }
        };
        return new Response(JSON.stringify(responsePayload), { headers: { "Content-Type": "application/json" } });
      }
      const hasDocument = body.messages.some(msg =>
        Array.isArray(msg.content) &&
        msg.content.some(part => part.type === "document" || this.isDocumentContent(part.document?.url))
      );
      const hasImages = body.messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "image_url"));
      
      let inputImages: any[] = [];
      if (hasImages) {
        body.messages.forEach(msg => {
          if (Array.isArray(msg.content)) {
            msg.content.forEach(part => {
              if (part.type === "image_url" && part.image_url) inputImages.push({ url: part.image_url.url });
            });
          }
        });
      }
      let responseText: string;

      // Routing logic based on keywords and content types
      if (hasDocument) {
        responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel, maxTokens);
      } else if (this.googleAI.isImageEditingModel(requestedModel) && hasImages) {
        responseText = await this.googleAI.generateOrEditImage(content, requestedModel, inputImages);
      } else if (this.googleAI.isImageGenerationModel(requestedModel)) {
        responseText = await this.googleAI.generateOrEditImage(content, requestedModel);
      } else if (requestedModel.endsWith("-search")) {
        const searchMessages = [{ ...lastMessage, content: content }];
        responseText = await this.googleAI.generateContentWithGrounding(searchMessages, requestedModel.slice(0, -"-search".length), maxTokens);
      } else {
        responseText = await this.googleAI.generateContent(body.messages, requestedModel, maxTokens, false);
      }

      if (stream) {
        const streamResponse = await this.streamStringAsOpenAIResponse(responseText, requestedModel);
        return new Response(streamResponse, {
          headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive", "Access-Control-Allow-Origin": "*" }
        });
      } else {
        const responsePayload = {
          id: `chatcmpl-${Date.now()}`, object: "chat.completion", created: Math.floor(Date.now() / 1000), model: requestedModel,
          choices: [{ index: 0, message: { role: "assistant", content: responseText }, finish_reason: "stop" }],
          usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }
        };
        return new Response(JSON.stringify(responsePayload), { headers: { "Content-Type": "application/json" } });
      }
    } catch (error) {
      console.error("Error in chat completions:", error.message);
      const status = error.message.includes("exceeds the limit") || error.message.includes("Invalid") ? 400 : 500;
      return new Response(
        JSON.stringify({
          error: {
            message: error.message,
            type: status === 400 ? "invalid_request_error" : "api_error",
            code: null
          }
        }),
        { status, headers: { "Content-Type": "application/json" } }
      );
    }
  }
  
  private async streamStringAsOpenAIResponse(content: string, modelName: string): Promise<ReadableStream<Uint8Array>> {
    const encoder = new TextEncoder();
    const streamId = `chatcmpl-${Date.now()}`;
    const creationTime = Math.floor(Date.now() / 1000);
    const chunkSize = 256; // 设置块大小为256个字符
    let position = 0;

    return new ReadableStream({
        start(controller) {
            const initialChunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: { role: 'assistant', content: '' }, finish_reason: null }] };
            controller.enqueue(encoder.encode(`data: ${JSON.stringify(initialChunk)}\n\n`));
        },
        pull(controller) {
            if (position >= content.length) {
                const finalChunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] };
                controller.enqueue(encoder.encode(`data: ${JSON.stringify(finalChunk)}\n\n`));
                controller.enqueue(encoder.encode('data: [DONE]\n\n'));
                controller.close();
                return;
            }
            
            const chunkContent = content.substring(position, position + chunkSize);
            position += chunkSize;

            const chunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: { content: chunkContent }, finish_reason: null }] };
            controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`));
        }
    });
  }
  
  private async handleModels(): Promise<Response> {
    try {
      const googleModels = await this.googleAI.fetchOfficialModels();
      const models = {
        object: "list",
        data: googleModels.map(model => {
          const modelId = model.name.replace('models/', '');
          return {
            id: modelId, object: "model", created: Math.floor(Date.now() / 1000), owned_by: "google",
            description: model.description || model.displayName, maxTokens: model.inputTokenLimit || model.maxTokens
          };
        })
      };
      return new Response(JSON.stringify(models), { headers: { "Content-Type": "application/json" } });
    } catch (error) {
      console.error("Error fetching models:", error);
      return new Response(JSON.stringify({ error: { message: "Failed to fetch models." } }), { status: 500 });
    }
  }
  
  private async handleStatus(): Promise<Response> {
      const status = {
          status: "healthy", timestamp: new Date().toISOString(), version: "2.5.0",
          api_keys_loaded: this.googleAI.apiKeys.length,
          models_in_cache: this.googleAI.cachedModels.length,
          models_last_fetched: this.googleAI.modelsLastFetch > 0 ? new Date(this.googleAI.modelsLastFetch).toISOString() : "never"
      };
      return new Response(JSON.stringify(status), { headers: { "Content-Type": "application/json" } });
  }

  async handleRequest(request: Request): Promise<Response> {
    const corsHeaders = {
      "Access-Control-Allow-Origin": "*",
      "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
      "Access-Control-Allow-Headers": "Content-Type, Authorization",
    };

    if (request.method === "OPTIONS") {
      return new Response(null, { headers: corsHeaders });
    }

    const url = new URL(request.url);
    let response: Response;

    // Handle routes
    if (url.pathname === "/health" || url.pathname === "/status") {
      response = await this.handleStatus();
    } else if (!this.authenticate(request)) {
      response = new Response(JSON.stringify({ error: { message: "Unauthorized" } }), { status: 401 });
    } else if (url.pathname === "/v1/audio/speech" && request.method === "POST") {
      response = await this.handleAudioSpeech(request);
    } else if (url.pathname === "/v1/chat/completions" && request.method === "POST") {
      response = await this.handleChatCompletions(request);
    } else if (url.pathname === "/v1/models" && request.method === "GET") {
      response = await this.handleModels();
    } else {
      response = new Response("Not Found", { status: 404 });
    }

    // Add CORS headers to all responses
    const finalHeaders = new Headers(response.headers);
    for (const [key, value] of Object.entries(corsHeaders)) {
      finalHeaders.set(key, value);
    }

    return new Response(response.body, { status: response.status, headers: finalHeaders });
  }
}

// --- 服务器启动 ---
const server = new OpenAICompatibleServer();

console.log("🚀 OpenAI Compatible Server with Google AI starting on port 8000...");
console.log(`✅ Loaded ${server.googleAI.apiKeys.length} API key(s).`);
console.log(`📄 Max document size set to ${MAX_DOCUMENT_SIZE_MB}MB.`);

// Pre-fetch models at startup
server.googleAI.fetchOfficialModels().then(models => {
  console.log(`✅ Successfully fetched ${models.length} models from Google AI.`);
}).catch(error => {
  console.warn(`⚠️ Could not pre-fetch models: ${error.message}. Will use fallbacks or fetch on first request.`);
});

console.log("\n🔗 Endpoints:");
console.log("   POST /v1/chat/completions");
console.log("   POST /v1/audio/speech");
console.log("   GET  /v1/models");
console.log("   GET  /status");

await serve(
  (request: Request) => server.handleRequest(request),
  { port: 7860 }
);