import pytest
from cosmosmesh_privacy_preserving_federated.energi_bridge import (
    LocalProblemEP,
    SharedVariableEP,
    DualVariableEP,
    PlanDeltaEP,
    EnergiBridge,
)


def test_energi_bridge_basic_encoding():
    lp = LocalProblemEP(
        problem_id="lp-001",
        assets=["rover-1", "drone-A"],
        objective="maximize_task_completion",
        constraints=["energy<=1000", "time_window=08:00-12:00"],
        data_contracts={"Telemetry": {"version": 1}},
    )

    sv1 = SharedVariableEP(channel="energy_budget", version=1, payload={"value": 42})
    dv1 = DualVariableEP(channel="lag_energy", version=1, payload={"value": 0.5})

    delta = PlanDeltaEP(delta_id="d-1", changes={"rover-1": {"move_to": "A"}}, timestamp=123456.0)

    bridge = EnergiBridge()
    ir = bridge.to_ir(lp, [sv1], [dv1], [delta])

    # Basic shape checks
    assert ir["Version"] == "0.1"
    assert "Objects" in ir
    assert ir["Objects"]["LocalProblem"]["problem_id"] == "lp-001"
    assert "Morphisms" in ir
    # We expect at least one SharedVariable and one DualVariable in the Morphisms list
    morphs = ir["Morphisms"]
    assert any("SharedVariable" in m.get("Morphisms", {}) for m in morphs) or any("DualVariable" in m.get("Morphisms", {}) for m in morphs)
    assert "PlanDeltas" in ir["Objects"]