eh-v2/main.py

# main.py
import os
import json
import logging
from pathlib import Path
from typing import Dict, Any, List

import aiofiles
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 常量定义
DOWNLOADS_DIR = "downloads"
MAX_FILENAME_LENGTH = 100
INVALID_FILENAME_CHARS = '<>:"/\\|?*'

# FastAPI应用
app = FastAPI(title="eh-v2")

# 数据模型
class SaveDataRequest(BaseModel):
    url: str
    title: str
    all_images: Dict[str, str]
    total_images: int

class GalleryInfo(BaseModel):
    title: str
    path: str
    total_images: int
    downloaded_images: int

# 工具函数
def setup_downloads_directory() -> Path:
    """创建并返回下载目录路径"""
    downloads_path = Path(DOWNLOADS_DIR)
    downloads_path.mkdir(exist_ok=True)
    logger.info(f"下载目录已准备: {downloads_path.absolute()}")
    return downloads_path

def sanitize_filename(filename: str) -> str:
    """清理文件名，移除非法字符并限制长度"""
    sanitized = filename
    for char in INVALID_FILENAME_CHARS:
        sanitized = sanitized.replace(char, '_')
    
    # 限制文件名长度
    if len(sanitized) > MAX_FILENAME_LENGTH:
        sanitized = sanitized[:MAX_FILENAME_LENGTH]
    
    return sanitized

def create_title_directory(base_path: Path, title: str) -> Path:
    """创建标题对应的目录"""
    safe_title = sanitize_filename(title)
    title_dir = base_path / safe_title
    title_dir.mkdir(exist_ok=True)
    logger.info(f"创建标题目录: {title_dir}")
    return title_dir

async def save_data_to_file(file_path: Path, data: Dict[str, Any]) -> None:
    """异步保存数据到JSON文件"""
    async with aiofiles.open(file_path, 'w', encoding='utf-8') as f:
        await f.write(json.dumps(data, ensure_ascii=False, indent=2))

