"""指数成分 Repository 的 SQLAlchemy 实现（B1）。"""

from __future__ import annotations

from collections.abc import Sequence
from datetime import date

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.domain.entities.index import IndexWeight
from app.infrastructure.persistence.sqlalchemy.models.index import IndexWeightModel


class SqlAlchemyIndexConstituentRepository:
    def __init__(self, session: Session) -> None:
        self._session = session

    def upsert_many(self, rows: Sequence[IndexWeight]) -> int:
        if not rows:
            return 0
        existing = {
            (r.index_code, r.trade_date, r.symbol): r
            for r in self._session.scalars(
                select(IndexWeightModel).where(
                    IndexWeightModel.index_code.in_({x.index_code for x in rows})
                )
            )
        }
        changed = 0
        for row in rows:
            key = (row.index_code, row.trade_date, row.symbol)
            model = existing.get(key)
            if model is None:
                self._session.add(IndexWeightModel(**row.model_dump()))
                changed += 1
            else:
                for k, v in row.model_dump(exclude={"index_code", "trade_date", "symbol"}).items():
                    setattr(model, k, v)
                changed += 1
        return changed

    def latest_date(self, index_code: str) -> date | None:
        return self._session.scalar(
            select(IndexWeightModel.trade_date)
            .where(IndexWeightModel.index_code == index_code)
            .order_by(IndexWeightModel.trade_date.desc())
            .limit(1)
        )

    def members_at(self, index_code: str, as_of: date) -> set[str]:
        """<= as_of 最近一期快照的成分（历史成分；无快照 → 空集）。"""
        latest = self._session.scalar(
            select(IndexWeightModel.trade_date)
            .where(
                IndexWeightModel.index_code == index_code,
                IndexWeightModel.trade_date <= as_of,
            )
            .order_by(IndexWeightModel.trade_date.desc())
            .limit(1)
        )
        if latest is None:
            return set()
        rows = self._session.scalars(
            select(IndexWeightModel.symbol).where(
                IndexWeightModel.index_code == index_code,
                IndexWeightModel.trade_date == latest,
            )
        ).all()
        return set(rows)