import unreal


MAP_PATH = "/Game/Agrarian/Maps/L_GroundZeroTerrain_Test"
FOLIAGE_LABEL = "AGR_GroundZeroFoliage_FirstPass"
EXPECTED_COUNTS = {
    "trees": 96,
    "shrubs": 220,
    "grass": 420,
}


def get_actor_label(actor):
    try:
        return actor.get_actor_label()
    except Exception:
        return actor.get_name()


def main():
    if not unreal.EditorLevelLibrary.load_level(MAP_PATH):
        raise RuntimeError(f"Could not load map: {MAP_PATH}")

    actors = unreal.EditorLevelLibrary.get_all_level_actors()
    foliage_actors = [actor for actor in actors if get_actor_label(actor) == FOLIAGE_LABEL]
    if len(foliage_actors) != 1:
        raise RuntimeError(f"Expected exactly one {FOLIAGE_LABEL}, found {len(foliage_actors)}")

    foliage = foliage_actors[0]
    actual_counts = {
        "trees": foliage.get_tree_instance_count(),
        "shrubs": foliage.get_shrub_instance_count(),
        "grass": foliage.get_grass_instance_count(),
    }

    failures = [
        f"{kind} expected {expected}, got {actual_counts[kind]}"
        for kind, expected in EXPECTED_COUNTS.items()
        if actual_counts[kind] != expected
    ]
    if failures:
        raise RuntimeError("Ground Zero foliage verification failed: " + "; ".join(failures))

    unreal.log(
        "Ground Zero foliage verification complete: "
        f"{actual_counts['trees']} trees, "
        f"{actual_counts['shrubs']} shrubs, "
        f"{actual_counts['grass']} grass clumps."
    )


main()