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",
}
VARIATION_PREFIX = "AGR_GZ_EnvVar_"
EXPECTED_VARIATION_COUNT = 11
EXPECTED_VARIATION_MATERIALS = {
    "AGR_GZ_EnvVar_Tree_Canopy": "tree",
    "AGR_GZ_EnvVar_Tree_Trunk": "wood_resource",
    "AGR_GZ_EnvVar_Bush": "shrub",
    "AGR_GZ_EnvVar_Grass": "grass",
    "AGR_GZ_EnvVar_Rock": "stone_resource",
    "AGR_GZ_EnvVar_Water": "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 variation_material_key_for_label(label):
    for prefix, material_key in EXPECTED_VARIATION_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}")

    variation_actors = [actor for actor in actors if get_actor_label(actor).startswith(VARIATION_PREFIX)]
    if len(variation_actors) != EXPECTED_VARIATION_COUNT:
        failures.append(f"expected {EXPECTED_VARIATION_COUNT} environment variation actors, found {len(variation_actors)}")

    variation_meshes = set()
    variation_scales = set()
    variation_material_keys = set()
    for actor in variation_actors:
        label = get_actor_label(actor)
        expected_material_key = variation_material_key_for_label(label)
        if not expected_material_key:
            failures.append(f"{label} does not match a known variation material family")
            continue

        mesh_components = actor.get_components_by_class(unreal.StaticMeshComponent)
        if not mesh_components:
            failures.append(f"{label} has no static mesh component")
            continue

        component = mesh_components[0]
        mesh = component.get_editor_property("static_mesh")
        variation_meshes.add(mesh.get_path_name().split(".", 1)[0] if mesh else "")
        scale = actor.get_actor_scale3d()
        variation_scales.add((round(scale.x, 2), round(scale.y, 2), round(scale.z, 2)))
        variation_material_keys.add(expected_material_key)

        assigned = material_path(component.get_material(0))
        expected = MATERIALS[expected_material_key]
        if assigned != expected:
            failures.append(f"{label} material expected {expected}, got {assigned}")

    if len(variation_meshes) < 4:
        failures.append(f"expected at least 4 variation mesh silhouettes, got {len(variation_meshes)}")
    if len(variation_scales) < EXPECTED_VARIATION_COUNT:
        failures.append("environment variation actors should use unique scale profiles")
    if not {"tree", "wood_resource", "shrub", "grass", "stone_resource", "fresh_water"}.issubset(variation_material_keys):
        failures.append("environment variation actors do not cover tree, shrub, grass, resource, and water families")

    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, "asset variation"),
        (roadmap, "[x] Replace grey-box environment presentation with an MVP natural environment pass"),
        (roadmap, "[x] Add first-pass environment asset variation"),
    ]:
        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), "
        f"{checked_resource_actors} dressed resource/water actor(s), {len(variation_actors)} variation actor(s)."
    )


main()
