dpb/backend/app/api/v1/module_bre/plan/service.py

from typing import Any

from fastapi import UploadFile
from sqlalchemy import false, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.base_schema import AuthSchema, PageResultSchema, ImportResultSchema
from app.core.exceptions import CustomException
from app.core.logger import logger
from app.core.base_crud import assert_no_children
from app.utils.common_util import search_to_dict
from app.utils.excel_util import ExcelUtil

from .crud import BreedingPlanCRUD
from .schema import (
    PlanCreateSchema,
    PlanOutSchema,
    PlanQueryParam,
    PlanUpdateSchema,
)
from app.api.v1.module_bre.target.model import TargetModel
from app.api.v1.module_bre.target.crud import BreedingTargetCRUD
from app.api.v1.module_bre.trial.model import TrialModel
from app.api.v1.module_bre.cross_combination.model import CrossCombinationModel
from app.api.v1.module_bre.trial_study.model import TrialStudyModel
from app.api.v1.module_bre.tree.model import TreeModel


def _is_blank(v: Any) -> bool:
    return v is None or (isinstance(v, str) and v.strip() == "")


def _none_if_blank(v: Any) -> Any:
    if _is_blank(v):
        return None
    return str(v).strip() if isinstance(v, str) else v


def _to_float(v: Any) -> float | None:
    if _is_blank(v):
        return None
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


def _to_int(v: Any) -> int | None:
    if _is_blank(v):
        return None
    try:
        return int(float(v))
    except (TypeError, ValueError):
        return None