def get_all_galleries() -> List[GalleryInfo]:
    """获取所有画廊信息"""
    galleries = []
    downloads_path = Path(DOWNLOADS_DIR)
    
    if not downloads_path.exists():
        return galleries
    
    for gallery_dir in downloads_path.iterdir():
        if gallery_dir.is_dir():
            data_file = gallery_dir / "data.json"
            if data_file.exists():
                try:
                    with open(data_file, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    
                    # 计算已下载的图片数量
                    downloaded_count = 0
                    if 'all_images' in data:
                        for filename, url in data['all_images'].items():
                            image_path = gallery_dir / filename
                            if image_path.exists():
                                downloaded_count += 1
                    
                    galleries.append(GalleryInfo(
                        title=data.get('title', gallery_dir.name),
                        path=str(gallery_dir),
                        total_images=data.get('total_images', 0),
                        downloaded_images=downloaded_count
                    ))
                except Exception as e:
                    logger.error(f"读取画廊数据失败 {gallery_dir}: {e}")
    
    return galleries

# 初始化
downloads_path = setup_downloads_directory()

# API路由
@app.post("/save_url")
async def save_url(data: SaveDataRequest):
    """保存URL数据到文件系统"""
    try:
        logger.info("收到保存数据请求")
        logger.info(f"标题: {data.title}, URL: {data.url}, 图片数量: {data.total_images}")
        
        # 创建标题目录
        title_dir = create_title_directory(downloads_path, data.title)
        
        # 数据文件路径
        data_file = title_dir / "data.json"
        
        # 异步保存数据
        await save_data_to_file(data_file, data.dict())
        
        logger.info(f"数据已保存到: {data_file}")
        
        return {
            "status": "success",
            "message": "数据保存成功",
            "file_path": str(data_file),
            "title": data.title,
            "total_images": data.total_images
        }
        
    except Exception as e:
        error_msg = f"保存数据时出错: {str(e)}"
        logger.error(error_msg)
        logger.exception("详细错误信息:")
        raise HTTPException(status_code=500, detail=error_msg)

@app.get("/", response_class=HTMLResponse)
async def read_gallery_manager():
    """画廊管理页面"""
    return """
    <!DOCTYPE html>
    <html lang="zh-CN">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>画廊下载管理器</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }
            body {
                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                padding: 20px;
            }
            .container {
                max-width: 1200px;
                margin: 0 auto;
                background: white;
                border-radius: 15px;
                box-shadow: 0 20px 40px rgba(0,0,0,0.1);
                overflow: hidden;
            }
            .header {
                background: linear-gradient(135deg, #2c3e50, #34495e);
                color: white;
                padding: 30px;
                text-align: center;
            }
            .header h1 {
                font-size: 2.5em;
                margin-bottom: 10px;
            }
            .header p {
                opacity: 0.8;
                font-size: 1.1em;
            }
            .controls {
                padding: 20px;
                background: #f8f9fa;
                border-bottom: 1px solid #e9ecef;
                display: flex;
                gap: 15px;
                flex-wrap: wrap;
            }
            .btn {
                padding: 12px 24px;
                border: none;
                border-radius: 8px;
                font-size: 16px;
                font-weight: 600;
                cursor: pointer;
                transition: all 0.3s ease;
                display: inline-flex;
                align-items: center;
                gap: 8px;
            }
            .btn-primary {
                background: #007bff;
                color: white;
            }
            .btn-primary:hover {
                background: #0056b3;
                transform: translateY(-2px);
            }
            .btn-success {
                background: #28a745;
                color: white;
            }
            .btn-success:hover {
                background: #1e7e34;
                transform: translateY(-2px);
            }
            .btn-danger {
                background: #dc3545;
                color: white;
            }
            .btn-danger:hover {
                background: #c82333;
                transform: translateY(-2px);
            }
            .btn:disabled {
                background: #6c757d;
                cursor: not-allowed;
                transform: none;
            }
            .gallery-list {
                padding: 20px;
            }
            .gallery-item {
                background: white;
                border: 1px solid #e9ecef;
                border-radius: 10px;
                padding: 20px;
                margin-bottom: 15px;
                transition: all 0.3s ease;
                display: flex;
                justify-content: space-between;
                align-items: center;
            }
            .gallery-item:hover {
                box-shadow: 0 5px 15px rgba(0,0,0,0.1);
                transform: translateY(-2px);
            }
            .gallery-info {
                flex: 1;
            }
            .gallery-title {
                font-size: 1.3em;
                font-weight: 600;
                color: #2c3e50;
                margin-bottom: 8px;
            }
            .gallery-stats {
                display: flex;
                gap: 20px;
                color: #6c757d;
                font-size: 0.9em;
            }
            .gallery-actions {
                display: flex;
                gap: 10px;
            }
            .progress-bar {
                width: 200px;
                height: 8px;
                background: #e9ecef;
                border-radius: 4px;
                overflow: hidden;
                margin-top: 8px;
            }
            .progress-fill {
                height: 100%;
                background: linear-gradient(90deg, #28a745, #20c997);
                transition: width 0.3s ease;
            }
            .empty-state {
                text-align: center;
                padding: 60px 20px;
                color: #6c757d;
            }
            .empty-state h3 {
                margin-bottom: 10px;
                font-size: 1.5em;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <div class="header">
                <h1>🎨 画廊下载管理器</h1>
                <p>管理您的画廊下载任务</p>
            </div>
            
            <div class="controls">
                <button class="btn btn-primary" onclick="loadGalleries()">
                    📁 读取文件夹
                </button>
                <button class="btn btn-success" onclick="startDownload()" id="downloadBtn">
                    ⬇️ 开始下载
                </button>
                <button class="btn btn-danger" onclick="deleteJsonFiles()">
                    🗑️ 删除JSON文件
                </button>
            </div>
            
            <div class="gallery-list" id="galleryList">
                <div class="empty-state">
                    <h3>暂无画廊数据</h3>
                    <p>点击"读取文件夹"按钮加载数据</p>
                </div>
            </div>
        </div>

        <script>
            let currentGalleries = [];
            
            async function loadGalleries() {
                try {
                    const response = await fetch('/api/galleries');
                    const galleries = await response.json();
                    currentGalleries = galleries;
                    displayGalleries(galleries);
                } catch (error) {
                    alert('读取文件夹失败: ' + error);
                }
            }
            
            function displayGalleries(galleries) {
                const galleryList = document.getElementById('galleryList');
                
                if (galleries.length === 0) {
                    galleryList.innerHTML = `
                        <div class="empty-state">
                            <h3>暂无画廊数据</h3>
                            <p>未找到任何画廊数据文件</p>
                        </div>
                    `;
                    return;
                }
                
                galleryList.innerHTML = galleries.map(gallery => `
                    <div class="gallery-item">
                        <div class="gallery-info">
                            <div class="gallery-title">${gallery.title}</div>
                            <div class="gallery-stats">
                                <span>总图片: ${gallery.total_images}</span>
                                <span>已下载: ${gallery.downloaded_images}</span>
                                <span>进度: ${Math.round((gallery.downloaded_images / gallery.total_images) * 100)}%</span>
                            </div>
                            <div class="progress-bar">
                                <div class="progress-fill" style="width: ${(gallery.downloaded_images / gallery.total_images) * 100}%"></div>
                            </div>
                        </div>
                        <div class="gallery-actions">
                            <button class="btn btn-primary" onclick="downloadGallery('${gallery.title}')">
                                下载
                            </button>
                            <button class="btn btn-danger" onclick="deleteGallery('${gallery.title}')">
                                删除
                            </button>
                        </div>
                    </div>
                `).join('');
            }
            
            async function startDownload() {
                const btn = document.getElementById('downloadBtn');
                btn.disabled = true;
                btn.innerHTML = '⏳ 下载中...';
                
                try {
                    // 这里可以添加批量下载逻辑
                    for (const gallery of currentGalleries) {
                        if (gallery.downloaded_images < gallery.total_images) {
                            await downloadGallery(gallery.title);
                        }
                    }
                    alert('所有下载任务已完成！');
                } catch (error) {
                    alert('下载失败: ' + error);
                } finally {
                    btn.disabled = false;
                    btn.innerHTML = '⬇️ 开始下载';
                    await loadGalleries(); // 刷新列表
                }
            }
            
            async function downloadGallery(title) {
                try {
                    const response = await fetch(`/api/download/${encodeURIComponent(title)}`, {
                        method: 'POST'
                    });
                    const result = await response.json();
                    
                    if (result.status === 'success') {
                        alert(`开始下载: ${title}`);
                        // 这里可以添加实时进度更新
                    } else {
                        alert(`下载失败: ${result.message}`);
                    }
                } catch (error) {
                    alert('下载请求失败: ' + error);
                }
            }
            
            async function deleteJsonFiles() {
                if (!confirm('确定要删除所有JSON文件吗？此操作不可恢复！')) {
                    return;
                }
                
                try {
                    const response = await fetch('/api/cleanup', {
                        method: 'DELETE'
                    });
                    const result = await response.json();
                    alert(result.message);
                    await loadGalleries(); // 刷新列表
                } catch (error) {
                    alert('删除失败: ' + error);
                }
            }
            
            async function deleteGallery(title) {
                if (!confirm(`确定要删除画廊"${title}"吗？此操作不可恢复！`)) {
                    return;
                }
                
                try {
                    const response = await fetch(`/api/galleries/${encodeURIComponent(title)}`, {
                        method: 'DELETE'
                    });
                    const result = await response.json();
                    alert(result.message);
                    await loadGalleries(); // 刷新列表
                } catch (error) {
                    alert('删除失败: ' + error);
                }
            }
            
            // 页面加载时自动读取
            document.addEventListener('DOMContentLoaded', loadGalleries);
        </script>
    </body>
    </html>
    """

@app.get("/api/galleries")
async def get_galleries():
    """获取所有画廊信息"""
    galleries = get_all_galleries()
    return galleries

@app.post("/api/download/{title}")
async def download_gallery(title: str):
    """开始下载指定画廊的图片"""
    try:
        # 这里实现图片下载逻辑
        # 遍历 all_images 字典，下载每个图片
        return {
            "status": "success",
            "message": f"开始下载画廊: {title}",
            "title": title
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"下载失败: {str(e)}")

@app.delete("/api/cleanup")
async def cleanup_json_files():
    """删除所有JSON文件（保留图片）"""
    try:
        deleted_count = 0
        downloads_path = Path(DOWNLOADS_DIR)
        
        for gallery_dir in downloads_path.iterdir():
            if gallery_dir.is_dir():
                data_file = gallery_dir / "data.json"
                if data_file.exists():
                    data_file.unlink()
                    deleted_count += 1
        
        return {
            "status": "success",
            "message": f"已删除 {deleted_count} 个JSON文件",
            "deleted_count": deleted_count
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"清理失败: {str(e)}")

@app.delete("/api/galleries/{title}")
async def delete_gallery(title: str):
    """删除指定画廊的所有文件"""
    try:
        safe_title = sanitize_filename(title)
        gallery_path = downloads_path / safe_title
        
        if gallery_path.exists():
            # 删除整个画廊目录
            import shutil
            shutil.rmtree(gallery_path)
            return {
                "status": "success",
                "message": f"已删除画廊: {title}"
            }
        else:
            raise HTTPException(status_code=404, detail="画廊不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除失败: {str(e)}")

@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=5100,
        reload=True
    )