main.ts

import { serve } from "https://deno.land/std@0.140.0/http/server.ts";
import { Md5 } from "https://deno.land/std@0.140.0/hash/md5.ts";

const API_DOMAIN = 'https://ai-api.dangbei.net';
const USER_AGENTS = [
  'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36',
  'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36',
  'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36',
  'Mozilla/5.0 (iPhone; CPU iPhone OS 16_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.0 Mobile/15E148 Safari/604.1',
  'Mozilla/5.0 (iPad; CPU OS 16_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.0 Mobile/15E148 Safari/604.1'
];
const VALID_API_KEY = Deno.env.get('VALID_API_KEY');
const MAX_CONVERSATIONS_PER_DEVICE = 10;  //  每个设备最多创建的会话数 

class ChatManage {
  private currentDeviceId: string | null = null;
  private currentConversationId: string | null = null;
  private conversationCount = 0;
  private currentUserAgent: string;

  constructor() {
    this.currentUserAgent = this.getRandomUserAgent();
  }

  private getRandomUserAgent(): string {
    return USER_AGENTS[Math.floor(Math.random() * USER_AGENTS.length)];
  }

  getOrCreateIds(forceNew = false) {
    let newDeviceId = this.currentDeviceId;
    let newConversationId = this.currentConversationId;
  
    if (forceNew || !newDeviceId || this.conversationCount >= MAX_CONVERSATIONS_PER_DEVICE) {
      newDeviceId = this.generateDeviceId();
      newConversationId = null;
      this.conversationCount = 0;
      // 在生成新设备ID时更新 User-Agent
      this.currentUserAgent = this.getRandomUserAgent();
    }
  
    this.currentDeviceId = newDeviceId;
    this.currentConversationId = newConversationId;
  
    return {
      deviceId: newDeviceId,
      conversationId: newConversationId,
      userAgent: this.currentUserAgent
    };
  }

  updateConversationId(conversationId: string) {
    this.currentConversationId = conversationId;
    this.conversationCount++;
  }
  
  generateDeviceId() {
    const uuid = crypto.randomUUID();
    const urlAlphabet = 'useandom26T198340PX75pxJACKVERYMINDBUSHWOLF_GQZbfghjklqvwyzrict';
    const nanoid = Array.from(crypto.getRandomValues(new Uint8Array(20)))
      .map(b => urlAlphabet[b % urlAlphabet.length])
      .join('');
    return `${uuid.replace(/-/g, '')}_${nanoid}`;
  }
}

class Pipe {
  private dataPrefix = 'data:';
  private chatManage = new ChatManage();
  private searchModels: Record<string, string> = {
    'DeepSeek-R1-Search': 'deepseek',
    'DeepSeek-V3-Search': 'deepseek',
    'Doubao-Search': 'doubao',
    'Qwen-Search': 'qwen'
  };

  // 创建新的会话
  async _create_conversation(deviceId: string) {
    const { userAgent } = this.chatManage.getOrCreateIds(false);
    const payload = { botCode: "AI_SEARCH" };
    const timestamp = Math.floor(Date.now() / 1000).toString();
    const nonce = this.nanoid(21);
    const sign = await this.generateSign(timestamp, payload, nonce);

    const headers = {
      "Origin": "https://ai.dangbei.com",
      "Referer": "https://ai.dangbei.com/",
      "User-Agent": userAgent,
      "deviceId": deviceId,
      "nonce": nonce,
      "sign": sign,
      "timestamp": timestamp,
      "Content-Type": "application/json"
    };

    try {
      console.log('Creating conversation with:', {
        url: `${API_DOMAIN}/ai-search/conversationApi/v1/create`,
        headers,
        payload
      });

      const response = await fetch(`${API_DOMAIN}/ai-search/conversationApi/v1/create`, {
        method: 'POST',
        headers,
        body: JSON.stringify(payload),
      });

      console.log('Response status:', response.status);
      const responseText = await response.text();
      console.log('Response body:', responseText);

      if (response.ok) {
        try {
          const data = JSON.parse(responseText);
          if (data.success) {
            console.log('Successfully created conversation:', data.data.conversationId);
            return data.data.conversationId;
          } else {
            console.error('API returned success: false:', data);
          }
        } catch (e) {
          console.error('Failed to parse response:', e);
        }
      } else {
        console.error('HTTP error:', response.status, responseText);
      }
    } catch (e) {
      console.error('Error creating conversation:', e);
    }
    return null;
  }

