import importlib.util
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest

from k1link.observatory.modular_composition import CompositionError
from k1link.simulation.ai_polygon.composition import compose, registry
from k1link.simulation.ai_polygon.terrain_contract import terrain_matches

ROOT = Path(__file__).resolve().parents[1] / "simulation/ai-polygon"


def module(name):
    spec = importlib.util.spec_from_file_location(name, ROOT / (name + ".py"))
    result = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(result)
    return result


def test_simulation_composition_has_causal_dependencies_and_separate_authority():
    graph = compose(ROOT)
    assert graph.as_dict()["execution"]["mode"] == "worker-local-simulation"
    motion = graph.nodes[-1]
    assert motion.module.group == "motion"
    assert dict(motion.inputs)["segmentation.surface"] == "simulation-segformer-ade"
    assert dict(motion.inputs)["detection.boxes"] == "simulation-rf-detr"
    assert motion.module.state_policy == "causal-reset-at-source-start"
    selection = graph.selection_document()
    selection["selections"] = [r for r in selection["selections"] if r["group"] != "segmentation"]
    with pytest.raises(CompositionError, match="segmentation.surface"):
        compose(ROOT, selection)


def test_surface_providers_are_interchangeable_in_the_shared_constructor():
    selection = compose(ROOT).selection_document()
    reference = next(m for m in registry(ROOT).modules if m.module_id == "simulation-ddrnet-goose")
    for row in selection["selections"]:
        if row["group"] == "segmentation":
            row.update(module_id=reference.module_id, module_sha256=reference.sha256)
    graph = compose(ROOT, selection)
    assert dict(graph.nodes[-1].inputs)["segmentation.surface"] == reference.module_id


def test_composition_rejects_stale_module_identity():
    selection = compose(ROOT).selection_document()
    selection["selections"][0]["module_sha256"] = "0" * 64
    with pytest.raises(CompositionError, match="not installed"):
        compose(ROOT, selection)


def test_shared_constructor_import_needs_no_core_or_third_party_runtime():
    # -S removes site-packages, as in the minimal Windows coordinator. Loading
    # a contract must not load the POSIX-only artifact gateway through __init__.
    source = str(ROOT.parents[1] / "src")
    result = subprocess.run(
        [
            sys.executable,
            "-S",
            "-c",
            f"import sys; sys.path.insert(0, {source!r}); "
            "from k1link.observatory.modular_composition import ModuleRegistry; "
            "from k1link.simulation.ai_polygon.composition import compose; "
            "assert 'k1link.artifact_gateway' not in sys.modules",
        ],
        capture_output=True,
        text=True,
    )
    assert result.returncode == 0, result.stderr


def test_semantic_goal_uses_range_and_correct_square_camera_crop():
    nav = module("navigation_client")
    points = np.array(
        [[x, y, 0] for x in np.linspace(1.5, 3, 20) for y in np.linspace(-0.3, 0.3, 9)],
        dtype=np.float32,
    )
    pose = [0, 0, 0.27, 0, 0, 0, 1]
    calibration = {
        "origin": [0.38, 0, 0.8],
        "rotation": np.eye(3).reshape(-1).tolist(),
        "intrinsics": [800 * 24 / 36, 800 * 24 / 36, 400, 300],
    }
    leaves = np.ones((512, 512), dtype=bool)
    goal = nav.visual_goal(leaves, points, pose, calibration)
    assert goal is not None and 1.5 < goal[0] < 3 and abs(goal[1]) < 0.3
    assert nav.visual_goal(np.zeros_like(leaves), points, pose, calibration) is None
    assert nav.visual_goal(leaves, points + [0, 0, 2], pose, calibration) is None
    # A previously valid goal cannot authorize motion through newly unknown RGB.
    assert nav.visual_goal(np.zeros_like(leaves), points, pose, calibration, goal) is None
    # An explicit-route waypoint entering the camera blind strip is retained,
    # but losing all current visual surface support still forbids movement.
    close = [0.65, 0, 0]
    assert nav.visual_goal(leaves, points, pose, calibration, close, target=[0.65, 0]) == close
    assert (
        nav.visual_goal(np.zeros_like(leaves), points, pose, calibration, close, target=[0.65, 0])
        is None
    )


