from __future__ import annotations

from sqlalchemy import func, select
from sqlalchemy.dialects.mysql import insert

from supply_infra.db.models.global_tree_category import GlobalTreeCategory
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 1000


class GlobalTreeCategoryRepository(BaseRepository[GlobalTreeCategory]):
    """全局树分类表 — 按 source_id 去重后增量写入。"""

    model = GlobalTreeCategory

    def list_active_categories(self) -> list[GlobalTreeCategory]:
        """返回所有未删除的分类，按 level、id 排序。"""
        stmt = (
            select(GlobalTreeCategory)
            .where(GlobalTreeCategory.is_delete == 0)
            .order_by(GlobalTreeCategory.level, GlobalTreeCategory.id)
        )
        return list(self.session.scalars(stmt).all())

    def list_active_by_ids(self, category_ids: list[int]) -> list[GlobalTreeCategory]:
        """按 id 列表返回未删除分类。"""
        if not category_ids:
            return []
        stmt = select(GlobalTreeCategory).where(
            GlobalTreeCategory.is_delete == 0,
            GlobalTreeCategory.id.in_(category_ids),
        )
        return list(self.session.scalars(stmt).all())

    def get_source_id_map(self) -> dict[int, int]:
        """返回全量 source_id → MySQL id 映射，用于增量去重与 parent 关联。"""
        stmt = select(GlobalTreeCategory.source_id, GlobalTreeCategory.id)
        return {int(row.source_id): int(row.id) for row in self.session.execute(stmt).all()}

    def get_max_id(self) -> int:
        """返回当前表内最大 id，用于分配新 id（兼容无 source_id 的历史数据）。"""
        max_id = self.session.scalar(select(func.max(GlobalTreeCategory.id)))
        return int(max_id) if max_id is not None else 0

    def bulk_insert_ignore(self, rows: list[dict]) -> int:
        """批量插入，MySQL 按 source_id 唯一索引忽略已存在行（历史数据不变）。"""
        if not rows:
            return 0

        inserted = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(GlobalTreeCategory).values(batch).prefix_with("IGNORE")
            result = self.session.execute(stmt)
            inserted += result.rowcount
        return inserted