class PlanService:
    """育种计划 模块服务层"""

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

    async def _attach_fk_labels(self, items: list[PlanOutSchema]) -> None:
        if not items:
            return
        crud = BreedingPlanCRUD(self.auth, self.db)

    async def detail(self, id: int) -> PlanOutSchema:
        obj = await BreedingPlanCRUD(self.auth, self.db).get(id=id)
        if not obj:
            raise CustomException(msg="该育种计划不存在")
        out = PlanOutSchema.model_validate(obj)
        await self._attach_fk_labels([out])
        return out

    async def get_list(
        self,
        search: PlanQueryParam | None = None,
        order_by: list[dict[str, str]] | None = None,
    ) -> list[PlanOutSchema]:
        obj_list = await BreedingPlanCRUD(self.auth, self.db).get_list(
            search=search_to_dict(search), order_by=order_by
        )
        outs = [PlanOutSchema.model_validate(obj) for obj in obj_list]
        await self._attach_fk_labels(outs)
        return outs

    async def page(
        self,
        page_no: int,
        page_size: int,
        search: PlanQueryParam | None = None,
        order_by: list[dict[str, str]] | None = None,
    ) -> PageResultSchema[PlanOutSchema]:
        offset = (page_no - 1) * page_size
        result = await BreedingPlanCRUD(self.auth, self.db).page(
            offset=offset,
            limit=page_size,
            order_by=order_by or [{"id": "asc"}],
            search=search_to_dict(search, {}),
            out_schema=PlanOutSchema,
        )
        await self._attach_fk_labels(result.items)
        return result

    async def create(self, data: PlanCreateSchema) -> PlanOutSchema:
        exist_obj = await BreedingPlanCRUD(self.auth, self.db).get(plan_code=data.plan_code)
        if exist_obj:
            raise CustomException(msg="创建失败，计划编号已存在")
        obj = await BreedingPlanCRUD(self.auth, self.db).create(data=data)
        out = PlanOutSchema.model_validate(obj)
        return out

    async def update(self, id: int, data: PlanUpdateSchema) -> PlanOutSchema:
        obj = await BreedingPlanCRUD(self.auth, self.db).get(id=id)
        if not obj:
            raise CustomException(msg="更新失败，该育种计划不存在")
        if data.plan_code is not None:
            exist_obj = await BreedingPlanCRUD(self.auth, self.db).get(plan_code=data.plan_code)
            if exist_obj and exist_obj.id != id:
                raise CustomException(msg="更新失败，计划编号重复")
        obj = await BreedingPlanCRUD(self.auth, self.db).update(id=id, data=data)
        out = PlanOutSchema.model_validate(obj)
        return out

    async def delete(self, ids: list[int]) -> None:
        if not ids:
            raise CustomException(msg="删除失败，删除对象不能为空")
        objs = await BreedingPlanCRUD(self.auth, self.db).get_list(search={"id": ("in", ids)})
        obj_map = {o.id: o for o in objs}
        for id_ in ids:
            if id_ not in obj_map:
                raise CustomException(msg="删除失败，该育种计划不存在")
        await assert_no_children(
            self.db,
            ids,
            [
                (TargetModel, 'plan_id', '育种目标'),
                (TrialModel, 'plan_id', '试验'),
            ],
        )
        await BreedingPlanCRUD(self.auth, self.db).delete(ids=ids)

    async def trace(self, plan_id: int) -> dict[str, Any]:
        """计划→目标→组合→树 与 计划→试验→试验点→树 双链可追溯汇总。

        返回 plan + 下游 targets(含组合/树数) + trials(含试验点/树数) + totals。
        """
        plan = await BreedingPlanCRUD(self.auth, self.db).get(id=plan_id)
        if not plan:
            raise CustomException(msg="该育种计划不存在")

        targets = (
            await self.db.execute(
                select(TargetModel).where(TargetModel.plan_id == plan_id, TargetModel.is_deleted == false())
            )
        ).scalars().all()
        target_ids = [t.id for t in targets]

        combos_by_target: dict[int, list[dict[str, Any]]] = {}
        if target_ids:
            combos = (
                await self.db.execute(
                    select(CrossCombinationModel).where(
                        CrossCombinationModel.bre_target_id.in_(target_ids),
                        CrossCombinationModel.is_deleted == false(),
                    )
                )
            ).scalars().all()
            combo_ids = [c.id for c in combos]
            combo_trees: dict[int, int] = {}
            if combo_ids:
                for cid, n in (
                    await self.db.execute(
                        select(TreeModel.combination_id, func.count())
                        .where(TreeModel.combination_id.in_(combo_ids), TreeModel.is_deleted == false())
                        .group_by(TreeModel.combination_id)
                    )
                ).all():
                    combo_trees[cid] = n
            for c in combos:
                combos_by_target.setdefault(c.bre_target_id, []).append(
                    {"id": c.id, "combination_code": c.combination_code, "trees": combo_trees.get(c.id, 0)}
                )

        trials = (
            await self.db.execute(
                select(TrialModel).where(TrialModel.plan_id == plan_id, TrialModel.is_deleted == false())
            )
        ).scalars().all()
        trial_ids = [r.id for r in trials]
        target_names: dict[int, str] = {}
        tri_target_ids = {r.target_id for r in trials if r.target_id}
        if tri_target_ids:
            tri_targets = (
                await self.db.execute(
                    select(TargetModel).where(TargetModel.id.in_(tri_target_ids), TargetModel.is_deleted == false())
                )
            ).scalars().all()
            target_names = {t.id: t.target_name for t in tri_targets}

        studies_by_trial: dict[int, list[dict[str, Any]]] = {}
        if trial_ids:
            studies = (
                await self.db.execute(
                    select(TrialStudyModel).where(
                        TrialStudyModel.trial_id.in_(trial_ids), TrialStudyModel.is_deleted == false()
                    )
                )
            ).scalars().all()
            study_ids = [s.id for s in studies]
            study_trees: dict[int, int] = {}
            if study_ids:
                for sid, n in (
                    await self.db.execute(
                        select(TreeModel.trial_study_id, func.count())
                        .where(TreeModel.trial_study_id.in_(study_ids), TreeModel.is_deleted == false())
                        .group_by(TreeModel.trial_study_id)
                    )
                ).all():
                    study_trees[sid] = n
            for s in studies:
                studies_by_trial.setdefault(s.trial_id, []).append(
                    {"id": s.id, "study_name": s.study_name, "trees": study_trees.get(s.id, 0)}
                )

        target_out = []
        for t in targets:
            combos = combos_by_target.get(t.id, [])
            target_out.append({
                "id": t.id,
                "target_name": t.target_name,
                "plan_id": t.plan_id,
                "combinations": combos,
                "n_combinations": len(combos),
                "n_trees": sum(c["trees"] for c in combos),
            })

        trial_out = []
        for r in trials:
            studies = studies_by_trial.get(r.id, [])
            trial_out.append({
                "id": r.id,
                "trial_name": r.trial_name,
                "plan_id": r.plan_id,
                "target_id": r.target_id,
                "target_name": target_names.get(r.target_id),
                "trial_studies": studies,
                "n_trial_studies": len(studies),
                "n_trees": sum(s["trees"] for s in studies),
            })

        return {
            "plan": {"id": plan.id, "plan_code": plan.plan_code, "plan_name": plan.plan_name},
            "targets": target_out,
            "trials": trial_out,
            "totals": {
                "targets": len(target_out),
                "trials": len(trial_out),
                "combinations": sum(t["n_combinations"] for t in target_out),
                "trial_studies": sum(r["n_trial_studies"] for r in trial_out),
                "trees": sum(t["n_trees"] for t in target_out) + sum(r["n_trees"] for r in trial_out),
            },
        }

    async def list_options(self) -> list[dict[str, Any]]:
        """供前端下拉选择使用：返回 [{value, label}]。"""
        obj_list = await BreedingPlanCRUD(self.auth, self.db).get_list(order_by=[{"id": "asc"}])
        return [{"value": o.id, "label": o.plan_code} for o in obj_list]

    @staticmethod
    def batch_export(obj_list: list[dict[str, Any]]) -> bytes:
        mapping_dict = {
            "plan_code": "计划编号",
            "plan_name": "计划名称",
            "objective": "计划目标",
            "start_year": "起始年份",
            "end_year": "结束年份",
            "leader": "负责人",
            "remark": "备注",
            "created_time": "创建时间",
            "created_by": "创建者",
        }
        data = [dict(item) for item in obj_list]
        for item in data:
            creator = item.get("created_by")
            item["created_by"] = creator.get("name", "未知") if isinstance(creator, dict) else "未知"
        return ExcelUtil.export_list2excel(list_data=data, mapping_dict=mapping_dict)

    async def batch_import(self, file: UploadFile, update_support: bool = False) -> ImportResultSchema:
        header_dict = {
            "计划编号": "plan_code",
            "计划名称": "plan_name",
            "计划目标": "objective",
            "起始年份": "start_year",
            "结束年份": "end_year",
            "负责人": "leader",
            "备注": "remark",
        }
        try:
            contents = await file.read()
            rows = ExcelUtil.read_excel_to_dicts(contents)
            await file.close()
            if not rows:
                raise CustomException(msg="导入文件为空")
            missing_headers = [h for h in header_dict if h not in rows[0]]
            if missing_headers:
                raise CustomException(msg=f"导入文件缺少必要的列: {', '.join(missing_headers)}")
            mapped_rows = []
            for row in rows:
                mapped_rows.append({en: row.get(ch) for ch, en in header_dict.items()})
            required_fields = ["plan_code", "plan_name"]
            errors = []
            for field in required_fields:
                missing_indices = [i + 1 for i, r in enumerate(mapped_rows) if _is_blank(r.get(field))]
                if missing_indices:
                    field_name = next(k for k, v in header_dict.items() if v == field)
                    rows_str = "、".join(str(i) for i in missing_indices)
                    errors.append(f"{field_name}不能为空，第{rows_str}行")
            if errors:
                raise CustomException(msg=f"导入失败，以下行缺少必要字段：\n{'; '.join(errors)}")
            error_msgs: list[str] = []
            success_count = 0
            crud = BreedingPlanCRUD(self.auth, self.db)
            for i, row in enumerate(mapped_rows, start=1):
                try:
                    fields = {
                        "plan_code": _none_if_blank(row.get("plan_code")),
                        "plan_name": _none_if_blank(row.get("plan_name")),
                        "objective": _none_if_blank(row.get("objective")),
                        "start_year": _to_int(row.get("start_year")),
                        "end_year": _to_int(row.get("end_year")),
                        "leader": _none_if_blank(row.get("leader")),
                        "remark": _none_if_blank(row.get("remark")),
                    }
                    unique_kwargs = {"plan_code": fields["plan_code"]}
                    create_data = PlanCreateSchema(**fields)
                    exist_obj = await crud.get(**unique_kwargs)
                    if exist_obj:
                        if update_support:
                            await crud.update(id=exist_obj.id, data=PlanUpdateSchema(**fields))
                            success_count += 1
                        else:
                            error_msgs.append(f"第{i}行: 计划编号 {fields['plan_code']} 已存在")
                    else:
                        await crud.create(data=create_data)
                        success_count += 1
                except Exception as e:
                    error_msgs.append(f"第{i}行: {e!s}")
                    continue
            return ImportResultSchema(
                valid_count=success_count,
                invalid_count=len(error_msgs),
                message_list=error_msgs,
            )
        except Exception as e:
            logger.error(f"批量导入育种计划失败: {e!s}")
            raise CustomException(msg=f"导入失败: {e!s}")

    @staticmethod
    def import_template_download() -> bytes:
        header_list = [
            "计划编号",
            "计划名称",
            "计划目标",
            "起始年份",
            "结束年份",
            "负责人",
            "备注",
        ]
        selector_header_list = []
        option_list = [
        ]
        return ExcelUtil.get_excel_template(
            header_list=header_list,
            selector_header_list=selector_header_list,
            option_list=option_list,
        )