from __future__ import annotations

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

from supply_infra.db.models.category_tree_weight import CategoryTreeWeight
from supply_infra.db.repositories.base import BaseRepository

_BATCH_SIZE = 500

_DIM_KEYS = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
    "real_rov_7d",
    "real_vov_7d",
)

_SCORE_COLUMNS = (
    "ext_pop_score",
    "plat_sust_pop_score",
    "plat_ly_pop_score",
    "recent_pop_score",
    "total_score",
)

_UPSERT_COLUMNS = tuple(
    f"{dim}_{suffix}" for dim in _DIM_KEYS for suffix in ("avg", "count")
) + ("hung_word_count",)


class CategoryTreeWeightRepository(BaseRepository[CategoryTreeWeight]):
    """类目树节点热度权重表 repository。"""

    model = CategoryTreeWeight

    def get_latest_biz_dt(self) -> str | None:
        """返回权重表中最新业务日；无数据时返回 None。"""
        stmt = select(func.max(CategoryTreeWeight.biz_dt))
        return self.session.scalar(stmt)

    def has_biz_dt(self, biz_dt: str) -> bool:
        """指定业务日是否存在权重数据。"""
        stmt = select(CategoryTreeWeight.id).where(
            CategoryTreeWeight.biz_dt == biz_dt
        ).limit(1)
        return self.session.scalars(stmt).first() is not None

    def list_by_biz_dt(self, biz_dt: str) -> list[CategoryTreeWeight]:
        """返回指定业务日的全部节点权重行。"""
        stmt = select(CategoryTreeWeight).where(CategoryTreeWeight.biz_dt == biz_dt)
        return list(self.session.scalars(stmt).all())

    def get_by_category_ids(
        self,
        category_ids: list[int],
        biz_dt: str | None = None,
    ) -> list[CategoryTreeWeight]:
        """按 category_id 列表查询权重行；未指定 biz_dt 时取每个节点各自的最新业务日。"""
        if not category_ids:
            return []

        if biz_dt:
            stmt = select(CategoryTreeWeight).where(
                CategoryTreeWeight.category_id.in_(category_ids),
                CategoryTreeWeight.biz_dt == biz_dt,
            )
            return list(self.session.scalars(stmt).all())

        latest_dt_subq = (
            select(
                CategoryTreeWeight.category_id,
                func.max(CategoryTreeWeight.biz_dt).label("max_biz_dt"),
            )
            .where(CategoryTreeWeight.category_id.in_(category_ids))
            .group_by(CategoryTreeWeight.category_id)
            .subquery()
        )
        stmt = select(CategoryTreeWeight).join(
            latest_dt_subq,
            (CategoryTreeWeight.category_id == latest_dt_subq.c.category_id)
            & (CategoryTreeWeight.biz_dt == latest_dt_subq.c.max_biz_dt),
        )
        return list(self.session.scalars(stmt).all())

    def get_by_category_biz_dt(
        self, category_id: int, biz_dt: str
    ) -> CategoryTreeWeight | None:
        """按类目 + 业务日查询单行。"""
        stmt = select(CategoryTreeWeight).where(
            CategoryTreeWeight.category_id == category_id,
            CategoryTreeWeight.biz_dt == biz_dt,
        )
        return self.session.scalars(stmt).first()

    def upsert_rows(self, rows: list[dict]) -> int:
        """按 (category_id, biz_dt) 批量 upsert。"""
        if not rows:
            return 0

        affected = 0
        for i in range(0, len(rows), _BATCH_SIZE):
            batch = rows[i : i + _BATCH_SIZE]
            stmt = insert(CategoryTreeWeight).values(batch)
            stmt = stmt.on_duplicate_key_update(
                **{col: stmt.inserted[col] for col in _UPSERT_COLUMNS}
            )
            result = self.session.execute(stmt)
            affected += result.rowcount or 0
        return affected

    def update_rank_scores(self, rows: list[dict]) -> int:
        """按 (category_id, biz_dt) 批量更新四维排名分与 total_score。"""
        if not rows:
            return 0

        from sqlalchemy import update

        affected = 0
        for row in rows:
            stmt = (
                update(CategoryTreeWeight)
                .where(
                    CategoryTreeWeight.category_id == row["category_id"],
                    CategoryTreeWeight.biz_dt == row["biz_dt"],
                )
                .values(
                    **{col: row[col] for col in _SCORE_COLUMNS}
                )
            )
            result = self.session.execute(stmt)
            affected += result.rowcount or 0
        return affected
