server.js

const express = require('express');
const puppeteer = require('puppeteer');
const cors = require('cors');
const helmet = require('helmet');
const rateLimit = require('express-rate-limit');

const app = express();
const PORT = process.env.PORT || 7860;

// 中间件配置 - HF Spaces 优化
app.use(helmet({
  contentSecurityPolicy: false // HF Spaces 需要
}));
app.use(cors());
app.use(express.json({ limit: '10mb' }));

// 速率限制 - HF Spaces 调整
const limiter = rateLimit({
  windowMs: 15 * 60 * 1000,
  max: 30, // 进一步降低限制
  message: {
    error: 'Too many requests, please try again later.'
  }
});
app.use('/screenshot', limiter);

// 健康检查端点
app.get('/', (req, res) => {
  res.json({
    message: 'Page Screenshot API - Hugging Face Spaces',
    version: '1.0.0',
    status: 'running',
    platform: 'HuggingFace Spaces',
    endpoints: {
      screenshot: 'POST /screenshot',
      demo: 'GET /demo',
      health: 'GET /'
    }
  });
});

// 截图API端点 - 增强错误处理
app.post('/screenshot', async (req, res) => {
  const { url, width = 1280, height = 720, quality = 75 } = req.body;

  // 参数验证
  if (!url) {
    return res.status(400).json({
      error: 'URL is required',
      example: { url: 'https://example.com', width: 1280, height: 720 }
    });
  }

  // URL格式验证
  try {
    const urlObj = new URL(url);
    // 检查协议
    if (!['http:', 'https:'].includes(urlObj.protocol)) {
      return res.status(400).json({
        error: 'Only HTTP and HTTPS URLs are supported'
      });
    }
  } catch (error) {
    return res.status(400).json({
      error: 'Invalid URL format'
    });
  }

  // 分辨率验证 - HF Spaces 更严格限制
  if (width < 100 || width > 1600 || height < 100 || height > 1200) {
    return res.status(400).json({
      error: 'Width must be 100-1600px, height must be 100-1200px for HF Spaces'
    });
  }

  let browser;
  try {
    // 启动浏览器 - HF Spaces 专用配置
    const browserOptions = {
      headless: 'new',
      args: [
        '--no-sandbox',
        '--disable-setuid-sandbox',
        '--disable-dev-shm-usage',
        '--disable-gpu',
        '--no-first-run',
        '--no-zygote',
        '--single-process',
        '--disable-extensions',
        '--disable-background-timer-throttling',
        '--disable-backgrounding-occluded-windows',
        '--disable-renderer-backgrounding',
        '--disable-features=TranslateUI',
        '--disable-default-apps',
        '--no-default-browser-check',
        '--disable-background-networking'
      ]
    };

    // 在 HF Spaces 中使用系统 Chrome
    if (process.env.PUPPETEER_EXECUTABLE_PATH) {
      browserOptions.executablePath = process.env.PUPPETEER_EXECUTABLE_PATH;
    }

    console.log('Launching browser...');
    browser = await puppeteer.launch(browserOptions);

    const page = await browser.newPage();
    
    // 设置视窗大小
    await page.setViewport({
      width: parseInt(width),
      height: parseInt(height),
      deviceScaleFactor: 1
    });

    // 设置用户代理和其他页面选项
    await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36');
    
    // 拦截不必要的资源以提高性能
    await page.setRequestInterception(true);
    page.on('request', (req) => {
      const resourceType = req.resourceType();
      if (['font', 'media'].includes(resourceType)) {
        req.abort();
      } else {
        req.continue();
      }
    });

    console.log(`Navigating to: ${url}`);
    
    // 访问页面 - HF Spaces 更短超时
    await page.goto(url, {
      waitUntil: 'domcontentloaded', // 更快的等待条件
      timeout: 15000 // 15秒超时
    });

    // 等待页面稳定
    await page.waitForTimeout(1000);

    console.log('Taking screenshot...');
    
    // 截图
    const screenshot = await page.screenshot({
      type: 'jpeg',
      quality: Math.max(10, Math.min(100, parseInt(quality))),
      fullPage: false
    });

    console.log(`Screenshot taken: ${screenshot.length} bytes`);

    // 设置响应头
    res.set({
      'Content-Type': 'image/jpeg',
      'Content-Length': screenshot.length,
      'Cache-Control': 'no-cache',
      'Content-Disposition': `inline; filename="screenshot-${Date.now()}.jpg"`
    });

    res.send(screenshot);

  } catch (error) {
    console.error('Screenshot error:', error.message);
    const errorResponse = {
      error: 'Failed to capture screenshot',
      message: error.message
    };
    
    // 根据错误类型提供更好的错误信息
    if (error.message.includes('timeout')) {
      errorResponse.suggestion = 'Try a simpler webpage or reduce timeout';
    } else if (error.message.includes('net::')) {
      errorResponse.suggestion = 'Check if the URL is accessible';
    }
    
    res.status(500).json(errorResponse);
  } finally {
    if (browser) {
      try {
        await browser.close();
        console.log('Browser closed');
      } catch (closeError) {
        console.error('Error closing browser:', closeError.message);
      }
    }
  }
});

