src/lightweight-client-express.js

import express from 'express';
import dotenv from 'dotenv';
import { randomUUID } from 'crypto';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import chalk from 'chalk';
import {
  ChatMessage, ChatCompletionRequest, Choice, ChoiceDelta, ChatCompletionChunk
} from './models.js';
import { cookieManager } from './CookieManager.js';
import { PassThrough } from 'stream';
import fetch from 'node-fetch';

// 获取当前文件的目录路径
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

// 加载环境变量
dotenv.config({ path: join(dirname(__dirname), '.env') });

// 初始化状态变量
let INITIALIZED_SUCCESSFULLY = false;

// 活跃流管理器 - 用于跟踪和管理当前活跃的请求流
const activeStreams = new Map();

// 日志配置
const logger = {
  info: (message) => console.log(chalk.blue(`[info] ${message}`)),
  error: (message) => console.error(chalk.red(`[error] ${message}`)),
  warning: (message) => console.warn(chalk.yellow(`[warn] ${message}`)),
  success: (message) => console.log(chalk.green(`[success] ${message}`)),
  request: (method, path, status, time) => {
    const statusColor = status >= 500 ? chalk.red : 
                        status >= 400 ? chalk.yellow : 
                        status >= 300 ? chalk.cyan : 
                        status >= 200 ? chalk.green : chalk.white;
    console.log(`${chalk.magenta(`[${method}]`)} - ${path} ${statusColor(status)} ${chalk.gray(`${time}ms`)}`);
  }
};

// 认证配置
const EXPECTED_TOKEN = process.env.PROXY_AUTH_TOKEN || "default_token";

// 代理配置
const PROXY_URL = process.env.PROXY_URL;

async function validateProxy() {
  if (!PROXY_URL) {
    logger.info('未配置代理，将使用直连访问');
    return true;
  }
  
  try {
    const { HttpsProxyAgent } = await import('https-proxy-agent');
    
    const testResponse = await fetch('https://httpbin.org/ip', {
      method: 'GET',
      agent: new HttpsProxyAgent(PROXY_URL),
      timeout: 15000
    });
    
    if (testResponse.ok) {
      const ipInfo = await testResponse.json();
      const safeProxyUrl = PROXY_URL.replace(/:([^:@]+)@/, ':****@');
      logger.success(`✅ 代理验证成功 - IP: ${ipInfo.origin} - 代理: ${safeProxyUrl}`);
      return true;
    } else {
      logger.error(`❌ 代理验证失败 - 状态码: ${testResponse.status}`);
      return false;
    }
  } catch (error) {
    logger.error(`❌ 代理验证异常: ${error.message}`);
    return false;
  }
}

// 初始化函数
async function initialize() {
  logger.info('开始系统初始化...');
  
  try {
    const envCookie = process.env.NOTION_COOKIE;
    const envCookieFile = process.env.COOKIE_FILE;
    
    if (envCookie) {
      logger.info('发现环境变量中的NOTION_COOKIE，正在初始化...');
      const success = await cookieManager.initialize(envCookie);
      if (success) {
        logger.success('Cookie初始化成功');
        return true;
      } else {
        logger.error('Cookie初始化失败');
        return false;
      }
    } else if (envCookieFile) {
      logger.info(`发现环境变量中的COOKIE_FILE: ${envCookieFile}，正在加载...`);
      const success = await cookieManager.loadFromFile(envCookieFile);
      if (success) {
        logger.success('从文件加载Cookie成功');
        return true;
      } else {
        logger.error('从文件加载Cookie失败');
        return false;
      }
    } else {
      logger.warning('未找到NOTION_COOKIE或COOKIE_FILE环境变量');
      return false;
    }
  } catch (error) {
    logger.error(`初始化过程出错: ${error.message}`);
    return false;
  }
}

