app.py

# ⚙️ 后端逻辑/核心服务端.py (Hugging Face Spaces app.py)
from fastapi import FastAPI, File, UploadFile, Form, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, JSONResponse, FileResponse 
from sqlalchemy.orm import Session
from pydantic import BaseModel
from huggingface_hub import hf_hub_download, HfApi
import hashlib
import urllib.parse
import urllib.request
import urllib.error
import os
import mimetypes 
import 数据库连接 as db

from router_users import router as users_router
from router_items import router as items_router
from router_comments import router as comments_router
from router_messages import router as messages_router
from router_wallet import router as wallet_router
from router_proxy import router as proxy_router

from database_sql import init_sql_db, get_db
from models_sql import Ownership

app = FastAPI(title="ComfyUI Ranking Community API")

@app.get("/")
def health_check():
    return {"status": "ok", "message": "ComfyUI Ranking API is running perfectly!"}

@app.on_event("startup")
def on_startup():
    init_sql_db()
    print("关系型数据库加载完毕，金融表同步完成。")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False, 
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(users_router)
app.include_router(items_router)
app.include_router(comments_router)
app.include_router(messages_router)
app.include_router(wallet_router)
app.include_router(proxy_router)

# ==========================================
# 🟢 私有图床代理中心 (Image Proxy)
# 解决 Private 仓库下，本地客户端报 401 Unauthorized 的终极方案
# ==========================================
@app.get("/api/image_proxy")
def proxy_hf_image(url: str = None, path: str = None):
    """云端图片代理：使用云端的 HF_TOKEN 提取私有图床图片并返回给没有任何权限的本地端"""
    
    # 兼容处理：如果本地发来的是已经被污染的老版本 HF 直链，我们自动将其转换为相对路径
    if url and url.startswith("https://huggingface.co/datasets/"):
        try:
            path = url.split("resolve/main/")[-1]
            path = urllib.parse.unquote(path)
        except:
            raise HTTPException(status_code=400, detail="无效的 HF 原链接格式")
            
    if not path:
        raise HTTPException(status_code=400, detail="缺少路径参数")
        
    # 🛡️ 绝对的安全红线：限制只能代理下载图片目录，严禁黑客通过此接口下载 users.json 或账本数据！
    allowed_dirs = ["uploads/", "avatars/", "covers/"]
    if not any(path.startswith(d) for d in allowed_dirs):
        raise HTTPException(status_code=403, detail="非法访问：该接口仅允许代理图片资源")
        
    hf_token = os.environ.get("HF_TOKEN")
    dataset_repo_id = db.DATASET_REPO_ID 
    
    try:
        # hf_hub_download 会自动利用云端容器的缓存，只有第一次会去真实请求 Dataset
        cached_file_path = hf_hub_download(
            repo_id=dataset_repo_id,
            repo_type="dataset",
            filename=path,
            token=hf_token
        )
        # 智能识别文件类型 (image/jpeg, image/png 等)
        content_type, _ = mimetypes.guess_type(cached_file_path)
        return FileResponse(cached_file_path, media_type=content_type or "application/octet-stream")
    except Exception as e:
        return JSONResponse(content={"error": f"代理获取图片失败: {str(e)}"}, status_code=404)

