import sys
from pathlib import Path

# Add framework src to path for imports
framework_src_path = Path(__file__).parent.parent.parent.parent / "src"
sys.path.insert(0, str(framework_src_path))

from layer2_simulation.single_runner import SingleRunner
from demos.school_demo.src.utils import print_header, print_initial_state, print_final_state


class SchoolSingleRunner(SingleRunner):
    """Runner for school simulations."""

    def __init__(self, world, param_dict) -> None:
        """Initialize SchoolSingleRunner.

        Args:
            world: The SchoolWorld instance (already created)
            param_dict: Simulation parameters
        """
        self.param_dict = param_dict
        super().__init__(world)

    def run_simulation(self, verbose: bool = True):
        """Run the school simulation.

        Args:
            verbose: Whether to print detailed output

        Returns:
            The SchoolWorld instance after simulation
        """
        max_steps = self.param_dict.get("max_steps", 1000)

        if verbose:
            print_header(f"RUNNING SIMULATION ({max_steps} steps)")
            print("\nSimulation in progress...")

        # Define stopping condition
        def stop_at_max_steps(schedule):
            return schedule.get_current_step() >= max_steps

        # Run the simulation using parent class method
        self.run(stopping_condition=stop_at_max_steps)

        if verbose:
            print_final_state(self.world)

        return self.world