// 流式响应函数
async function streamNotionResponse(notionRequestBody) {
  const stream = new PassThrough();
  let streamClosed = false;
  
  const originalEnd = stream.end;
  stream.end = function(...args) {
    if (streamClosed) return;
    streamClosed = true;
    return originalEnd.apply(this, args);
  };
  
  stream.write(':\n\n');
  
  const timeoutId = setTimeout(() => {
    if (streamClosed) return;
    
    logger.warning(`请求超时，30秒内未收到响应`);
    try {
      const endChunk = new ChatCompletionChunk({
        choices: [
          new Choice({
            delta: new ChoiceDelta({ content: "请求超时，未收到Notion响应。" }),
            finish_reason: "timeout"
          })
        ]
      });
      stream.write(`data: ${JSON.stringify(endChunk)}\n\n`);
      stream.write('data: [DONE]\n\n');
      stream.end();
    } catch (error) {
      logger.error(`发送超时消息时出错: ${error}`);
      if (!streamClosed) stream.end();
    }
  }, 30000);
  
  fetchNotionResponse(stream, notionRequestBody, timeoutId).catch((error) => {
    if (streamClosed) return;
    
    logger.error(`流处理出错: ${error}`);
    clearTimeout(timeoutId);
    
    try {
      const errorChunk = new ChatCompletionChunk({
        choices: [
          new Choice({
            delta: new ChoiceDelta({ content: `处理请求时出错: ${error.message}` }),
            finish_reason: "error"
          })
        ]
      });
      stream.write(`data: ${JSON.stringify(errorChunk)}\n\n`);
      stream.write('data: [DONE]\n\n');
    } catch (e) {
      logger.error(`发送错误消息时出错: ${e}`);
    } finally {
      if (!streamClosed) stream.end();
    }
  });
  
  return stream;
}

// 实际请求Notion API的函数
async function fetchNotionResponse(chunkQueue, notionRequestBody, timeoutId) {
  let responseReceived = false;
  
  const isStreamClosed = () => {
    return chunkQueue.destroyed || (typeof chunkQueue.closed === 'boolean' && chunkQueue.closed);
  };
  
  const safeWrite = (data) => {
    if (!isStreamClosed()) {
      try {
        return chunkQueue.write(data);
      } catch (error) {
        logger.error(`流写入错误: ${error.message}`);
        return false;
      }
    }
    return false;
  };
  
  try {
    const cookieData = cookieManager.getNext();
    if (!cookieData) {
      throw new Error('没有可用的cookie');
    }
    
    const headers = {
      'Content-Type': 'application/json',
      'accept': 'application/x-ndjson',
      'accept-language': 'en-US,en;q=0.9',
      'notion-audit-log-platform': 'web',
      'notion-client-version': '23.13.0.3686',
      'origin': 'https://www.notion.so',
      'referer': 'https://www.notion.so/chat',
      'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36',
      'x-notion-active-user-header': cookieData.userId,
      'x-notion-space-id': cookieData.spaceId,
      'Cookie': cookieData.cookie
    };
    
    const fetchOptions = {
      method: 'POST',
      headers: headers,
      body: JSON.stringify(notionRequestBody),
    };
    
    // 添加代理配置
    if (PROXY_URL) {
      const { HttpsProxyAgent } = await import('https-proxy-agent');
      fetchOptions.agent = new HttpsProxyAgent(PROXY_URL);
      logger.info(`使用固定代理连接Notion API`);
    }
    
    const response = await fetch('https://www.notion.so/api/v3/runInferenceTranscript', fetchOptions);
    
    if (response.status === 401) {
      logger.error(`收到401未授权错误，cookie可能已失效`);
      cookieManager.markAsInvalid(cookieData.userId);
      throw new Error('Cookie已失效');
    }
    
    if (!response.ok) {
      throw new Error(`HTTP error! status: ${response.status}`);
    }
    
    if (!response.body) {
      throw new Error("Response body is null");
    }
    
    const reader = response.body;
    let buffer = '';
    
    reader.on('data', (chunk) => {
      if (isStreamClosed()) {
        try {
          reader.destroy();
        } catch (error) {
          logger.error(`销毁reader时出错: ${error.message}`);
        }
        return;
      }
      
      try {
        if (!responseReceived) {
          responseReceived = true;
          logger.info(`已连接Notion API`);
          clearTimeout(timeoutId);
        }
        
        const text = chunk.toString('utf8');
        buffer += text;
        
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';
        
        for (const line of lines) {
          if (!line.trim()) continue;
          
          try {
            const jsonData = JSON.parse(line);
            
            if (jsonData?.type === "markdown-chat" && typeof jsonData?.value === "string") {
              const content = jsonData.value;
              if (!content) continue;
              
              const chunk = new ChatCompletionChunk({
                choices: [
                  new Choice({
                    delta: new ChoiceDelta({ content }),
                    finish_reason: null
                  })
                ]
              });
              
              const dataStr = `data: ${JSON.stringify(chunk)}\n\n`;
              if (!safeWrite(dataStr)) {
                try {
                  reader.destroy();
                } catch (error) {
                  logger.error(`写入失败后销毁reader时出错: ${error.message}`);
                }
                return;
              }
            }
          } catch (parseError) {
            logger.error(`解析JSON失败: ${parseError.message}`);
          }
        }
      } catch (error) {
        logger.error(`处理数据块时出错: ${error.message}`);
      }
    });
    
    reader.on('end', () => {
      logger.info('Notion API响应流结束');
      clearTimeout(timeoutId);
      
      try {
        const endChunk = new ChatCompletionChunk({
          choices: [
            new Choice({
              delta: new ChoiceDelta({ content: "" }),
              finish_reason: "stop"
            })
          ]
        });
        safeWrite(`data: ${JSON.stringify(endChunk)}\n\n`);
        safeWrite('data: [DONE]\n\n');
      } catch (error) {
        logger.error(`发送结束标记时出错: ${error.message}`);
      }
      
      if (!isStreamClosed()) {
        chunkQueue.end();
      }
    });
    
    reader.on('error', (error) => {
      logger.error(`读取响应流时出错: ${error.message}`);
      clearTimeout(timeoutId);
      if (!isStreamClosed()) {
        chunkQueue.end();
      }
    });
    
  } catch (error) {
    logger.error(`请求Notion API失败: ${error.message}`);
    clearTimeout(timeoutId);
    if (!isStreamClosed()) {
      chunkQueue.end();
    }
    throw error;
  }
}

