WonderQ-Project/WonderQ-Admin/app/seed.py

from __future__ import annotations

import argparse
import json
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from .auth import hash_password
from .content import DEMAND_FEATURE_CARDS, DEMAND_FORM, DEMAND_HERO, DESTINATION_HERO, HERO_SLIDES, VEHICLE_OPTIONS, VEHICLE_SERVICE
from .database import Base, SessionLocal, engine
from .models import (
    AdminUser,
    DemandFeatureCard,
    DemandForm,
    DemandHero,
    DestinationHero,
    HeroSlide,
    MediaAsset,
    SiteVersion,
    VehicleOption,
    VehicleServiceConfig,
    utc_now,
)
def create_media(db: Session, url: str | None, group: str, name: str | None = None) -> None:
    if not url:
        return
    media = db.scalar(select(MediaAsset).where(MediaAsset.url == url))
    if media is None:
        media = next((item for item in db.new if isinstance(item, MediaAsset) and item.url == url), None)
    if media:
        media.group = group
        media.name = name
    else:
        db.add(MediaAsset(url=url, group=group, name=name))


def reset_guizhou_content(db: Session) -> dict:
    for model in [SiteVersion, DemandForm, DemandFeatureCard, DemandHero, VehicleOption, VehicleServiceConfig, DestinationHero, HeroSlide, MediaAsset]:
        db.execute(delete(model))
    db.flush()

    for index, slide in enumerate(HERO_SLIDES):
        create_media(db, slide["image"], "hero", slide["title"])
        db.add(
            HeroSlide(
                title=slide["title"],
                kicker=slide["kicker"],
                actionLabel=slide["action"],
                image=slide["image"],
                sortOrder=index,
            )
        )

    for index, item in enumerate(DESTINATION_HERO):
        create_media(db, item.get("image"), "destination-hero", item["title"])
        db.add(
            DestinationHero(
                title=item["title"],
                kicker=item.get("kicker"),
                image=item.get("image"),
                sortOrder=index,
            )
        )

    for index, option in enumerate(VEHICLE_OPTIONS):
        create_media(db, option["image"], "vehicle-option", option["title"])
        db.add(
            VehicleOption(
                title=option["title"],
                description=option.get("description"),
                image=option.get("image"),
                sortOrder=index,
            )
        )

    db.add(VehicleServiceConfig(**VEHICLE_SERVICE))

    for index, item in enumerate(DEMAND_HERO):
        db.add(
            DemandHero(
                title=item["title"],
                kicker=item.get("kicker"),
                description=item.get("description"),
                steps=item.get("steps", []),
                sortOrder=index,
            )
        )

    for index, item in enumerate(DEMAND_FEATURE_CARDS):
        db.add(
            DemandFeatureCard(
                title=item["title"],
                description=item.get("description"),
                sortOrder=index,
            )
        )

    db.add(DemandForm(**DEMAND_FORM))

    snapshot = SiteVersion(
        title="guizhou-content-reset",
        status="published",
        publishedAt=utc_now(),
        snapshot={
            "heroSlides": len(HERO_SLIDES),
            "destinationHero": len(DESTINATION_HERO),
            "vehicleOptions": len(VEHICLE_OPTIONS),
            "vehicleService": 1,
            "demandHero": len(DEMAND_HERO),
            "demandFeatureCards": len(DEMAND_FEATURE_CARDS),
            "demandForm": 1,
        },
    )
    db.add(snapshot)
    db.flush()
    return {
        "heroSlides": len(HERO_SLIDES),
        "destinationHero": len(DESTINATION_HERO),
        "vehicleOptions": len(VEHICLE_OPTIONS),
        "vehicleService": 1,
        "demandHero": len(DEMAND_HERO),
        "demandFeatureCards": len(DEMAND_FEATURE_CARDS),
        "demandForm": 1,
        "siteVersionId": snapshot.id,
    }


def seed_database(reset: bool) -> dict:
    Base.metadata.create_all(bind=engine)
    with SessionLocal() as db:
        user = db.scalar(select(AdminUser).where(AdminUser.email == "admin@example.com"))
        if user:
            user.passwordHash = hash_password("ChangeMe123!")
            user.isActive = True
            user.name = "后台管理员"
            user.role = "super_admin"
        else:
            db.add(
                AdminUser(
                    email="admin@example.com",
                    name="后台管理员",
                    passwordHash=hash_password("ChangeMe123!"),
                    role="super_admin",
                )
            )
        result = reset_guizhou_content(db) if reset else {"reset": False}
        db.commit()
        return result


def main() -> None:
    parser = argparse.ArgumentParser(description="Seed WonderQ Admin database")
    parser.add_argument("--no-reset", action="store_true", help="Only ensure admin user exists")
    args = parser.parse_args()
    result = seed_database(reset=not args.no_reset)
    print(f"Seed complete. Admin login: admin@example.com / ChangeMe123!")
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()