def test_ground_placement_uses_actual_triangle_intersection():
    terrain = module("terrain")
    vertices = np.array([[0, 0, 0], [1, 0, 0.2], [0, 1, 0]], dtype=np.float32)
    faces = np.array([[0, 1, 2]], dtype=np.int32)
    assert terrain.ground_intersections(vertices, faces, 0.25, 0.25)[0] == pytest.approx(0.05)
    assert len(terrain.ground_intersections(vertices, faces, 0.9, 0.9)) == 0


def test_observed_route_goal_keeps_task_position_and_cannot_run_away_from_it():
    choose = module("navigation_client").visual_goal
    points = np.array(
        [[x, y, 0] for x in np.arange(1.2, 3.1, 0.05) for y in np.arange(-0.5, 0.51, 0.05)]
    )
    pose = [0, 0, 0.37, 0, 0, 0, 1]
    calibration = {
        "origin": [0.38, 0, 0.8],
        "rotation": np.eye(3).reshape(-1).tolist(),
        "intrinsics": [800 * 24 / 36, 800 * 24 / 36, 400, 300],
        "body_contact_height_m": 0.37,
    }
    surface = np.ones((512, 512), dtype=bool)
    target = [2.03, 0.27]
    assert choose(surface, points, pose, calibration, target=target) == pytest.approx([*target, 0])
    # The close, already observed waypoint can enter the camera blind strip.
    advanced = [1.5, 0, 0.37, 0, 0, 0, 1]
    camera = {**calibration, "origin": [1.88, 0, 0.8]}
    prior = [*target, 0]
    assert choose(surface, points + [1.5, 0, 0], advanced, camera, prior, target) == prior
    # The actual camera loses nearby ground beyond the old hardcoded 0.8 m.
    advanced = [1.1, 0, 0.37, 0, 0, 0, 1]
    camera = {**calibration, "origin": [1.48, 0, 0.8]}
    assert choose(surface, points + [1.1, 0, 0], advanced, camera, prior, target) == prior
    assert choose(np.zeros_like(surface), points, pose, calibration, target=target) is None
    # Clear road ahead is not permission to drive away from a missed waypoint.
    assert choose(surface, points, pose, calibration, target=[-1, 0]) is None
    # A distant task may still use an observed local goal towards it.
    far = choose(surface, points, pose, calibration, target=[8, 0])
    assert far is not None and 1.2 <= far[0] <= 3.1


def test_collision_identity_follows_geometry_and_tile_coverage_not_camera_or_start():
    settings = dict(
        meters_per_unit=1,
        rotation_degrees=[-90, 0, 180],
        spawn_xy=[0, 0],
        ground_z=0,
        camera_height_m=0.8,
        max_speed_mps=0.15,
    )
    terrain = dict(source_sha256="a" * 64, generator_sha256="b" * 64, settings=settings)
    world = dict(sha256="a" * 64, settings={**settings, "spawn_xy": [1, 1], "camera_height_m": 1})
    assert terrain_matches(terrain, world, "b" * 64)
    assert not terrain_matches(terrain, world, "c" * 64)
    assert not terrain_matches(terrain, {**world, "sha256": "d" * 64})
    for change in ({"spawn_xy": [20, 0]}, {"meters_per_unit": 2}, {"rotation_degrees": [0, 0, 0]}):
        assert not terrain_matches(terrain, {**world, "settings": {**world["settings"], **change}})


def test_paired_full_scene_does_not_inherit_generated_tile_bounds():
    settings = dict(meters_per_unit=1, rotation_degrees=[90, 0, 0], spawn_xy=[0, 0], ground_z=5)
    terrain = dict(
        generator="paired-source",
        source_sha256="a" * 64,
        collider_sha256="b" * 64,
        settings=settings,
    )
    world = dict(
        sha256="a" * 64,
        collider_sha256="b" * 64,
        settings={**settings, "spawn_xy": [210, 30], "ground_z": 1.5},
    )
    assert terrain_matches(terrain, world)
    assert not terrain_matches(terrain, {**world, "collider_sha256": "c" * 64})
    assert not terrain_matches(
        terrain, {**world, "settings": {**world["settings"], "meters_per_unit": 2}}
    )
    assert not terrain_matches({**terrain, "generator": "generated-tile"}, world)