# ==========================================
# 上传接口 (将返回的 URL 替换为 Proxy 代理链接)
# ==========================================
@app.post("/api/upload")
def upload_file(file: UploadFile = File(...), file_type: str = Form(...)):
    content = file.file.read() 
    
    # 🟢 动态文件大小风控
    max_size = 10 * 1024 * 1024  
    if file_type == "avatar":
        max_size = 2 * 1024 * 1024   
    elif file_type == "cover":
        max_size = 5 * 1024 * 1024   
        
    if len(content) > max_size:
        raise HTTPException(status_code=400, detail=f"文件过大，{file_type} 类型请限制在 {max_size // (1024*1024)}MB 以内")
        
    ext = file.filename.split(".")[-1].lower()
    if ext not in ["jpg", "jpeg", "png", "gif", "webp", "json", "mp4"]:
        raise HTTPException(status_code=400, detail="不支持的文件格式")
        
    file_hash = hashlib.md5(content).hexdigest()[:10]
    safe_filename = f"{file_type}_{file_hash}.{ext}"
    
    local_tmp_path = f"/tmp/{safe_filename}"
    with open(local_tmp_path, "wb") as f:
        f.write(content)
        
    hf_token = os.environ.get("HF_TOKEN")
    dataset_repo_id = "ZHIWEI666/ComfyUI-Ranking"
    
    try:
        api = HfApi()
        api.upload_file(
            path_or_fileobj=local_tmp_path,
            path_in_repo=f"uploads/{file_type}/{safe_filename}",
            repo_id=dataset_repo_id,
            repo_type="dataset",
            token=hf_token,
            commit_message=f"Upload media: {safe_filename}"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"图床同步失败: {str(e)}")
    finally:
        if os.path.exists(local_tmp_path):
            os.remove(local_tmp_path)
            
    # 🚀 核心修复：不再返回暴露隐私且报 401 的 HF 直链，而是返回我们刚写好的 Proxy 代理链接
    permanent_url = f"https://zhiwei666-comfyui-ranking-api.hf.space/api/image_proxy?path=uploads/{file_type}/{safe_filename}"
    return {"status": "success", "url": permanent_url}


class ValidateResourceRequest(BaseModel):
    url: str
    item_id: str
    account: str

@app.post("/api/validate_resource")
def validate_resource(req_data: ValidateResourceRequest, sql_db: Session = Depends(get_db)):
    target_url = req_data.url
    if not target_url.startswith("https://huggingface.co/datasets/") and not target_url.startswith("https://github.com/"):
        return JSONResponse(content={"error": "无效的下载链接"}, status_code=400)
        
    items_db = db.load_data("items.json", default_data=[])
    item = next((i for i in items_db if i["id"] == req_data.item_id), None)
    
    if not item: return JSONResponse(content={"error": "资源不存在或已被删除"}, status_code=404)
        
    price = int(item.get("price", 0))
    author = item.get("author")
    
    if price > 0 and req_data.account != author:
        owned = sql_db.query(Ownership).filter(Ownership.account == req_data.account, Ownership.item_id == req_data.item_id).first()
        if not owned:
            return JSONResponse(content={"error": "🚨 非法下载：云端数据库未找到您的购买凭证！"}, status_code=403)

    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token: return JSONResponse(content={"error": "云端环境变量未配置 HF_TOKEN"}, status_code=401)
        
    try:
        if target_url.startswith("https://github.com/"):
            creator_token = item.get("github_token")
            fallback_token = os.environ.get("GITHUB_PAT")
            active_token = creator_token if creator_token else fallback_token
            
            headers = {"User-Agent": "ComfyUI-Ranking-SaaS"}
            if active_token:
                headers["Authorization"] = f"Bearer {active_token}"
                
            repo_parts = target_url.rstrip("/").split("/")
            if len(repo_parts) < 2: return JSONResponse(content={"error": "无效的仓库地址格式"}, status_code=400)
            owner, repo = repo_parts[-2], repo_parts[-1]
            api_url = f"https://api.github.com/repos/{owner}/{repo}"
            
            req = urllib.request.Request(api_url, headers=headers)
            with urllib.request.urlopen(req) as response:
                if response.status != 200:
                    return JSONResponse(content={"error": "资源仓库不可访问，可能已被作者删除或设为私有"}, status_code=404)
            return {"status": "success", "message": "资源有效"}
            
        elif target_url.startswith("https://huggingface.co/datasets/"):
            repo_path_encoded = target_url.split("resolve/main/")[-1]
            repo_path = urllib.parse.unquote(repo_path_encoded)
            
            cached_file_path = hf_hub_download(
                repo_id=db.DATASET_REPO_ID,
                repo_type="dataset",
                filename=repo_path,
                token=hf_token
            )
            if not os.path.exists(cached_file_path):
                return JSONResponse(content={"error": "云端文件不存在，可能已被作者删除"}, status_code=404)
                
            return {"status": "success", "message": "资源有效"}
            
    except urllib.error.HTTPError as e:
        return JSONResponse(content={"error": f"资源探测失败，源站返回: {e.code}。请联系作者处理。"}, status_code=400)
    except Exception as e:
        return JSONResponse(content={"error": f"资源探测异常: {str(e)}"}, status_code=500)
        
    return {"status": "success"}