  // 新增方法：构建完整提示
  _buildFullPrompt(messages: any[]): string {
    if (!messages || messages.length === 0) {
      return '';
    }

    let systemPrompt = '';
    const history: string[] = [];
    let lastUserMessage = '';

    for (const msg of messages) {
      if (msg.role === 'system' && !systemPrompt) {
        systemPrompt = msg.content;
      } else if (msg.role === 'user') {
        history.push(`user: ${msg.content}`);
        lastUserMessage = msg.content;
      } else if (msg.role === 'assistant') {
        history.push(`assistant: ${msg.content}`);
      }
    }

    const parts: string[] = [];
    if (systemPrompt) {
      parts.push(`[System Prompt]\n${systemPrompt}`);
    }
    if (history.length > 1) {
      parts.push(`[Chat History]\n${history.slice(0, -1).join('\n')}`);
    }
    parts.push(`[Question]\n${lastUserMessage}`);

    return parts.join('\n\n');
  }

  async* pipe(body: any) {
    const thinkingState = { thinking: -1 };
    
    // Build full prompt
    const fullPrompt = this._buildFullPrompt(body.messages);
    
    // Check if we need to force new conversation
    let forceNew = false;
    const messages = body.messages;
    if (messages.length === 1) {
      forceNew = true;
    } else if (messages.length >= 2) {
      const lastTwo = messages.slice(-2);
      if (lastTwo[0].role === 'user' && lastTwo[1].role === 'user') {
        forceNew = true;
      }
    }

    // Get or create device ID and conversation ID with User-Agent
    const { deviceId, conversationId: storedConversationId, userAgent } = this.chatManage.getOrCreateIds(forceNew);
    let conversationId = storedConversationId;

    // Create new conversation if needed
    if (!conversationId) {
      conversationId = await this._create_conversation(deviceId);
      if (!conversationId) {
        yield { error: 'Failed to create conversation' };
        return;
      }
      this.chatManage.updateConversationId(conversationId);
    }

    // Model name handling
    let modelName;
    const isSearchModel = body.model.endsWith('-Search');
    if (isSearchModel) {
      modelName = this.searchModels[body.model] || body.model.replace('-Search', '').toLowerCase();
    } else {
      const isDeepSeekModel = ['DeepSeek-R1', 'DeepSeek-V3'].includes(body.model);
      modelName = isDeepSeekModel ? 'deepseek' : body.model.toLowerCase();
    }

    // 确定 userAction 参数
    let userAction = '';
    if (body.model.includes('DeepSeek-R1')) {
      userAction = 'deep';
    }
    if (isSearchModel) {
      userAction = userAction ? `${userAction},online` : 'online';
    }

    const payload = {
      stream: true,
      botCode: 'AI_SEARCH',
      userAction,
      model: modelName,
      conversationId: conversationId,
      question: fullPrompt,
    };

    const timestamp = Math.floor(Date.now() / 1000).toString();
    const nonce = this.nanoid(21);
    const sign = await this.generateSign(timestamp, payload, nonce);

    const headers = {
      'Origin': 'https://ai.dangbei.com',
      'Referer': 'https://ai.dangbei.com/',
      'User-Agent': userAgent,
      'deviceId': deviceId,
      'nonce': nonce,
      'sign': sign,
      'timestamp': timestamp,
      'Content-Type': 'application/json',
    };

    try {
      const response = await fetch(`${API_DOMAIN}/ai-search/chatApi/v1/chat`, {
        method: 'POST',
        headers,
        body: JSON.stringify(payload),
      });

      if (!response.ok) {
        const error = await response.text();
        console.error('HTTP Error:', response.status, error);
        yield { error: `HTTP ${response.status}: ${error}` };
        return;
      }

      const reader = response.body!.getReader();
      const decoder = new TextDecoder();
      let buffer = '';
      let cardMessages: string[] = [];

      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.startsWith(this.dataPrefix)) continue;

          try {
            const data = JSON.parse(line.slice(this.dataPrefix.length));
            if (data.type === 'answer') {
              const content = data.content;
              const contentType = data.content_type;

              if (thinkingState.thinking === -1 && contentType === 'thinking') {
                thinkingState.thinking = 0;
                yield { choices: [{ delta: { content: '<think>\n\n' }, finish_reason: null }] };
              } else if (thinkingState.thinking === 0 && contentType === 'text') {
                thinkingState.thinking = 1;
                yield { choices: [{ delta: { content: '\n' }, finish_reason: null }] };
                yield { choices: [{ delta: { content: '</think>' }, finish_reason: null }] };
                yield { choices: [{ delta: { content: '\n\n' }, finish_reason: null }] };
              }

              if (contentType === 'card') {
                try {
                  const cardContent = JSON.parse(content);
                  const cardItems = cardContent.cardInfo.cardItems;
                  let markdownOutput = '\n\n---\n\n';

                  const searchKeywords = cardItems.find((item: any) => item.type === '2001');
                  if (searchKeywords) {
                    const keywords = JSON.parse(searchKeywords.content);
                    markdownOutput += `搜索关键字：${keywords.join('; ')}\n\n`;
                  }

                  const searchResults = cardItems.find((item: any) => item.type === '2002');
                  if (searchResults) {
                    const results = JSON.parse(searchResults.content);
                    markdownOutput += `共找到 ${results.length} 个搜索结果：\n\n`;
                    
                    results.forEach((result: any) => {
                      markdownOutput += `[${result.idIndex}] [${result.name}](${result.url})  来源：${result.siteName}\n`;
                    });
                  }

                  cardMessages.push(markdownOutput);
                } catch (e) {
                  console.error('Error processing card:', e);
                }
              }

              if (content && (contentType === 'text' || contentType === 'thinking')) {
                yield { choices: [{ delta: { content }, finish_reason: null }] };
              }
            }
          } catch (e) {
            console.error('Parse error:', e, 'Line:', line);
            yield { error: `JSONDecodeError: ${(e as Error).message}` };
            return;
          }
        }
      }

      if (cardMessages.length > 0) {
        yield { choices: [{ delta: { content: cardMessages.join('') }, finish_reason: null }] };
      }

      yield {
        choices: [{
          delta: {
            meta: {
              device_id: deviceId,
              conversation_id: conversationId
            }
          },
          finish_reason: null
        }]
      };

    } catch (e) {
      console.error('Error in pipe:', e);
      yield { error: `${(e as Error).name}: ${(e as Error).message}` };
    }
  }

  nanoid(size = 21) {
    const urlAlphabet = 'useandom26T198340PX75pxJACKVERYMINDBUSHWOLF_GQZbfghjklqvwyzrict';
    const bytes = new Uint8Array(size);
    crypto.getRandomValues(bytes);
    return Array.from(bytes).reverse().map(b => urlAlphabet[b & 63]).join('');
  }

  async generateSign(timestamp: string, payload: any, nonce: string) {
    const payloadStr = JSON.stringify(payload);
    const signStr = `${timestamp}${payloadStr}${nonce}`;
    console.log('Sign string:', signStr);
    
    // 使用 Deno 标准库的 MD5
    const sign = new Md5()
      .update(signStr)
      .toString()
      .toUpperCase();
    
    console.log('Generated sign:', sign);
    return sign;
  }
}