// 构建Notion请求的函数
function buildNotionRequest(requestData) {
  const cookieData = cookieManager.getNext();
  if (!cookieData) {
    throw new Error('没有可用的cookie');
  }

  const now = new Date().toISOString();
  const transcript = [];
  
  // 添加配置
  if (requestData.model === 'google-gemini-2.5-pro') {
    transcript.push({
      id: randomUUID(),
      type: "config",
      value: { model: 'vertex-gemini-2.5-pro' }
    });
  } else if (requestData.model === 'google-gemini-2.5-flash') {
    transcript.push({
      id: randomUUID(),
      type: "config",
      value: { model: 'vertex-gemini-2.5-flash' }
    });
  } else {
    transcript.push({
      id: randomUUID(),
      type: "config",
      value: { model: requestData.model }
    });
  }
  
  // 添加上下文
  transcript.push({
    id: randomUUID(),
    type: "context",
    value: {
      userId: cookieData.userId,
      spaceId: cookieData.spaceId,
      surface: "home_module",
      timezone: "America/Los_Angeles",
      userName: `User${Math.floor(Math.random() * 900) + 100}`,
      spaceName: `Project ${Math.floor(Math.random() * 99) + 1}`,
      spaceViewId: randomUUID(),
      currentDatetime: now
    }
  });
  
  // 添加消息
  for (const message of requestData.messages) {
    let content = message.content;
    if (Array.isArray(content)) {
      content = content.map(part => part.type === 'text' ? part.text : '').join('');
    }
    
    if (message.role === "user" || message.role === "system") {
      transcript.push({
        id: randomUUID(),
        type: "user",
        value: [[content]],
        userId: cookieData.userId,
        createdAt: now
      });
    } else if (message.role === "assistant") {
      transcript.push({
        id: randomUUID(),
        type: "markdown-chat",
        value: content,
        traceId: randomUUID(),
        createdAt: now
      });
    }
  }
  
  return {
    spaceId: cookieData.spaceId,
    transcript: transcript,
    createThread: true,
    traceId: randomUUID(),
    debugOverrides: {
      cachedInferences: {},
      annotationInferences: {},
      emitInferences: false
    },
    generateTitle: false,
    saveAllThreadOperations: false
  };
}

// 创建Express应用
const app = express();
app.use(express.json({ limit: '50mb' }));
app.use(express.urlencoded({ extended: true, limit: '50mb' }));

// 请求日志中间件
app.use((req, res, next) => {
  const start = Date.now();
  
  // 保存原始的 end 方法
  const originalEnd = res.end;
  
  // 重写 end 方法以记录请求完成时间
  res.end = function(...args) {
    const duration = Date.now() - start;
    logger.request(req.method, req.path, res.statusCode, duration);
    return originalEnd.apply(this, args);
  };
  
  next();
});

