X-Financial/server/scripts/seed_half_year_expense_demo.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import sys
from datetime import date
from pathlib import Path

from sqlalchemy import select

SERVER_DIR = Path(__file__).resolve().parents[1]
SRC_DIR = SERVER_DIR / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from app.db.session import get_session_factory  # noqa: E402
from app.models.employee import Employee  # noqa: E402
from app.services.demo_company_simulation_filters import is_admin_employee_like  # noqa: E402
from app.services.demo_company_simulation_seed import (  # noqa: E402
    HalfYearExpenseSimulationSeeder,
    SimulationConfig,
)
from app.services.employee_behavior_profile_service import (  # noqa: E402
    EmployeeBehaviorProfileService,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Seed half-year simulated reimbursement, budget, and employee data.",
    )
    parser.add_argument("--target-employees", type=int, default=100)
    parser.add_argument("--start-date", type=date.fromisoformat, default=date(2026, 1, 1))
    parser.add_argument("--end-date", type=date.fromisoformat, default=date(2026, 6, 2))
    parser.add_argument("--months", type=int, default=6)
    parser.add_argument("--seed", type=int, default=20260602)
    parser.add_argument("--apply", action="store_true", help="Write data. Default is dry-run only.")
    parser.add_argument(
        "--refresh-profiles",
        action="store_true",
        help="After --apply, refresh employee behavior profile snapshots for simulated employees.",
    )
    parser.add_argument("--profile-limit", type=int, default=120)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = SimulationConfig(
        target_employees=args.target_employees,
        start_date=args.start_date,
        end_date=args.end_date,
        months=args.months,
        seed=args.seed,
    )
    session_factory = get_session_factory()
    with session_factory() as db:
        seeder = HalfYearExpenseSimulationSeeder(db, config)
        try:
            summary = seeder.apply() if args.apply else seeder.preview()
            profile_refresh = None
            if args.apply and args.refresh_profiles:
                profile_refresh = refresh_company_profiles(db, limit=args.profile_limit)
            elif args.apply:
                db.commit()
            payload = summary.to_dict()
            if profile_refresh is not None:
                payload["profile_refresh"] = profile_refresh
            print(json.dumps(payload, ensure_ascii=False, indent=2))
            if not args.apply:
                print("dry-run only; pass --apply after confirmation to write simulated data.")
            elif not args.refresh_profiles:
                print("pass --refresh-profiles to generate employee behavior profile snapshots.")
        except Exception:
            db.rollback()
            raise


def refresh_company_profiles(db, *, limit: int) -> dict[str, object]:
    capped_limit = max(1, min(int(limit or 120), 500))
    employees = list(
        db.scalars(select(Employee).order_by(Employee.employee_no.asc())).all()
    )
    employee_ids = [
        employee.id
        for employee in employees
        if not is_admin_employee_like(employee)
    ][:capped_limit]
    service = EmployeeBehaviorProfileService(db)
    snapshot_count = 0
    for employee_id in employee_ids:
        snapshots = service.refresh_employee_profiles(
            employee_id=employee_id,
            window_days=(30, 90, 180),
            expense_type_scope="overall",
            source_task_type="half_year_expense_demo_seed",
            commit=False,
        )
        snapshot_count += len(snapshots)

    db.commit()
    return {
        "target_employee_count": len(employee_ids),
        "snapshot_count": snapshot_count,
        "window_days": [30, 90, 180],
        "source_task_type": "half_year_expense_demo_seed",
        "scope": "all_non_admin_employees",
    }


if __name__ == "__main__":
    main()