import asyncio from datetime import date from decimal import Decimal from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from app.db.session import async_session_factory from app.models.car import Car, CarMake, CarModel from app.models.expense import FuelEntry, ServiceEntry, ServiceType from app.models.user import User from app.services.catalog_data import CAR_CATALOG MOCK_PLATE_PREFIX = "MOCK" MOCK_CARS = [ ("KIA Sportage", "KIA", "Sportage", 2021, "gasoline", 36200, Decimal("2450000")), ("Toyota Camry", "Toyota", "Camry", 2020, "gasoline", 58400, Decimal("2850000")), ("Hyundai Tucson", "Hyundai", "Tucson", 2022, "gasoline", 27100, Decimal("2750000")), ("Volkswagen Tiguan", "Volkswagen", "Tiguan", 2019, "gasoline", 73400, Decimal("2300000")), ("BMW X3", "BMW", "X3", 2021, "diesel", 48900, Decimal("4350000")), ("Mercedes GLC", "Mercedes", "GLC", 2020, "gasoline", 52200, Decimal("4500000")), ("Nissan X-Trail", "Nissan", "X-Trail", 2018, "gasoline", 91400, Decimal("1850000")), ("Skoda Octavia", "Skoda", "Octavia", 2021, "gasoline", 46800, Decimal("2050000")), ("Tesla Model 3", "Tesla", "Model 3", 2022, "electric", 33800, Decimal("3900000")), ("Haval Jolion", "Haval", "Jolion", 2023, "gasoline", 19600, Decimal("2150000")), ] def month_shift(base: date, months_back: int) -> date: month_index = base.year * 12 + base.month - 1 - months_back year = month_index // 12 month = month_index % 12 + 1 return date(year, month, min(base.day, 24)) async def seed_catalog(session: AsyncSession) -> None: result = await session.execute(select(CarMake).options(selectinload(CarMake.models))) existing = {make.name: make for make in result.scalars()} for make_name, model_names in CAR_CATALOG.items(): make = existing.get(make_name) if make is None: make = CarMake(name=make_name) session.add(make) await session.flush() existing_models = set() else: existing_models = {model.name for model in make.models} for model_name in model_names: if model_name not in existing_models: session.add(CarModel(make_id=make.id, name=model_name)) async def pick_owner(session: AsyncSession) -> User: result = await session.execute(select(User).where(User.telegram_id == 1)) user = result.scalar_one_or_none() if user: return user user = User(telegram_id=1, username="demo", first_name="Demo") session.add(user) await session.flush() return user async def clear_previous_mock(session: AsyncSession) -> None: result = await session.execute( select(Car.id).where(Car.plate_number.like(f"{MOCK_PLATE_PREFIX}-%")) ) car_ids = list(result.scalars()) if not car_ids: return await session.execute(delete(ServiceEntry).where(ServiceEntry.car_id.in_(car_ids))) await session.execute(delete(FuelEntry).where(FuelEntry.car_id.in_(car_ids))) await session.execute(delete(Car).where(Car.id.in_(car_ids))) async def seed_mock_usage(session: AsyncSession, owner: User) -> None: await clear_previous_mock(session) today = date.today() for index, (name, make, model, year, fuel_type, start_odo, price) in enumerate(MOCK_CARS, start=1): car = Car( owner_id=owner.id, name=name, make=make, model=model, year=year, plate_number=f"{MOCK_PLATE_PREFIX}-{index:02d}", fuel_type=fuel_type, purchase_date=date(year, min(index, 12), 10), purchase_price=price, current_odometer=start_odo, ) session.add(car) await session.flush() odometer = start_odo monthly_km = 820 + index * 95 consumption = Decimal("7.2") + Decimal(index % 5) * Decimal("0.45") if fuel_type == "electric": consumption = Decimal("0") for months_back in range(11, -1, -1): base_day = month_shift(today.replace(day=15), months_back) odometer += monthly_km if fuel_type == "electric": energy_cost = Decimal(110 + index * 8 + months_back % 3 * 15) session.add( ServiceEntry( car_id=car.id, entry_date=base_day, odometer=odometer, service_type=ServiceType.other, title="Зарядка и парковка", category="charging", vendor="EV network", total_cost=energy_cost, ) ) else: liters_per_month = Decimal(monthly_km) * consumption / Decimal(100) price_per_liter = Decimal("58.50") + Decimal(index % 4) * Decimal("2.10") for fill in range(2): fill_date = date(base_day.year, base_day.month, 8 if fill == 0 else 22) liters = (liters_per_month / 2).quantize(Decimal("0.001")) session.add( FuelEntry( car_id=car.id, entry_date=fill_date, odometer=odometer - (monthly_km // 2 if fill == 0 else 0), liters=liters, price_per_liter=price_per_liter, total_cost=(liters * price_per_liter).quantize(Decimal("0.01")), station=["Shell", "Lukoil", "Gazprom", "Rosneft", "Neste"][index % 5], is_full_tank=True, ) ) if months_back in {11, 5}: session.add( ServiceEntry( car_id=car.id, entry_date=date(base_day.year, base_day.month, 12), odometer=odometer, service_type=ServiceType.maintenance, title="Плановое ТО", category="regular", vendor="Service center", total_cost=Decimal(12500 + index * 950), next_due_odometer=odometer + 10000, ) ) if months_back in {8, 2}: session.add( ServiceEntry( car_id=car.id, entry_date=date(base_day.year, base_day.month, 18), odometer=odometer, service_type=ServiceType.tire, title="Сезонная замена шин", category="tires", vendor="Tire shop", total_cost=Decimal(4200 + index * 180), ) ) if months_back == 6: session.add( ServiceEntry( car_id=car.id, entry_date=date(base_day.year, base_day.month, 20), odometer=odometer, service_type=ServiceType.insurance, title="ОСАГО / страховка", category="insurance", vendor="Insurance", total_cost=Decimal(18200 + index * 1150), ) ) if months_back == index % 10: session.add( ServiceEntry( car_id=car.id, entry_date=date(base_day.year, base_day.month, 25), odometer=odometer, service_type=ServiceType.repair, title=["Тормозные колодки", "Диагностика подвески", "Замена АКБ", "Ремонт кондиционера"][index % 4], category="repair", vendor="Garage", total_cost=Decimal(7200 + index * 1350), ) ) car.current_odometer = odometer async def main() -> None: async with async_session_factory() as session: await seed_catalog(session) owner = await pick_owner(session) await seed_mock_usage(session, owner) await session.commit() print(f"Seeded catalog and mock usage for owner_id={owner.id}") if __name__ == "__main__": asyncio.run(main())