csg-server/app/api/v1/module_energy/variable/service.py

from app.api.v1.module_energy.chapter.crud import ChapterCRUD
from app.core.base_schema import AuthSchema
from app.core.exceptions import CustomException

from .crud import VariableCRUD
from .schema import VariableGroupOutSchema, VariableSaveSchema


class VariableService:
    """能源报告变量关联服务。"""

    def __init__(self, auth: AuthSchema) -> None:
        self.auth = auth

    async def list_by_template(self, template_id: int) -> list[dict]:
        records = await VariableCRUD(self.auth).get_list(
            search={"template_id": template_id},
            order_by=[{"variable_name": "asc"}, {"id": "asc"}],
        )
        chapter_map = await self._chapter_summary_map(template_id=template_id)
        group_map: dict[str, VariableGroupOutSchema] = {}
        for record in records:
            key = record.variable_name
            if key not in group_map:
                group_map[key] = VariableGroupOutSchema(
                    variable_name=record.variable_name,
                    chapter_names=[],
                    chapters=[],
                )
            group_map[key].chapter_names.append(record.chapter_name)
            group_map[key].chapters.append(
                chapter_map.get(record.chapter_name, {"chapter_name": record.chapter_name})
            )
        return [group.model_dump() for group in group_map.values()]

    async def list_by_chapter(self, chapter_id: int) -> dict:
        chapter = await ChapterCRUD(self.auth).get_or_404(id=chapter_id, msg="章节标注不存在")
        records = await VariableCRUD(self.auth).get_list(
            search={"template_id": chapter.template_id},
            order_by=[{"variable_name": "asc"}, {"id": "asc"}],
        )
        rule_name = chapter.rule.name if hasattr(chapter, "rule") and chapter.rule else None
        variable_map: dict[str, dict] = {}
        for record in records:
            if record.variable_name not in variable_map:
                variable_map[record.variable_name] = {
                    "variable_name": record.variable_name,
                    "chapter_names": [],
                    "chapters": [],
                    "checked": False,
                }
            variable_map[record.variable_name]["chapter_names"].append(record.chapter_name)

        chapter_map = await self._chapter_summary_map(template_id=chapter.template_id)
        for variable in variable_map.values():
            variable["chapter_names"] = self._unique_names(variable["chapter_names"])
            variable["chapters"] = [
                chapter_map.get(chapter_name, {"chapter_name": chapter_name})
                for chapter_name in variable["chapter_names"]
            ]
            variable["checked"] = chapter.chapter_name in variable["chapter_names"]

        return {
            "chapter_annotation": {
                "id": chapter.id,
                "chapter_name": chapter.chapter_name,
                "rule_id": chapter.rule_id,
                "rule_name": rule_name,
            },
            "variables": list(variable_map.values()),
        }

    async def save(self, data: VariableSaveSchema) -> list[dict]:
        chapter_names = self._collect_chapter_names(data)
        await self._validate_chapters(template_id=data.template_id, chapter_names=chapter_names)

        existing = await VariableCRUD(self.auth).get_list(search={"template_id": data.template_id})
        if existing:
            await VariableCRUD(self.auth).hard_delete(ids=[record.id for record in existing])

        for variable in data.variables:
            for chapter_name in self._unique_names(variable.chapter_names):
                await VariableCRUD(self.auth).create(
                    data={
                        "template_id": data.template_id,
                        "variable_name": variable.variable_name,
                        "chapter_name": chapter_name,
                    }
                )

        return await self.list_by_template(template_id=data.template_id)

    async def delete_by_variable(self, template_id: int, variable_name: str) -> None:
        records = await VariableCRUD(self.auth).get_list(
            search={"template_id": template_id, "variable_name": variable_name},
        )
        if not records:
            raise CustomException(msg="变量关联不存在")
        await VariableCRUD(self.auth).hard_delete(ids=[record.id for record in records])

    def _collect_chapter_names(self, data: VariableSaveSchema) -> list[str]:
        chapter_names: list[str] = []
        for variable in data.variables:
            chapter_names.extend(variable.chapter_names)
        return self._unique_names(chapter_names)

    def _unique_names(self, names: list[str]) -> list[str]:
        return list(dict.fromkeys(name.strip() for name in names if name.strip()))

    async def _validate_chapters(self, template_id: int, chapter_names: list[str]) -> None:
        if not chapter_names:
            return
        chapters = await ChapterCRUD(self.auth).get_list(search={"template_id": template_id})
        exists_names = {chapter.chapter_name for chapter in chapters}
        invalid_names = [chapter_name for chapter_name in chapter_names if chapter_name not in exists_names]
        if invalid_names:
            raise CustomException(msg=f"章节不存在或不属于当前模板: {invalid_names}")

    async def _chapter_summary_map(self, template_id: int) -> dict[str, dict]:
        chapters = await ChapterCRUD(self.auth).get_list(
            search={"template_id": template_id},
            order_by=[{"chapter_order": "asc"}, {"id": "asc"}],
        )
        result: dict[str, dict] = {}
        for chapter in chapters:
            if chapter.chapter_name not in result:
                result[chapter.chapter_name] = {
                    "chapter_name": chapter.chapter_name,
                    "annotation_ids": [],
                    "rules": [],
                }
            rule_name = chapter.rule.name if hasattr(chapter, "rule") and chapter.rule else None
            result[chapter.chapter_name]["annotation_ids"].append(chapter.id)
            result[chapter.chapter_name]["rules"].append(
                {
                    "chapter_annotation_id": chapter.id,
                    "rule_id": chapter.rule_id,
                    "rule_name": rule_name,
                }
            )
        return result