from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified

from app.models.reading import Reading
from app.services.bazi_calculator import BaziCalculator


async def backfill_bazi_wuxing_balance(db: AsyncSession) -> int:
    result = await db.execute(select(Reading).where(Reading.reading_type == "bazi"))
    readings = result.scalars().all()
    calculator = BaziCalculator()
    updated = 0

    for reading in readings:
        input_data = dict(reading.input_data or {})
        chart = input_data.get("chart")
        if not isinstance(chart, dict):
            continue

        pillars = chart.get("pillars")
        if not isinstance(pillars, dict):
            continue

        current = chart.get("wuxing_balance")
        if isinstance(current, dict) and any(int(current.get(name) or 0) for name in ("木", "火", "土", "金", "水")):
            continue

        chart["wuxing_balance"] = calculator._count_wuxing(pillars)
        input_data["chart"] = chart
        reading.input_data = input_data
        flag_modified(reading, "input_data")
        updated += 1

    await db.flush()
    return updated