Profeto/src/db/repositories.py

"""Repository 层:封装数据访问。

Repository 只负责查询,不负责事务提交。
事务由 Application Service 通过 UnitOfWork 控制。
"""
from __future__ import annotations

from datetime import datetime
from sqlalchemy import select
from sqlalchemy.orm import selectinload
from sqlalchemy.ext.asyncio import AsyncSession

from src.db.models import League, Match, Prediction, Team


class MatchRepository:
    """比赛数据访问。"""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_id(self, match_id: int) -> Match | None:
        return await self._session.get(Match, match_id)

    async def get_with_relations(self, match_id: int) -> Match | None:
        stmt = (
            select(Match)
            .options(
                selectinload(Match.league),
                selectinload(Match.home_team),
                selectinload(Match.away_team),
                selectinload(Match.stats),
            )
            .where(Match.id == match_id)
        )
        return (await self._session.execute(stmt)).scalar_one_or_none()

    async def find_by_teams_and_date(
        self, league_id: int, home_team_id: int, away_team_id: int, date
    ) -> Match | None:
        """按联赛+主队+客队+日期查找比赛(天级匹配)。

        预加载 stats:调用方(统计回填)会读取 existing.stats,
        async session 下惰性加载会抛 MissingGreenlet。

        P2-3: 使用 match_date_date(已建索引)做等值匹配,避免 func.date()
        导致的全表扫描。
        """
        if isinstance(date, datetime):
            date = date.date()
        elif hasattr(date, "date"):
            date = date.date()
        else:
            # 字符串等其它格式,尝试转换
            date = datetime.fromisoformat(str(date)).date()

        stmt = (
            select(Match)
            .options(selectinload(Match.stats))
            .where(Match.league_id == league_id)
            .where(Match.home_team_id == home_team_id)
            .where(Match.away_team_id == away_team_id)
            .where(Match.match_date_date == date)
        )
        return (await self._session.execute(stmt)).scalar_one_or_none()

    async def find_by_league_and_date_range(
        self, league_id: int, start, end
    ) -> list[Match]:
        """批量预加载某联赛日期范围内的比赛(ingest 管线内存去重用)。"""
        stmt = (
            select(Match)
            .where(Match.league_id == league_id)
            .where(Match.match_date >= start)
            .where(Match.match_date <= end)
        )
        return (await self._session.execute(stmt)).scalars().all()

    async def find_finished_with_stats(self, league_ids: list[int], *, limit: int) -> list[Match]:
        """已完赛且有上游 event id 的比赛(按日期倒序),供统计回填逐场拉取。

        预加载 stats:调用方需读取 existing.stats 判断是否跳过。
        """
        stmt = (
            select(Match)
            .options(selectinload(Match.stats))
            .where(Match.match_status == "finished")
            .where(Match.source_event_id.is_not(None))
            .where(Match.league_id.in_(league_ids))
            .order_by(Match.match_date.desc())
            .limit(limit)
        )
        return (await self._session.execute(stmt)).scalars().all()

    async def add(self, match: Match) -> None:
        self._session.add(match)
        await self._session.flush()


class TeamRepository:
    """球队数据访问。"""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_name(self, name: str) -> Team | None:
        stmt = select(Team).where(Team.name == name)
        return (await self._session.execute(stmt)).scalar_one_or_none()

    async def get_or_create(self, name: str, *, name_zh: str | None = None) -> Team:
        """按名获取球队,不存在则创建(name_zh 供 bzzoiro 管线写中文名)。"""
        team = await self.get_by_name(name)
        if team is None:
            team = Team(name=name, name_zh=name_zh)
            self._session.add(team)
            await self._session.flush()
        return team

    async def get_all_by_names(self, names: list[str]) -> dict[str, Team]:
        """批量获取球队,返回 name → Team 映射。"""
        if not names:
            return {}
        stmt = select(Team).where(Team.name.in_(names))
        teams = (await self._session.execute(stmt)).scalars().all()
        return {t.name: t for t in teams}

    async def add(self, team: Team) -> None:
        self._session.add(team)
        await self._session.flush()


class LeagueRepository:
    """联赛数据访问。"""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_code(self, code: str) -> League | None:
        stmt = select(League).where(League.code == code)
        return (await self._session.execute(stmt)).scalar_one_or_none()

    async def get_or_create(self, code: str, name: str, country: str | None = None) -> League:
        league = await self.get_by_code(code)
        if league is None:
            league = League(code=code, name=name, country=country)
            self._session.add(league)
            await self._session.flush()
        return league

    async def add(self, league: League) -> None:
        self._session.add(league)
        await self._session.flush()


class PredictionRepository:
    """预测记录数据访问。"""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_id(self, prediction_id: int) -> Prediction | None:
        return await self._session.get(Prediction, prediction_id)

    async def add(self, prediction: Prediction) -> None:
        self._session.add(prediction)
        await self._session.flush()