from .models import NoOp, DeployResource, RestrictMovement, CrisisworldcortexAction, CrisisworldcortexObservation, PublicCommunication
from .server import CrisisWorld as CrisisWorldEnvironment

def logger(count, obs: CrisisworldcortexObservation, action: CrisisworldcortexAction | None) -> None:
    if action:
        print(f"[Action Taken]: {action.model_dump(exclude_none=True)}")
    else:
        print("Initial Observation:")
    print(f"turn={count}, done={obs.done}, reward={obs.reward}")
    for reg in obs.regions:
        print(
            f"{reg.region_id} pop={reg.population} "
            f"inf={reg.infected} rec={reg.recovered} "
            f"dec={reg.deceased}"
        )
    print("-" * 40)
    print()

env = CrisisWorldEnvironment(config=EnvConfig(max_turns=6, num_regions=4))
obs = env.reset(seed=42, episode_id="manual-test")
logger(0, obs, None)

# Turn 1 (NoOp)
action = DeployResource(resource="medical", region_id="r0", amount=10)
# action = NoOp()
obs = env.step(action)
logger(1, obs, action)

# Turn 2 (NoOp)
action = NoOp()
obs = env.step(action)
logger(2, obs, action)

# Turn 3 (DeployResource to r0)
action = RestrictMovement(region_id="r0", level=2)
obs = env.step(action)
logger(3, obs, action)

# Turn 4 (PublicCommunication)
action = PublicCommunication(audience="public", message="Stay safe! Wear masks and maintain social distancing. Don't move between regions, stay home if you can, and seek medical care if you have symptoms.")
obs = env.step(action)
logger(4, obs, action)

# Turn 5 (NoOp)
action = DeployResource(resource="medical", region_id="r0", amount=40)
# action = NoOp()
obs = env.step(action)
logger(5, obs, action)