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 {
  initialize,
  buildNotionRequest,
  streamNotionResponse,
  INITIALIZED_SUCCESSFULLY
} from './lightweight-client.js';
import { cookieManager } from './CookieManager.js';

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

dotenv.config({ path: join(dirname(__dirname), '.env') });

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 app = express();
app.use(express.json({ limit: '50mb' }));
app.use(express.urlencoded({ extended: true, limit: '50mb' }));

// 请求计数器和内存监控
let requestCount = 0;
let lastResetTime = Date.now();
const startTime = Date.now();
const MAX_UPTIME = 2 * 60 * 60 * 1000; // 2小时
const MAX_MEMORY = 400 * 1024 * 1024; // 400MB

// 请求日志中间件
app.use((req, res, next) => {
  const start = Date.now();
  
  const originalEnd = res.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();
});

// 请求频率限制中间件
app.use('/v1/chat/completions', (req, res, next) => {
  const now = Date.now();
  
  // 每分钟重置计数
  if (now - lastResetTime > 60000) {
    requestCount = 0;
    lastResetTime = now;
  }
  
  requestCount++;
  
  // 限制每分钟最多15个请求
  if (requestCount > 15) {
    logger.warning(`请求频率过高，当前: ${requestCount}/分钟`);
    return res.status(429).json({
      error: {
        message: "Too many requests, please try again later",
        type: "rate_limit_error"
      }
    });
  }
  
  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();
}

// 根路径
app.get('/', (req, res) => {
  const memUsage = process.memoryUsage();
  const uptime = Date.now() - startTime;
  
  res.json({
    message: "Notion2API NodeJS Service",
    version: "1.0.0",
    status: "running",
    uptime: Math.floor(uptime / 1000),
    memory: {
      used: Math.round(memUsage.heapUsed / 1024 / 1024),
      total: Math.round(memUsage.heapTotal / 1024 / 1024)
    },
    endpoints: {
      health: "/health",
      models: "/v1/models",
      chat: "/v1/chat/completions"
    }
  });
});

// 获取模型列表 - 添加了所有原始项目中的模型
app.get('/v1/models', authenticate, (req, res) => {
  const modelList = {
    object: "list",
    data: [
      {
        id: "openai-gpt-4.1",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      },
      {
        id: "anthropic-opus-4",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      },
      {
        id: "anthropic-sonnet-4",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      },
      {
        id: "anthropic-sonnet-3.x-stable",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      },
      {
        id: "google-gemini-2.5-pro",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      },
      {
        id: "google-gemini-2.5-flash",
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: "notion"
      }
    ]
  };
  
  res.json(modelList);
});

// 聊天完成端点
app.post('/v1/chat/completions', authenticate, async (req, res) => {
  try {
    if (!INITIALIZED_SUCCESSFULLY) {
      return res.status(500).json({
        error: {
          message: "系统未成功初始化。请检查您的NOTION_COOKIE是否有效。",
          type: "server_error"
        }
      });
    }
    
    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"
        }
      });
    }
    
    if (!requestData.model) {
      requestData.model = "anthropic-sonnet-4";
    }
    
    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');
      res.setHeader('Access-Control-Allow-Origin', '*');
      res.setHeader('Access-Control-Allow-Headers', 'Cache-Control');
      
      logger.info(`开始流式响应`);
      const stream = await streamNotionResponse(notionRequestBody);
      stream.pipe(res);
      
      req.on('close', () => {
        logger.info('客户端断开连接，结束流');
        stream.destroy();
      });
      
      req.on('aborted', () => {
        logger.info('请求被中止，结束流');
        stream.destroy();
      });
    } else {
      logger.info(`开始非流式响应`);
      const chunks = [];
      const stream = await streamNotionResponse(notionRequestBody);
      
      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}`);
          if (!res.headersSent) {
            res.status(500).json({
              error: {
                message: `Stream error: ${error.message}`,
                type: "server_error"
              }
            });
          }
          reject(error);
        });
      });
    }
  } catch (error) {
    logger.error(`聊天完成端点错误: ${error}`);
    if (!res.headersSent) {
      res.status(500).json({
        error: {
          message: `Internal server error: ${error.message}`,
          type: "server_error"
        }
      });
    }
  }
});

// 增强的健康检查端点
app.get('/health', (req, res) => {
  const memUsage = process.memoryUsage();
  const memUsageMB = {
    rss: Math.round(memUsage.rss / 1024 / 1024),
    heapUsed: Math.round(memUsage.heapUsed / 1024 / 1024),
    heapTotal: Math.round(memUsage.heapTotal / 1024 / 1024),
    external: Math.round(memUsage.external / 1024 / 1024)
  };
  
  const uptime = Date.now() - startTime;
  const status = memUsage.heapUsed > MAX_MEMORY ? 'warning' : 'ok';
  
  res.json({
    status: status,
    timestamp: new Date().toISOString(),
    initialized: INITIALIZED_SUCCESSFULLY,
    valid_cookies: cookieManager.getValidCount(),
    uptime: Math.floor(uptime / 1000),
    memory: memUsageMB,
    requests_per_minute: requestCount,
    version: "1.0.0"
  });
});

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

// 内存监控和自动重启
setInterval(() => {
  const uptime = Date.now() - startTime;
  const memUsage = process.memoryUsage();
  
  // 记录内存使用情况
  if (memUsage.heapUsed > 300 * 1024 * 1024) { // 300MB警告
    logger.warning(`内存使用较高: ${Math.round(memUsage.heapUsed / 1024 / 1024)}MB`);
  }
  
  // 检查是否需要重启
  if (uptime > MAX_UPTIME || memUsage.heapUsed > MAX_MEMORY) {
    logger.warning('达到重启条件:');
    logger.warning(`运行时间: ${Math.floor(uptime / 1000 / 60)}分钟 (最大: ${Math.floor(MAX_UPTIME / 1000 / 60)}分钟)`);
    logger.warning(`内存使用: ${Math.round(memUsage.heapUsed / 1024 / 1024)}MB (最大: ${Math.round(MAX_MEMORY / 1024 / 1024)}MB)`);
    logger.warning('正在重启应用以释放内存...');
    
    // 优雅关闭
    setTimeout(() => {
      process.exit(0); // Hugging Face 会自动重启
    }, 1000);
  }
}, 5 * 60 * 1000); // 每5分钟检查一次

// 错误处理中间件
app.use((error, req, res, next) => {
  logger.error(`未处理的错误: ${error.message}`);
  if (!res.headersSent) {
    res.status(500).json({
      error: {
        message: "Internal server error",
        type: "server_error"
      }
    });
  }
});

// 404处理
app.use((req, res) => {
  res.status(404).json({
    error: {
      message: `Not found: ${req.method} ${req.path}`,
      type: "not_found_error"
    }
  });
});

const PORT = process.env.PORT || 7860;

// 初始化并启动服务器
initialize().then(() => {
  app.listen(PORT, '0.0.0.0', () => {
    logger.info(`服务已启动 - 端口: ${PORT}`);
    logger.info(`访问地址: http://0.0.0.0:${PORT}`);
    logger.info(`环境: ${process.env.NODE_ENV || 'development'}`);
    
    if (INITIALIZED_SUCCESSFULLY) {
      logger.success(`系统初始化状态: ✅`);
      logger.success(`可用cookie数量: ${cookieManager.getValidCount()}`);
    } else {
      logger.warning(`系统初始化状态: ❌`);
      logger.warning(`警告: 系统未成功初始化，API调用将无法正常工作`);
      logger.warning(`请检查COOKIE_FILE_CONTENT或NOTION_COOKIE配置是否有效`);
    }
  });
}).catch((error) => {
  logger.error(`初始化失败: ${error}`);
  process.exit(1);
});

// 优雅关闭处理
process.on('SIGTERM', () => {
  logger.info('收到SIGTERM信号，正在优雅关闭...');
  process.exit(0);
});

process.on('SIGINT', () => {
  logger.info('收到SIGINT信号，正在优雅关闭...');
  process.exit(0);
});