Pre_Seg_Server/backend/routers/templates.py

"""Template (Ontology) CRUD endpoints."""

import logging
from typing import List

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import or_
from sqlalchemy.orm import Session

from database import get_db
from models import Template, User
from routers.auth import get_current_user, require_editor
from schemas import TemplateCreate, TemplateOut, TemplateUpdate

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/templates", tags=["Templates"])
RESERVED_UNCLASSIFIED_CLASS = {
    "id": "reserved-unclassified",
    "name": "待分类",
    "color": "#000000",
    "zIndex": 0,
    "maskId": 0,
    "category": "系统保留",
}


def _is_reserved_class(item: dict) -> bool:
    return (
        item.get("id") == RESERVED_UNCLASSIFIED_CLASS["id"]
        or item.get("name") == RESERVED_UNCLASSIFIED_CLASS["name"]
        or item.get("maskId") == 0
    )


def _normalize_template_classes(classes: list[dict] | None) -> list[dict]:
    normalized = [item for item in (classes or []) if not _is_reserved_class(item)]
    return [*normalized, dict(RESERVED_UNCLASSIFIED_CLASS)]


def _pack_mapping_rules(data: dict) -> dict:
    """Pack classes/rules into mapping_rules for DB storage."""
    mapping = data.get("mapping_rules") or {}
    if "classes" in data and data["classes"] is not None:
        mapping["classes"] = _normalize_template_classes(data.pop("classes"))
    if "rules" in data and data["rules"] is not None:
        mapping["rules"] = data.pop("rules")
    if "classes" in mapping:
        mapping["classes"] = _normalize_template_classes(mapping.get("classes"))
    data["mapping_rules"] = mapping
    return data


def _unpack_template(template: Template) -> Template:
    """Unpack mapping_rules into classes/rules for response."""
    mapping = template.mapping_rules or {}
    # Set as attributes so Pydantic from_attributes can pick them up
    template.classes = _normalize_template_classes(mapping.get("classes", []))
    template.rules = mapping.get("rules", [])
    return template


@router.post(
    "",
    response_model=TemplateOut,
    status_code=status.HTTP_201_CREATED,
    summary="Create a new template",
)
def create_template(
    payload: TemplateCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> Template:
    """Create a new ontology template / segmentation class."""
    data = payload.model_dump()
    data = _pack_mapping_rules(data)
    template = Template(**data, owner_user_id=current_user.id)
    db.add(template)
    db.commit()
    db.refresh(template)
    _unpack_template(template)
    logger.info("Created template id=%s name=%s", template.id, template.name)
    return template


@router.get(
    "",
    response_model=List[TemplateOut],
    summary="List all templates",
)
def list_templates(
    skip: int = 0,
    limit: int = 100,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> List[Template]:
    """Retrieve all ontology templates."""
    templates = (
        db.query(Template)
        .filter(or_(Template.owner_user_id == current_user.id, Template.owner_user_id.is_(None)))
        .offset(skip)
        .limit(limit)
        .all()
    )
    for t in templates:
        _unpack_template(t)
    return templates


@router.get(
    "/{template_id}",
    response_model=TemplateOut,
    summary="Get a single template",
)
def get_template(
    template_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> Template:
    """Retrieve a template by its ID."""
    template = db.query(Template).filter(
        Template.id == template_id,
        or_(Template.owner_user_id == current_user.id, Template.owner_user_id.is_(None)),
    ).first()
    if not template:
        raise HTTPException(status_code=404, detail="Template not found")
    _unpack_template(template)
    return template


@router.patch(
    "/{template_id}",
    response_model=TemplateOut,
    summary="Update a template",
)
def update_template(
    template_id: int,
    payload: TemplateUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> Template:
    """Update template fields partially."""
    template = db.query(Template).filter(
        Template.id == template_id,
        or_(Template.owner_user_id == current_user.id, Template.owner_user_id.is_(None)),
    ).first()
    if not template:
        raise HTTPException(status_code=404, detail="Template not found")

    data = payload.model_dump(exclude_unset=True)
    if "classes" in data or "rules" in data:
        data = _pack_mapping_rules(data)

    for key, value in data.items():
        setattr(template, key, value)

    db.commit()
    db.refresh(template)
    _unpack_template(template)
    logger.info("Updated template id=%s", template_id)
    return template


@router.delete(
    "/{template_id}",
    status_code=status.HTTP_204_NO_CONTENT,
    summary="Delete a template",
)
def delete_template(
    template_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> None:
    """Delete a template. Associated annotations will have template_id set to NULL."""
    template = db.query(Template).filter(
        Template.id == template_id,
        or_(Template.owner_user_id == current_user.id, Template.owner_user_id.is_(None)),
    ).first()
    if not template:
        raise HTTPException(status_code=404, detail="Template not found")

    db.delete(template)
    db.commit()
    logger.info("Deleted template id=%s", template_id)