def test_square_footprint_fits_straight_corridor_and_rejects_corner_sweep():
    check = module("navigation/footprint").swept_footprint_clear
    pose = [0, 0, 0.27, 0, 0, 0, 1]
    path = [[0, 0, 0], [0.5, 0, 0], [1, 0, 0]]
    walls = np.array([[x, y, 0.5, 0.5] for x in np.arange(-1, 2, 0.1) for y in [-0.65, 0.65]])
    assert check(path, walls, pose)
    walls[:, 1] *= 0.45 / 0.65
    assert not check(path, walls, pose)
    # A diagonal turn sweeps a square corner into this obstacle, even though
    # the chassis at its initial and final straight poses does not contain it.
    assert not check([[0, 0, 0], [0.5, 0.5, 0]], [[0.7, 0, 0.5, 0.5]], pose)


def test_command_monitor_covers_deadman_braking_distance_and_rotation():
    check = module("navigation/footprint").command_footprint_clear
    pose = [0, 0, 0.27, 0, 0, 0, 1]
    assert check(0.15, 0, [[1.5, 0, 0.5, 0.5]], pose)
    assert not check(0.15, 0, [[0.7, 0, 0.5, 0.5]], pose)
    assert not check(0, 0.8, [[0.7, 0, 0.5, 0.5]], pose)


