prop-data-guard/backend/app/api/v1/classification.py

from typing import Optional
from fastapi import APIRouter, Depends, Query
from sqlalchemy.orm import Session

from app.core.database import get_db
from app.models.user import User
from app.schemas.classification import (
    CategoryCreate, CategoryUpdate, CategoryOut, CategoryTree,
    DataLevelOut, RecognitionRuleCreate, RecognitionRuleUpdate, RecognitionRuleOut,
    TemplateOut,
)
from app.schemas.common import ResponseModel, ListResponse
from app.services import classification_service, classification_engine
from app.api.deps import get_current_user

router = APIRouter()


@router.get("/categories/tree", response_model=ResponseModel[list])
def get_category_tree(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    tree = classification_service.build_category_tree(db)
    return ResponseModel(data=tree)


@router.get("/categories", response_model=ResponseModel[list[CategoryOut]])
def list_categories(
    parent_id: Optional[int] = Query(None),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    items = classification_service.list_categories(db, parent_id=parent_id)
    return ResponseModel(data=[CategoryOut.model_validate(i) for i in items])


@router.post("/categories", response_model=ResponseModel[CategoryOut])
def create_category(
    req: CategoryCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    item = classification_service.create_category(db, req)
    return ResponseModel(data=CategoryOut.model_validate(item))


@router.put("/categories/{category_id}", response_model=ResponseModel[CategoryOut])
def update_category(
    category_id: int,
    req: CategoryUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    db_obj = classification_service.get_category(db, category_id)
    if not db_obj:
        from fastapi import HTTPException, status
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="分类不存在")
    item = classification_service.update_category(db, db_obj, req)
    return ResponseModel(data=CategoryOut.model_validate(item))


@router.delete("/categories/{category_id}")
def delete_category(
    category_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    classification_service.delete_category(db, category_id)
    return ResponseModel(message="删除成功")


@router.get("/levels", response_model=ResponseModel[list[DataLevelOut]])
def list_levels(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    items = classification_service.list_data_levels(db)
    return ResponseModel(data=[DataLevelOut.model_validate(i) for i in items])


@router.get("/rules", response_model=ListResponse[RecognitionRuleOut])
def list_rules(
    template_id: Optional[int] = Query(None),
    keyword: Optional[str] = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=500),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    items, total = classification_service.list_rules(db, template_id=template_id, keyword=keyword, page=page, page_size=page_size)
    out = []
    for i in items:
        data = RecognitionRuleOut.model_validate(i)
        data.category_name = i.category.name if i.category else None
        data.level_name = i.level.name if i.level else None
        data.level_color = i.level.color if i.level else None
        out.append(data)
    return ListResponse(data=out, total=total, page=page, page_size=page_size)


@router.post("/rules", response_model=ResponseModel[RecognitionRuleOut])
def create_rule(
    req: RecognitionRuleCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    item = classification_service.create_rule(db, req)
    data = RecognitionRuleOut.model_validate(item)
    data.category_name = item.category.name if item.category else None
    data.level_name = item.level.name if item.level else None
    data.level_color = item.level.color if item.level else None
    return ResponseModel(data=data)


@router.put("/rules/{rule_id}", response_model=ResponseModel[RecognitionRuleOut])
def update_rule(
    rule_id: int,
    req: RecognitionRuleUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    db_obj = classification_service.get_rule(db, rule_id)
    if not db_obj:
        from fastapi import HTTPException, status
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="规则不存在")
    item = classification_service.update_rule(db, db_obj, req)
    data = RecognitionRuleOut.model_validate(item)
    data.category_name = item.category.name if item.category else None
    data.level_name = item.level.name if item.level else None
    data.level_color = item.level.color if item.level else None
    return ResponseModel(data=data)


@router.delete("/rules/{rule_id}")
def delete_rule(
    rule_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    classification_service.delete_rule(db, rule_id)
    return ResponseModel(message="删除成功")


@router.get("/templates", response_model=ResponseModel[list[TemplateOut]])
def list_templates(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    items = classification_service.list_templates(db)
    return ResponseModel(data=[TemplateOut.model_validate(i) for i in items])


@router.post("/auto-classify/{project_id}")
def auto_classify(
    project_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    result = classification_engine.run_auto_classification(db, project_id)
    return ResponseModel(data=result)