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, require_admin, require_labeler, require_guest_or_above, _is_admin, ROLE_PROJECT_MANAGER

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(require_admin),
):
    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(require_admin),
):
    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(require_admin),
):
    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(require_admin),
):
    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(require_admin),
):
    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(require_admin),
):
    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(require_guest_or_above),
):
    from app.services.project_service import list_results as _list_results
    from app.models.project import ClassificationProject, ClassificationTask
    from app.api.deps import _has_role

    # Data isolation: compute allowed project IDs for non-admin users
    allowed_project_ids = None
    if not _is_admin(current_user):
        if _has_role(current_user, ROLE_PROJECT_MANAGER):
            # Project managers see their own projects
            allowed_project_ids = [
                p.id for p in db.query(ClassificationProject.id).filter(
                    ClassificationProject.created_by == current_user.id
                ).all()
            ]
        elif _has_role(current_user, 'labeler') or _has_role(current_user, 'reviewer'):
            # Labelers/reviewers see projects where they have tasks
            task_projects = db.query(ClassificationTask.project_id).filter(
                ClassificationTask.assignee_id == current_user.id
            ).distinct().all()
            allowed_project_ids = [p[0] for p in task_projects]
        else:
            # Guests see all results (read-only)
            allowed_project_ids = None
        if allowed_project_ids is not None and not allowed_project_ids:
            allowed_project_ids = []

    # If a specific project_id is requested, check permission
    if project_id and allowed_project_ids is not None and project_id not in allowed_project_ids:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="无权查看此项目")

    items, total = _list_results(
        db, project_id=project_id, keyword=keyword, page=page, page_size=page_size,
        project_ids=allowed_project_ids,
    )

    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

        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)


@router.post("/ml-train")
def ml_train(
    background: bool = True,
    model_name: Optional[str] = None,
    algorithm: str = "logistic_regression",
    db: Session = Depends(get_db),
    current_user: User = Depends(require_admin),
):
    from app.tasks.ml_tasks import train_ml_model_task
    from app.services.ml_service import train_model

    if background:
        task = train_ml_model_task.delay(model_name=model_name, algorithm=algorithm)
        return ResponseModel(data={"task_id": task.id, "status": task.state})
    else:
        mv = train_model(db, model_name=model_name, algorithm=algorithm)
        if mv:
            return ResponseModel(data={"model_id": mv.id, "accuracy": mv.accuracy, "train_samples": mv.train_samples})
        return ResponseModel(message="训练失败：样本不足或发生错误")


@router.get("/ml-suggest/{project_id}")
def ml_suggest(
    project_id: int,
    column_ids: Optional[str] = Query(None),
    top_k: int = Query(3, ge=1, le=5),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    from app.services.ml_service import suggest_for_project_columns
    ids = None
    if column_ids:
        ids = [int(x) for x in column_ids.split(",") if x.strip().isdigit()]
    result = suggest_for_project_columns(db, project_id, column_ids=ids, top_k=top_k)
    if not result.get("success"):
        from fastapi import HTTPException, status
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=result.get("message"))
    return ResponseModel(data=result["suggestions"])