import time
from cosmosmesh_privacy_preserving_federated.cosmosmesh.models import LocalProblem, SharedVariables
from cosmosmesh_privacy_preserving_federated.cosmosmesh.delta_sync import reconcile_deltas
from cosmosmesh_privacy_preserving_federated.cosmosmesh.adapters import RoverPlannerAdapter, HabitatModuleAdapter


def test_models_basic_serialization():
    lp = LocalProblem(
        id="lp-1",
        domain="rover",
        assets=["rover-1"],
        objective="minimize_energy",
        constraints={"max_time": 3600},
        solver_hint="-quadratic",
        version=1,
    )

    sv = SharedVariables(version=1, forecasts={"energy": 100}, priors={"energy": 90})

    assert isinstance(lp.to_dict(), dict)
    assert isinstance(sv.to_dict(), dict)


def test_delta_sync_reconcile_basic():
    # Existing state has a delta for contract 'c1'
    existing = {
        "c1": {"delta": {"assignment": {"a1": "idle"}}, "timestamp": 1.0, "contract_id": "c1"},
    }
    new = [
        {"delta": {"assignment": {"a2": "patrol"}}, "timestamp": 2.0, "contract_id": "c1"},
        {"delta": {"assignment": {"a3": "scan"}}, "timestamp": 0.5, "contract_id": "c2"},
    ]
    updated = reconcile_deltas(existing, new)
    assert "c1" in updated
    assert updated["c1"]["delta"]["assignment"].get("a2") == "patrol"
    assert "c2" in updated


def test_adapters_mapping_exist():
    lp = LocalProblem(
        id="lp-2",
        domain="rover",
        assets=["rover-1", "rover-2"],
        objective="balance",
        constraints={},
    )
    rover = RoverPlannerAdapter()
    habitat = HabitatModuleAdapter()

    delta1 = rover.map_local_problem(lp)
    delta2 = habitat.map_local_problem(lp)

    assert isinstance(delta1, type(delta1))
    assert delta1.contract_id == lp.id
    assert isinstance(delta2, type(delta2))
    assert delta2.contract_id == lp.id