const pipe = new Pipe();

// 验证 API 密钥
function verifyApiKey(request: Request) {
  const authorization = request.headers.get('Authorization');
    // 检查环境变量是否配置
    if (!VALID_API_KEY) {
      return new Response(JSON.stringify({ error: 'API key not configured' }), {
        status: 500,
        headers: { 'Content-Type': 'application/json' },
      });
    }
    
  if (!authorization) {
    return new Response(JSON.stringify({ error: 'Missing API key' }), {
      status: 401,
      headers: {
        'Content-Type': 'application/json',
        'Access-Control-Allow-Origin': '*',
      },
    });
  }

  const apiKey = authorization.replace('Bearer ', '').trim();
  if (apiKey !== VALID_API_KEY) {
        return new Response(JSON.stringify({ error: 'Invalid API key' }), {
          status: 401,
          headers: {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
          },
        });
      }

  return null;
}

async function handleRequest(request: Request) {
  const url = new URL(request.url);

  // 添加根路径处理
  if (request.method === 'GET' && url.pathname === '/') {
    return new Response("it's work!", {
      headers: {
        'Content-Type': 'text/plain',
        'Access-Control-Allow-Origin': '*',
      },
    });
  }

  if (request.method === 'OPTIONS') {
    return new Response(null, {
      headers: {
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
        'Access-Control-Allow-Headers': 'Content-Type, Authorization',
      },
    });
  }

  // 验证 API 密钥（除了 OPTIONS 请求）
  const authError = verifyApiKey(request);
  if (authError) return authError;

  if (request.method === 'GET' && url.pathname === '/v1/models') {
    const currentTime = Math.floor(Date.now() / 1000);
    return new Response(JSON.stringify({
      object: 'list',
      data: [
        // Original models
        {
          id: 'DeepSeek-R1',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        {
          id: 'DeepSeek-V3',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        {
          id: 'Doubao',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        {
          id: 'Qwen',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        {
          id: 'Glm3',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        {
          id: 'Moonshot_v1',
          object: 'model',
          created: currentTime,
          owned_by: 'library'
        },
        // Search-enabled models
        {
          id: 'DeepSeek-R1-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        },
        {
          id: 'DeepSeek-V3-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        },
        {
          id: 'Doubao-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        },
        {
          id: 'Qwen-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        },
        {
          id: 'Glm3-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        },
        {
          id: 'Moonshot_v1-Search',
          object: 'model',
          created: currentTime,
          owned_by: 'library',
          features: ['online_search']
        }
      ]
    }), {
      headers: {
        'Content-Type': 'application/json',
        'Access-Control-Allow-Origin': '*',
      },
    });
  }

  if (request.method === 'POST' && url.pathname === '/v1/chat/completions') {
    const body = await request.json();
    const isStream = body.stream || false;

    if (isStream) {
      const stream = new ReadableStream({
        async start(controller) {
          try {
            for await (const chunk of pipe.pipe(body)) {
              controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
            }
            controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n'));
            controller.close();
          } catch (e) {
            console.error('Error in stream:', e);
            controller.error(e);
          }
        },
      });

      return new Response(stream, {
        headers: {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          'Connection': 'keep-alive',
          'Access-Control-Allow-Origin': '*',
        },
      });
    }

    if (!isStream) {
      let content = '';
      let meta = null;
      let thinking_content: string[] = [];
      let is_thinking = false;
      
      try {
        for await (const chunk of pipe.pipe(body)) {
          if (chunk.choices?.[0]?.delta?.content) {
            const content_chunk = chunk.choices[0].delta.content;
            if (content_chunk === '<think>\n\n') {
              is_thinking = true;
            } else if (content_chunk === '\n</think>\n\n') {
              is_thinking = false;
            } else if (is_thinking) {
              thinking_content.push(content_chunk);
            } else {
              content += content_chunk;
            }
          }
          if (chunk.choices?.[0]?.delta?.meta) {
            meta = chunk.choices[0].delta.meta;
          }
        }

        // 处理思考内容
        const reasoningContent = thinking_content.join('');
        
        return new Response(JSON.stringify({
          id: crypto.randomUUID(),
          object: 'chat.completion',
          created: Math.floor(Date.now() / 1000),
          model: body.model,
          choices: [{
            message: {
              role: 'assistant',
              reasoning_content: reasoningContent ? `<think>\n${reasoningContent}\n</think>` : '',
              content: content.trim(),
              meta: meta
            },
            finish_reason: 'stop'
          }]
        } as NonStreamResponse), {
          headers: {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
          },
        });
      } catch (e) {
        console.error('Error processing chat request:', e);
        return new Response(JSON.stringify({ error: 'Internal Server Error' }), {
          status: 500,
          headers: {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
          },
        });
      }
    }
  }

  return new Response('Not Found', { status: 404 });
}

serve(handleRequest, { port: 7860 });

interface Message {
  role: string;
  content: string;
}

interface ChatRequest {
  model: string;
  messages: Message[];
  stream: boolean;
  temperature?: number;
  top_p?: number;
  n?: number;
  max_tokens?: number;
  presence_penalty?: number;
  frequency_penalty?: number;
  user?: string;
}

interface DeltaContent {
  content?: string;
  meta?: {
    device_id: string;
    conversation_id: string;
  };
}

interface Choice {
  delta: DeltaContent;
  finish_reason: string | null;
}

interface StreamResponse {
  choices?: Choice[];
  error?: string;
}

interface NonStreamResponse {
  id: string;
  object: string;
  created: number;
  model: string;
  choices: Array<{
    message: {
      role: string;
      reasoning_content: string;
      content: string;
      meta: {
        device_id: string;
        conversation_id: string;
      };
    };
    finish_reason: string;
  }>;
}