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.get("/results", response_model=ListResponse)
def list_results(
    project_id: Optional[int] = Query(None),
    level_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),
):
    from app.services.project_service import list_results as _list_results
    items, total = _list_results(db, project_id=project_id, keyword=keyword, page=page, page_size=page_size)

    data = []
    for r in items:
        col = r.column
        table = col.table if col else None
        database = table.database if table else None
        source = database.source if database else None

        # Filter by level_id if specified
        if level_id and r.level_id != level_id:
            continue

        data.append({
            "id": r.id,
            "project_id": r.project_id,
            "column_id": col.id if col else None,
            "column_name": col.name if col else None,
            "data_type": col.data_type if col else None,
            "comment": col.comment if col else None,
            "table_name": table.name if table else None,
            "database_name": database.name if database else None,
            "source_name": source.name if source else None,
            "category_id": r.category_id,
            "category_name": r.category.name if r.category else None,
            "level_id": r.level_id,
            "level_name": r.level.name if r.level else None,
            "level_color": r.level.color if r.level else None,
            "source": r.source,
            "confidence": r.confidence,
            "status": r.status,
            "created_at": r.created_at.isoformat() if r.created_at else None,
        })

    return ListResponse(data=data, total=total, page=page, page_size=page_size)


@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)