// HF Spaces 演示界面 - 改进版
app.get('/demo', (req, res) => {
  res.send(`
    <!DOCTYPE html>
    <html>
    <head>
        <title>📸 Page Screenshot API Demo</title>
        <meta charset="utf-8">
        <meta name="viewport" content="width=device-width, initial-scale=1">
        <style>
            body { 
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', system-ui, sans-serif; 
                max-width: 800px; margin: 0 auto; padding: 20px; 
                background: #f8f9fa;
            }
            .container { background: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }
            .form-group { margin: 20px 0; }
            label { display: block; margin-bottom: 8px; font-weight: 600; color: #333; }
            input[type="text"], input[type="number"] { 
                width: 100%; padding: 12px; border: 2px solid #e1e5e9; 
                border-radius: 6px; font-size: 16px; box-sizing: border-box;
            }
            input:focus { border-color: #007bff; outline: none; }
            .input-row { display: flex; gap: 15px; }
            .input-row > div { flex: 1; }
            button { 
                background: linear-gradient(135deg, #007bff, #0056b3); 
                color: white; border: none; padding: 14px 28px; 
                border-radius: 6px; cursor: pointer; font-size: 16px; font-weight: 600;
                transition: transform 0.2s;
            }
            button:hover { transform: translateY(-1px); }
            button:disabled { background: #6c757d; cursor: not-allowed; transform: none; }
            #result { margin-top: 30px; }
            .loading { color: #007bff; font-weight: 500; }
            .error { color: #dc3545; background: #f8d7da; padding: 15px; border-radius: 6px; }
            .success img { max-width: 100%; border-radius: 6px; box-shadow: 0 4px 15px rgba(0,0,0,0.1); }
            .examples { margin: 20px 0; }
            .example-btn { 
                background: #e9ecef; color: #495057; border: none; 
                padding: 8px 12px; margin: 5px; border-radius: 4px; cursor: pointer; font-size: 14px;
            }
            .example-btn:hover { background: #dee2e6; }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>📸 Page Screenshot API</h1>
            <p>Enter a URL to capture a screenshot. Optimized for Hugging Face Spaces.</p>
            
            <div class="examples">
                <strong>Try these examples:</strong><br>
                <button class="example-btn" onclick="setExample('https://www.google.com')">Google</button>
                <button class="example-btn" onclick="setExample('https://www.github.com')">GitHub</button>
                <button class="example-btn" onclick="setExample('https://www.wikipedia.org')">Wikipedia</button>
                <button class="example-btn" onclick="setExample('https://news.ycombinator.com')">Hacker News</button>
            </div>
            
            <div class="form-group">
                <label for="url">URL:</label>
                <input type="text" id="url" placeholder="https://example.com" value="https://www.google.com">
            </div>
            
            <div class="input-row">
                <div>
                    <label for="width">Width (px):</label>
                    <input type="number" id="width" value="1280" min="100" max="1600">
                </div>
                <div>
                    <label for="height">Height (px):</label>
                    <input type="number" id="height" value="720" min="100" max="1200">
                </div>
                <div>
                    <label for="quality">Quality:</label>
                    <input type="number" id="quality" value="75" min="10" max="100">
                </div>
            </div>
            
            <button onclick="takeScreenshot()" id="captureBtn">Take Screenshot</button>
            
            <div id="result"></div>
        </div>
        
        <script>
            function setExample(url) {
                document.getElementById('url').value = url;
            }
            
            async function takeScreenshot() {
                const url = document.getElementById('url').value;
                const width = parseInt(document.getElementById('width').value);
                const height = parseInt(document.getElementById('height').value);
                const quality = parseInt(document.getElementById('quality').value);
                const btn = document.getElementById('captureBtn');
                
                if (!url) {
                    document.getElementById('result').innerHTML = '<div class="error">Please enter a URL</div>';
                    return;
                }
                
                btn.disabled = true;
                btn.textContent = 'Taking Screenshot...';
                document.getElementById('result').innerHTML = '<div class="loading">📸 Capturing screenshot, please wait...</div>';
                
                try {
                    const response = await fetch('/screenshot', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify({ url, width, height, quality })
                    });
                    
                    if (response.ok) {
                        const blob = await response.blob();
                        const imageUrl = URL.createObjectURL(blob);
                        const size = (blob.size / 1024).toFixed(1);
                        document.getElementById('result').innerHTML = 
                            '<div class="success"><h3>Screenshot Result:</h3>' +
                            '<p>Size: ' + size + ' KB | Dimensions: ' + width + 'x' + height + '</p>' +
                            '<img src="' + imageUrl + '" alt="Screenshot"><br><br>' +
                            '<a href="' + imageUrl + '" download="screenshot.jpg" style="background: #28a745; color: white; padding: 10px 20px; text-decoration: none; border-radius: 5px;">Download Image</a></div>';
                    } else {
                        const error = await response.json();
                        document.getElementById('result').innerHTML = 
                            '<div class="error"><strong>Error:</strong> ' + error.error + 
                            (error.suggestion ? '<br><strong>Suggestion:</strong> ' + error.suggestion : '') + '</div>';
                    }
                } catch (error) {
                    document.getElementById('result').innerHTML = 
                        '<div class="error"><strong>Network Error:</strong> ' + error.message + '</div>';
                } finally {
                    btn.disabled = false;
                    btn.textContent = 'Take Screenshot';
                }
            }
            
            // Enter key support
            document.getElementById('url').addEventListener('keypress', function(e) {
                if (e.key === 'Enter') {
                    takeScreenshot();
                }
            });
        </script>
    </body>
    </html>
  `);
});

// 错误处理中间件
app.use((error, req, res, next) => {
  console.error('Unhandled error:', error);
  res.status(500).json({
    error: 'Internal server error'
  });
});

// 404处理
app.use((req, res) => {
  res.status(404).json({
    error: 'Endpoint not found'
  });
});

app.listen(PORT, '0.0.0.0', () => {
  console.log(`Screenshot API server running on port ${PORT} for Hugging Face Spaces`);
});