def test_smooth_slope_is_distinct_from_a_step_or_vertical_terrain():
    costs = module("navigation/terrain_costs").supported_slope_costs
    xy = np.array([[x, y] for x in np.arange(-1, 1.01, 0.1) for y in np.arange(-1, 1.01, 0.1)])
    slope = np.column_stack((xy, xy[:, 0] * np.tan(np.radians(20)), np.full(len(xy), 0.15)))
    normalized, corrected = costs(slope)
    assert corrected > len(slope) * 0.9
    assert normalized[len(slope) // 2, 3] == 0
    assert np.array_equal(normalized[:, :3], slope[:, :3])
    # A 15 cm ledge across the initial footprint cannot become a traversable ramp.
    step = slope.copy()
    step[:, 2] = np.where(step[:, 0] >= 0, 0.15, 0)
    assert costs(step)[1] == 0
    cliff = slope.copy()
    cliff[:, 2] = np.where(cliff[:, 0] >= 0, -0.4, 0)
    assert costs(cliff)[1] == 0
    steep = slope.copy()
    steep[:, 2] = steep[:, 0] * np.tan(np.radians(35))
    assert costs(steep)[1] == 0
    assert costs(slope[np.abs(slope[:, 1]) < 0.01])[1] == 0  # Unobserved lateral support.


def test_underbody_support_does_not_clear_future_terrain_walls_or_drops():
    correct = module("navigation/terrain_costs").underbody_support_costs
    terrain = np.array(
        [
            [-0.375, -0.28, 0.066, 0.103],  # Low return already under the chassis.
            [0.46, 0, 0.066, 0.103],  # Inset excludes the leading edge.
            [0.75, 0, 0.066, 0.103],  # Never change future terrain from body pose.
            [0, 0, 0.12, 0.12],  # A real step within the footprint remains blocked.
            [0, 0, 0.5, 0.5],
            [0, 0, -0.4, 0.4],
        ]
    )
    original = terrain.copy()
    result, count = correct(terrain, [0, 0, 0.37, 0, 0, 0, 1])
    assert count == 1 and result[0, 3] == pytest.approx(0.066)
    assert np.array_equal(result[1:], original[1:])
    assert np.array_equal(terrain, original)  # Never erase the causal raw map.
    # Rotate both observations and the measured chassis; the result must agree.
    yaw = np.pi / 2
    rotated = terrain.copy()
    rotated[:, :2] = terrain[:, :2] @ np.array([[0, 1], [-1, 0]]) + [3, 4]
    pose = [3, 4, 0.37, 0, 0, np.sin(yaw / 2), np.cos(yaw / 2)]
    assert np.allclose(correct(rotated, pose)[0][:, 3], result[:, 3])
    assert correct(terrain, [0, 0, 0.37, 0, np.sin(np.pi / 12), 0, np.cos(np.pi / 12)])[1] == 0


def test_retreat_requires_observed_full_width_support_and_no_drop_or_step():
    choose = module("navigation_client").recovery_goal
    points = np.array(
        [[x, y, 0.0] for x in np.arange(-1.5, -0.39, 0.05) for y in np.arange(-0.85, 0.86, 0.05)]
    )
    pose = [0, 0, 0.37, 0, 0, 0, 1]
    calibration = {"body_contact_height_m": 0.37}
    assert choose(points, pose, calibration) == pytest.approx([-0.65, 0, 0])
    assert choose(points[points[:, 1] > -0.2], pose, calibration) is None
    assert choose(points[points[:, 0] < -0.9], pose, calibration) is None
    for height in (-0.4, 0.15):
        discontinuous = points.copy()
        discontinuous[points[:, 0] < -0.9, 2] = height
        assert choose(discontinuous, pose, calibration) is None
    slope = points.copy()
    slope[:, 2] = slope[:, 0] * np.tan(np.radians(10))
    assert choose(slope, pose, calibration) is not None
    assert choose(points, pose, calibration, [-0.65, 0.3, 0]) is None


def test_reverse_monitor_checks_behind_the_body():
    check = module("navigation/footprint").command_footprint_clear
    pose = [0, 0, 0.37, 0, 0, 0, 1]
    assert check(-0.1, 0, [[-1.5, 0, 0.5, 0.5]], pose)
    assert not check(-0.1, 0, [[-0.65, 0, 0.5, 0.5]], pose)


def test_legacy_coordinate_migration_is_an_exact_rigid_rotation():
    import json
    import struct

    migrate = module("navigation/migrate_terrain_coordinates")
    positions = [[1.0, -2.0, -3.0], [2.0, -2.0, -3.0], [1.0, -1.0, -3.0]]
    binary = struct.pack("<9f3I", *(v for p in positions for v in p), 0, 1, 2)
    document = {
        "nodes": [{"mesh": 0}],
        "meshes": [{"primitives": [{"attributes": {"POSITION": 0}, "indices": 1}]}],
        "accessors": [
            {
                "bufferView": 0,
                "componentType": 5126,
                "type": "VEC3",
                "count": 3,
                "min": [1, -2, -3],
                "max": [2, -1, -3],
            },
            {"bufferView": 1, "componentType": 5125, "type": "SCALAR", "count": 3},
        ],
        "bufferViews": [{"byteOffset": 0, "byteLength": 36}, {"byteOffset": 36, "byteLength": 12}],
    }
    raw = json.dumps(document).encode()
    raw += b" " * ((-len(raw)) % 4)
    glb = (
        struct.pack("<III", 0x46546C67, 2, 28 + len(raw) + len(binary))
        + struct.pack("<II", len(raw), 0x4E4F534A)
        + raw
        + struct.pack("<II", len(binary), 0x004E4942)
        + binary
    )
    corrected = migrate.rotate_glb(glb)
    length = struct.unpack_from("<I", corrected, 12)[0]
    result = np.array(struct.unpack_from("<9f", corrected, 28 + length)).reshape(-1, 3)
    assert np.array_equal(result, np.array(positions) * [-1, 1, -1])
    assert np.array_equal(
        result[:, [0, 2, 1]] * [1, -1, 1], [[-1, -3, -2], [-2, -3, -2], [-1, -3, -1]]
    )
    assert struct.unpack_from("<3I", corrected, 28 + length + 36) == (0, 1, 2)


def test_voxel_quantized_grade_does_not_become_a_wall_but_ledge_remains():
    costs = module("navigation/terrain_costs").supported_slope_costs
    xy = np.array([[x, y] for x in np.arange(-1, 1.01, 0.06) for y in np.arange(-1, 1.01, 0.06)])
    height = np.round(xy[:, 0] * np.tan(np.radians(20)) / 0.06) * 0.06
    surface = np.column_stack((xy, height, np.full(len(xy), 0.15)))
    assert costs(surface)[1] > len(surface) * 0.8
    for discontinuity in (0.12, 0.15, -0.4):
        ledge = surface.copy()
        ledge[:, 2] = np.where(xy[:, 0] >= 0, discontinuity, 0)
        corrected, _ = costs(ledge)
        near_edge = np.abs(xy[:, 0]) < 0.12
        assert np.all(corrected[near_edge, 3] > 0.1)
    stone = surface.copy()
    stone[:, 2] = 0
    stone[(abs(xy[:, 0]) < 0.12) & (abs(xy[:, 1]) < 0.12), 2] = 0.15
    assert costs(stone)[1] == 0


def test_grade_fit_cannot_bridge_an_unobserved_gap():
    costs = module("navigation/terrain_costs").supported_slope_costs
    points = np.array(
        [
            [x, y, 0 if x < 0 else 0.15, 0.15]
            for x in [-0.3, -0.2, 0.2, 0.3]
            for y in np.arange(-0.3, 0.31, 0.1)
        ]
    )
    assert costs(points)[1] == 0


def test_existing_reserve_overlap_only_allows_departure_not_approach_or_body_overlap():
    check = module("navigation/footprint").command_footprint_clear
    pose = [0, 0, 0.37, 0, 0, 0, 1]
    behind = [[-0.53, 0, 0.5, 0.5]]
    assert check(0.15, 0, behind, pose)
    assert not check(-0.1, 0, behind, pose)
    assert not check(0, 0.35, behind, pose)
    assert not check(0.15, 0, [[-0.49, 0, 0.5, 0.5]], pose)
    assert not check(0.15, 0, [[0.53, 0, 0.5, 0.5]], pose)
    # A longer admitted camera age must also enlarge the collision envelope.
    assert not check(0.15, 0, [[0.80, 0, 0.5, 0.5]], pose)


def test_terrain_fit_cache_invalidates_when_a_new_obstacle_is_observed():
    costs = module("navigation/terrain_costs")
    normalize = costs.TerrainCostNormalizer()
    xy = np.array([[x, y] for x in np.arange(-1, 1.01, 0.1) for y in np.arange(-1, 1.01, 0.1)])
    grade = np.column_stack((xy, xy[:, 0] * np.tan(np.radians(20)), np.full(len(xy), 0.15)))
    clear, count = normalize(grade)
    assert count > len(grade) * 0.9
    assert np.array_equal(normalize(grade)[0], clear)
    changed = np.vstack((grade, [0.02, 0.02, 0.3, 0.3]))
    cached, count = normalize(changed)
    fresh, expected_count = costs.supported_slope_costs(changed)
    assert np.array_equal(cached, fresh) and count == expected_count
    assert cached[len(grade) // 2, 3] > 0.1
    assert np.array_equal(normalize(grade)[0], clear)


def test_velocity_regulation_keeps_the_selected_arc_and_its_braking_clearance():
    monitor = module("navigation/footprint")
    pose = [0, 0, 0.37, 0, 0, 0, 1]
    hazard = [[0.834, -0.08, 0.15, 0.15]]
    speed, yaw, scale = monitor.regulate_command(0.15, -0.245, hazard, pose)
    assert 0.2 <= scale < 1
    assert speed / yaw == pytest.approx(0.15 / -0.245)
    assert monitor.command_footprint_clear(speed, yaw, hazard, pose)
    assert monitor.regulate_command(0.15, 0, [[0.56, 0, 0.2, 0.2]], pose) == (0, 0, 0)
    assert monitor.regulate_command(0.15, 0, [[0.49, 0, 0.2, 0.2]], pose) == (0, 0, 0)
    reverse, _, scale = monitor.regulate_command(-0.1, 0, [[-0.68, 0, 0.2, 0.2]], pose)
    assert 0.2 <= scale < 1 and reverse < 0
    assert monitor.command_footprint_clear(reverse, 0, [[-0.68, 0, 0.2, 0.2]], pose)