#!/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()