import unreal


MAP_PATH = "/Game/Agrarian/Maps/L_GroundZeroTerrain_Test"
FOLIAGE_LABEL = "AGR_GroundZeroFoliage_FirstPass"
MATERIALS = {
    "terrain": "/Game/Agrarian/Materials/M_AGR_GZ_Terrain_CoastalScrub",
    "tree": "/Game/Agrarian/Materials/M_AGR_GZ_Tree_CoastalOak",
    "shrub": "/Game/Agrarian/Materials/M_AGR_GZ_Shrub_CoyoteBrush",
    "grass": "/Game/Agrarian/Materials/M_AGR_GZ_Grass_DryCoastal",
    "wood_resource": "/Game/Agrarian/Materials/M_AGR_GZ_Wood_Resource",
    "fiber_resource": "/Game/Agrarian/Materials/M_AGR_GZ_Fiber_Resource",
    "stone_resource": "/Game/Agrarian/Materials/M_AGR_GZ_Stone_Sandstone",
    "fresh_water": "/Game/Agrarian/Materials/M_AGR_GZ_FreshWater",
}
EXPECTED_FOLIAGE_COUNTS = {
    "trees": 42,
    "shrubs": 96,
    "grass": 180,
}
RESOURCE_MATERIALS = {
    "AGR_GZ_Wood": "wood_resource",
    "AGR_GZ_Fiber": "fiber_resource",
    "AGR_GZ_Stone": "stone_resource",
    "AGR_DemoWoodResource": "wood_resource",
    "AGR_DemoFiberResource": "fiber_resource",
    "AGR_GZ_FreshWaterSource": "fresh_water",
}


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


def material_path(material):
    if not material:
        return ""
    return material.get_path_name().split(".", 1)[0]


def material_key_for_label(label):
    for prefix, material_key in RESOURCE_MATERIALS.items():
        if label.startswith(prefix):
            return material_key
    return None


def assert_asset(path):
    asset = unreal.EditorAssetLibrary.load_asset(path)
    if not asset:
        raise RuntimeError(f"Missing required environment material: {path}")
    return asset


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

    materials = {key: assert_asset(path) for key, path in MATERIALS.items()}
    actors = unreal.EditorLevelLibrary.get_all_level_actors()
    failures = []

    landscapes = [actor for actor in actors if isinstance(actor, unreal.Landscape)]
    if not landscapes:
        failures.append("no landscape actor found")
    for landscape in landscapes:
        assigned = material_path(landscape.get_editor_property("landscape_material"))
        expected = MATERIALS["terrain"]
        if assigned != expected:
            failures.append(f"landscape material expected {expected}, got {assigned}")

    foliage_actors = [actor for actor in actors if get_actor_label(actor) == FOLIAGE_LABEL]
    if len(foliage_actors) != 1:
        failures.append(f"expected one {FOLIAGE_LABEL}, found {len(foliage_actors)}")
    else:
        foliage = foliage_actors[0]
        counts = {
            "trees": foliage.get_tree_instance_count(),
            "shrubs": foliage.get_shrub_instance_count(),
            "grass": foliage.get_grass_instance_count(),
        }
        for key, expected_count in EXPECTED_FOLIAGE_COUNTS.items():
            if counts[key] != expected_count:
                failures.append(f"{key} expected {expected_count}, got {counts[key]}")

        component_expectations = {
            "tree_instances": "tree",
            "shrub_instances": "shrub",
            "grass_instances": "grass",
        }
        for property_name, material_key in component_expectations.items():
            component = foliage.get_editor_property(property_name)
            assigned = material_path(component.get_material(0))
            expected = MATERIALS[material_key]
            if assigned != expected:
                failures.append(f"{property_name} material expected {expected}, got {assigned}")

    checked_resource_actors = 0
    for actor in actors:
        label = get_actor_label(actor)
        material_key = material_key_for_label(label)
        if not material_key:
            continue

        checked_resource_actors += 1
        expected = MATERIALS[material_key]
        mesh_components = actor.get_components_by_class(unreal.StaticMeshComponent)
        if not mesh_components:
            failures.append(f"{label} has no static mesh component for material assignment")
            continue
        assigned_any = any(material_path(component.get_material(0)) == expected for component in mesh_components)
        if not assigned_any:
            failures.append(f"{label} did not use expected material {expected}")

    if checked_resource_actors < 10:
        failures.append(f"expected at least 10 dressed resource/water actors, checked {checked_resource_actors}")

    docs = unreal.Paths.convert_relative_path_to_full(unreal.Paths.project_dir()) + "Docs/Terrain/GroundZeroNaturalEnvironmentPass.md"
    roadmap = unreal.Paths.convert_relative_path_to_full(unreal.Paths.project_dir()) + "AGRARIAN_DEVELOPMENT_ROADMAP.md"
    for path, snippet in [
        (docs, "MVP natural environment pass"),
        (roadmap, "[x] Replace grey-box environment presentation with an MVP natural environment pass"),
    ]:
        with open(path, "r", encoding="utf-8") as handle:
            text = handle.read()
        if snippet not in text:
            failures.append(f"{path} missing {snippet}")

    if failures:
        raise RuntimeError("Ground Zero natural environment verification failed: " + "; ".join(failures))

    unreal.log(
        "Ground Zero natural environment verification complete: "
        f"{len(materials)} materials, {len(landscapes)} landscape(s), {checked_resource_actors} dressed resource/water actor(s)."
    )


main()