main.py

from fastapi import FastAPI, Depends, HTTPException, Request
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from database import engine, Base, get_db
import models
import uuid
import schemas
from celery_app import process_submission
from logger import get_logger
from presets import get_all_presets

logger = get_logger("main")

# Create tables
Base.metadata.create_all(bind=engine)

app = FastAPI(title="LabelHub AI 审核 Agent API (Enterprise Edition)", version="1.0.0")

def trace_id(request: Request) -> str:
    incoming = request.headers.get("x-trace-id")
    return incoming if incoming else uuid.uuid4().hex

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"System Error: {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={
            "code": "SYSTEM_ERROR",
            "message": "内部服务器错误，请联系管理员",
            "data": None,
            "traceId": trace_id(request)
        }
    )

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "code": "BAD_REQUEST",
            "message": str(exc.detail),
            "data": None,
            "traceId": trace_id(request)
        }
    )

@app.get("/health")
async def health(request: Request):
    """系统健康探针，供 k8s 或负载均衡器调用"""
    return {
        "code": "SUCCESS",
        "message": "success",
        "data": {"status": "UP"},
        "traceId": trace_id(request)
    }

@app.get("/api/config/presets")
def list_preset_templates():
    """获取系统预制好的 Prompt 模板和维度推荐，供不懂技术的“老板”用户一键套用"""
    logger.info("Fetching preset templates.")
    return {"presets": get_all_presets()}

@app.post("/api/config", response_model=schemas.ReviewConfigResponse)
def create_or_update_config(config: schemas.ReviewConfigCreate, db: Session = Depends(get_db)):
    """后台配置：任务负责人配置审核 Prompt 模板与评分维度"""
    logger.info(f"Configuring task: {config.task_name} with AI: {config.ai_model}")
    db_config = db.query(models.ReviewConfig).filter(models.ReviewConfig.task_name == config.task_name).first()
    if db_config:
        db_config.prompt_template = config.prompt_template
        db_config.dimensions = config.dimensions
        db_config.ai_model = config.ai_model
    else:
        db_config = models.ReviewConfig(
            task_name=config.task_name,
            prompt_template=config.prompt_template,
            dimensions=config.dimensions,
            ai_model=config.ai_model
        )
        db.add(db_config)
    db.commit()
    db.refresh(db_config)
    return db_config

@app.post("/api/submit", response_model=schemas.SubmissionResponse)
def submit_data(submission: schemas.SubmissionCreate, db: Session = Depends(get_db)):
    """标注员提交数据，自动入队"""
    logger.info(f"Received submission: {submission.submission_id}")
    # 幂等性：检查是否已存在 submission_id
    existing = db.query(models.Submission).filter(models.Submission.submission_id == submission.submission_id).first()
    if existing:
        logger.info(f"Submission {submission.submission_id} already exists, skipping.")
        return existing

    db_submission = models.Submission(
        submission_id=submission.submission_id,
        task_name=submission.task_name,
        annotator_id=submission.annotator_id,
        data=submission.data,
        status="pending"
    )
    db.add(db_submission)
    db.commit()
    db.refresh(db_submission)

    # 异步入队
    process_submission.delay(db_submission.id)
    logger.info(f"Enqueued task for submission: {submission.submission_id}")
    
    return db_submission

@app.get("/api/workbench/{submission_id}", response_model=schemas.WorkbenchResponse)
def get_workbench(submission_id: str, db: Session = Depends(get_db)):
    """审核工作台：查看审核结果与原数据，包含 AI 评语与原始 Prompt"""
    logger.info(f"Querying workbench for submission: {submission_id}")
    submission = db.query(models.Submission).filter(models.Submission.submission_id == submission_id).first()
    if not submission:
        raise HTTPException(status_code=404, detail="Submission not found")

    review = db.query(models.ReviewResult).filter(models.ReviewResult.submission_id == submission_id).first()

    return schemas.WorkbenchResponse(
        submission=submission,
        review=review
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)