// 认证中间件
function authenticate(req, res, next) {
  const authHeader = req.headers.authorization;
  
  if (!authHeader || !authHeader.startsWith('Bearer ')) {
    return res.status(401).json({
      error: {
        message: "Authentication required. Please provide a valid Bearer token.",
        type: "authentication_error"
      }
    });
  }
  
  const token = authHeader.split(' ')[1];
  
  if (token !== EXPECTED_TOKEN) {
    return res.status(401).json({
      error: {
        message: "Invalid authentication credentials",
        type: "authentication_error"
      }
    });
  }
  
  next();
}

// 流管理函数
function manageStream(clientId, stream) {
  // 如果该客户端已有活跃流，先关闭它
  if (activeStreams.has(clientId)) {
    try {
      const oldStream = activeStreams.get(clientId);
      logger.info(`关闭客户端 ${clientId} 的旧流`);
      oldStream.end();
    } catch (error) {
      logger.error(`关闭旧流时出错: ${error.message}`);
    }
  }
  
  // 注册新流
  activeStreams.set(clientId, stream);
  
  // 当流结束时从管理器中移除
  stream.on('end', () => {
    if (activeStreams.get(clientId) === stream) {
      activeStreams.delete(clientId);
      //logger.info(`客户端 ${clientId} 的流已结束并移除`);
    }
  });
  
  stream.on('error', (error) => {
    logger.error(`流错误: ${error.message}`);
    if (activeStreams.get(clientId) === stream) {
      activeStreams.delete(clientId);
    }
  });
  
  return stream;
}

// API路由

// 获取模型列表
app.get('/v1/models', authenticate, (req, res) => {
  // 返回可用模型列表
  const modelList = {
    data: [
      { id: "openai-gpt-4.1" },
      { id: "anthropic-opus-4" },
      { id: "anthropic-sonnet-4" },
      { id: "anthropic-sonnet-3.x-stable" },
      { id: "google-gemini-2.5-pro"}, //vertex-gemini-2.5-pro
      { id: "google-gemini-2.5-flash"}, //vertex-gemini-2.5-flash
    ]
  };
  
  res.json(modelList);
});

