from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import os

import os
from pathlib import Path


# 验证缓存目录可写
cache_dir = Path(os.getenv("HF_HOME", ""))
if not cache_dir.exists():
    cache_dir.mkdir(parents=True, exist_ok=True)
    
test_file = cache_dir / "permission_test.txt"
try:
    with open(test_file, "w") as f:
        f.write("test")
    os.remove(test_file)
    print("✅ Cache directory is writable")
except Exception as e:
    print(f"❌ Cache directory write failed: {str(e)}")
    raise

# 正确加载模型（从缓存或下载）
classifier = pipeline(
    "text-classification",
    model="mrm8488/codebert-base-finetuned-detect-insecure-code"
)

app = FastAPI()

class CodeRequest(BaseModel):
    code: str  # 输入参数定义

@app.post("/detect")
async def detect_insecure_code(request: CodeRequest):
    try:
        # 直接传递代码字符串到分类器
        result = classifier(request.code)
        return {"status": "success", "result": result[0]}
    except Exception as e:
        return {"status": "error", "message": str(e)}