polyguard-rl/app/env/scenario_loader.py

"""Scenario loading utilities."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

from app.common.enums import Difficulty
from app.common.types import PatientProfile
from app.simulator.scenario_generator import generate_patient_scenario


def _scenario_path(root: Path, difficulty: Difficulty, scenario_id: str) -> Path:
    return root / "data" / "scenarios" / difficulty.value / f"{scenario_id}.json"


def load_or_generate_scenario(
    root: Path,
    difficulty: Difficulty,
    scenario_id: Optional[str],
    patient_id: Optional[str],
    seed: int,
) -> PatientProfile:
    if scenario_id:
        path = _scenario_path(root, difficulty, scenario_id)
        if path.exists():
            payload = json.loads(path.read_text(encoding="utf-8"))
            return PatientProfile.model_validate(payload)
    return generate_patient_scenario(difficulty=difficulty, patient_id=patient_id, seed=seed)