// 聊天完成端点
app.post('/v1/chat/completions', authenticate, async (req, res) => {
  try {
    // 生成或获取客户端ID
    const clientId = req.headers['x-client-id'] || randomUUID();
    
    // 检查是否成功初始化
    if (!INITIALIZED_SUCCESSFULLY) {
      return res.status(500).json({
        error: {
          message: "系统未成功初始化。请检查您的NOTION_COOKIE是否有效。",
          type: "server_error"
        }
      });
    }
    
    // 检查是否有可用的cookie
    if (cookieManager.getValidCount() === 0) {
      return res.status(500).json({
        error: {
          message: "没有可用的有效cookie。请检查您的NOTION_COOKIE配置。",
          type: "server_error"
        }
      });
    }
    
    // 验证请求数据
    const requestData = req.body;
    
    if (!requestData.messages || !Array.isArray(requestData.messages) || requestData.messages.length === 0) {
      return res.status(400).json({
        error: {
          message: "Invalid request: 'messages' field must be a non-empty array.",
          type: "invalid_request_error"
        }
      });
    }
    
    // 构建Notion请求
    const notionRequestBody = buildNotionRequest(requestData);
    
    // 处理流式响应
    if (requestData.stream) {
      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');
      
      logger.info(`开始流式响应`);
      const stream = await streamNotionResponse(notionRequestBody);
      
      // 管理流
      manageStream(clientId, stream);
      
      stream.pipe(res);
      
      // 处理客户端断开连接
      req.on('close', () => {
        logger.info(`客户端 ${clientId} 断开连接`);
        if (activeStreams.has(clientId)) {
          try {
            activeStreams.get(clientId).end();
            activeStreams.delete(clientId);
          } catch (error) {
            logger.error(`关闭流时出错: ${error.message}`);
          }
        }
      });
    } else {
      // 非流式响应
      // 创建一个内部流来收集完整响应
      logger.info(`开始非流式响应`);
      const chunks = [];
      const stream = await streamNotionResponse(notionRequestBody);
      
      // 管理流
      manageStream(clientId, stream);
      
      return new Promise((resolve, reject) => {
        stream.on('data', (chunk) => {
          const chunkStr = chunk.toString();
          if (chunkStr.startsWith('data: ') && !chunkStr.includes('[DONE]')) {
            try {
              const dataJson = chunkStr.substring(6).trim();
              if (dataJson) {
                const chunkData = JSON.parse(dataJson);
                if (chunkData.choices && chunkData.choices[0].delta && chunkData.choices[0].delta.content) {
                  chunks.push(chunkData.choices[0].delta.content);
                }
              }
            } catch (error) {
              logger.error(`解析非流式响应块时出错: ${error}`);
            }
          }
        });
        
        stream.on('end', () => {
          const fullResponse = {
            id: `chatcmpl-${randomUUID()}`,
            object: "chat.completion",
            created: Math.floor(Date.now() / 1000),
            model: requestData.model,
            choices: [
              {
                index: 0,
                message: {
                  role: "assistant",
                  content: chunks.join('')
                },
                finish_reason: "stop"
              }
            ],
            usage: {
              prompt_tokens: null,
              completion_tokens: null,
              total_tokens: null
            }
          };
          
          res.json(fullResponse);
          resolve();
        });
        
        stream.on('error', (error) => {
          logger.error(`非流式响应出错: ${error}`);
          reject(error);
        });
        
        // 处理客户端断开连接
        req.on('close', () => {
          logger.info(`客户端 ${clientId} 断开连接（非流式）`);
          if (activeStreams.has(clientId)) {
            try {
              activeStreams.get(clientId).end();
              activeStreams.delete(clientId);
            } catch (error) {
              logger.error(`关闭流时出错: ${error.message}`);
            }
          }
        });
      });
    }
  } catch (error) {
    logger.error(`聊天完成端点错误: ${error}`);
    res.status(500).json({
      error: {
        message: `Internal server error: ${error.message}`,
        type: "server_error"
      }
    });
  }
});

// 健康检查端点
app.get('/health', (req, res) => {
  res.json({
    status: 'ok',
    timestamp: new Date().toISOString(),
    initialized: INITIALIZED_SUCCESSFULLY,
    valid_cookies: cookieManager.getValidCount(),
    active_streams: activeStreams.size
  });
});

// Cookie状态查询端点
app.get('/cookies/status', authenticate, (req, res) => {
  res.json({
    total_cookies: cookieManager.getValidCount(),
    cookies: cookieManager.getStatus()
  });
});

// 启动服务器
const PORT = process.env.PORT || 7860;

// 初始化并启动服务器
initialize().then(async (initResult) => {
  // 设置初始化状态
  INITIALIZED_SUCCESSFULLY = initResult !== false;
  
  // 验证代理配置
  if (PROXY_URL) {
    try {
      await validateProxy();
    } catch (error) {
      logger.warning(`代理验证失败: ${error.message}`);
    }
  }
  
  // 启动服务器
  app.listen(PORT, '0.0.0.0', () => {
    logger.info(`服务已在 0.0.0.0:${PORT} 上启动`);
    logger.info(`访问地址: http://0.0.0.0:${PORT}`);
    
    // 显示代理状态
    if (PROXY_URL) {
      const safeProxyUrl = PROXY_URL.replace(/:([^:@]+)@/, ':****@');
      logger.info(`代理配置: ${safeProxyUrl}`);
    } else {
      logger.info(`代理配置: 未配置代理，使用直连`);
    }
    
    if (INITIALIZED_SUCCESSFULLY) {
      logger.success(`系统初始化状态: ✅`);
      logger.success(`可用cookie数量: ${cookieManager.getValidCount()}`);
    } else {
      logger.warning(`系统初始化状态: ❌`);
      logger.warning(`警告: 系统未成功初始化，API调用将无法正常工作`);
      logger.warning(`请检查NOTION_COOKIE配置是否有效`);
    }
  });
}).catch((error) => {
  logger.error(`初始化失败: ${error}`);
  INITIALIZED_SUCCESSFULLY = false;
  
  // 即使初始化失败也启动服务器
  app.listen(PORT, '0.0.0.0', () => {
    logger.warning(`服务已在 0.0.0.0:${PORT} 上启动（初始化失败模式）`);
  });
});