"""Глаголы и профили запуска генератора.""" from __future__ import annotations import argparse import os import shlex import sys from dataclasses import dataclass from datetime import datetime, timedelta, timezone LAUNCH_ENV_KEYS = ( "GEN_SEED", "GEN_MODEL_T0", "GEN_MODEL_TIMEZONE", "GEN_MODEL_TIME_SPEED", "GEN_TICK_SECONDS", "GEN_LAMBDA_BASE_PER_MIN", "GEN_JITTER_PCT", "GEN_MIN_EVENTS_PER_TICK", "GEN_MAX_EVENTS_PER_TICK", ) @dataclass(frozen=True) class LaunchProfile: """Именованный набор настроек запуска.""" duration: str env: dict[str, str] PROFILES = { "ci": LaunchProfile( duration="6h", env={ "GEN_SEED": "4242", "GEN_MODEL_T0": "2026-01-01T00:00:00+00:00", "GEN_MODEL_TIMEZONE": "UTC", "GEN_MODEL_TIME_SPEED": "1", "GEN_TICK_SECONDS": "60", "GEN_LAMBDA_BASE_PER_MIN": "60", "GEN_JITTER_PCT": "0", "GEN_MIN_EVENTS_PER_TICK": "1", "GEN_MAX_EVENTS_PER_TICK": "1000", }, ), "daily-wave": LaunchProfile( duration="3d", env={ "GEN_SEED": "4242", "GEN_MODEL_T0": "2026-01-01T00:00:00+00:00", "GEN_MODEL_TIMEZONE": "UTC", "GEN_MODEL_TIME_SPEED": "60", "GEN_TICK_SECONDS": "1", "GEN_LAMBDA_BASE_PER_MIN": "60", "GEN_JITTER_PCT": "0", "GEN_MIN_EVENTS_PER_TICK": "1", "GEN_MAX_EVENTS_PER_TICK": "1000", }, ), } def parse_duration(value: str) -> timedelta: """Разбирает короткую длительность: 30m, 6h, 2d.""" if len(value) < 2: raise ValueError("Duration must look like 6h or 2d") amount_text = value[:-1] unit = value[-1] if not amount_text.isdigit(): raise ValueError("Duration amount must be a positive integer") amount = int(amount_text) if amount <= 0: raise ValueError("Duration amount must be > 0") if unit == "s": return timedelta(seconds=amount) if unit == "m": return timedelta(minutes=amount) if unit == "h": return timedelta(hours=amount) if unit == "d": return timedelta(days=amount) raise ValueError("Duration unit must be one of: s, m, h, d") def build_launch_env( verb: str, *, profile_name: str = "daily-wave", duration: str | None = None, overrides: dict[str, str] | None = None, ) -> dict[str, str]: """Возвращает старые env-переменные для нового глагола запуска.""" if verb not in {"backfill", "continue", "next-day", "reset"}: raise ValueError("глагол должен быть backfill, continue, next-day или reset") if profile_name not in PROFILES: known = ", ".join(sorted(PROFILES)) raise ValueError(f"Unknown launch profile {profile_name!r}; known: {known}") profile = PROFILES[profile_name] overrides = overrides or {} env = dict(profile.env) for key in LAUNCH_ENV_KEYS: if key in overrides and overrides[key] != "": env[key] = overrides[key] selected_duration = ( duration or overrides.get("GEN_HISTORY_DURATION") or profile.duration ) env["GEN_LAUNCH_VERB"] = verb env["GEN_LAUNCH_PROFILE"] = profile_name if verb == "backfill": env["GEN_RUN_MODE"] = "backfill" env["GEN_STATE_RESET"] = "true" env["GEN_HISTORY_DURATION"] = selected_duration env["GEN_MODEL_T_END"] = _model_t_end(env["GEN_MODEL_T0"], selected_duration) elif verb == "continue": env["GEN_RUN_MODE"] = "live" env["GEN_STATE_RESET"] = "false" env["GEN_HISTORY_DURATION"] = selected_duration env["GEN_MODEL_T_END"] = _model_t_end(env["GEN_MODEL_T0"], selected_duration) elif verb == "next-day": current_t_end = overrides.get("GEN_MODEL_T_END") if not current_t_end: raise ValueError("для next-day требуется GEN_MODEL_T_END") env["GEN_RUN_MODE"] = "next-day" env["GEN_STATE_RESET"] = "false" env["GEN_MODEL_T_END"] = current_t_end else: env["GEN_RUN_MODE"] = "live" env["GEN_STATE_RESET"] = "true" return env def shell_assignments(env: dict[str, str]) -> str: """Печатает env в виде, пригодном для eval в Bash.""" return "\n".join( f"export {key}={shlex.quote(value)}" for key, value in sorted(env.items()) ) def _model_t_end(model_t0: str, duration: str) -> str: t0 = datetime.fromisoformat(model_t0.replace("Z", "+00:00")) if t0.tzinfo is None: raise ValueError("GEN_MODEL_T0 must include timezone") return (t0.astimezone(timezone.utc) + parse_duration(duration)).isoformat() def _parse_args(argv: list[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Глаголы запуска генератора") parser.add_argument( "verb", choices=("backfill", "continue", "next-day", "reset"), ) parser.add_argument( "--profile", default=( os.getenv("PROFILE") or os.getenv("GEN_LAUNCH_PROFILE") or "daily-wave" ), choices=sorted(PROFILES), ) parser.add_argument("--duration", default=os.getenv("GEN_HISTORY_DURATION")) return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: """CLI для Bash-скриптов запуска.""" args = _parse_args(sys.argv[1:] if argv is None else argv) env = build_launch_env( args.verb, profile_name=args.profile, duration=args.duration, overrides=os.environ, ) print(shell_assignments(env)) return 0 if __name__ == "__main__": raise SystemExit(main())