import json

from policy_mesh.core import LocalPolicySet, GlobalConstraints, PolicyMesh


def test_basic_policymesh_flow():
    pm = PolicyMesh()

    # Create two local domains
    lp1 = LocalPolicySet(domain="district-a", policies={"energy_budget": 100.0, "ramp_rate": 1.0}, version="v0")
    lp2 = LocalPolicySet(domain="district-b", policies={"energy_budget": 80.0, "ramp_rate": 0.8}, version="v0")

    pm.add_local_policy(lp1)
    pm.add_local_policy(lp2)

    # Global constraint placeholder
    gc = GlobalConstraints(constraints={"mesh_budget": 180.0}, version="v0")
    pm.set_global_constraints(gc)

    # Run a couple of ADMM steps (mocked)
    pm.admm_step(rho=1.0)
    pm.admm_step(rho=1.0)

    state = pm.get_state()
    # basic assertions about stored state
    assert "local_problems" in state
    assert "district-a" in state["local_problems"]
    assert state["global_constraints"]["version"] == "v0" if state["global_constraints"] else True

    # ensure ledger has entries from admm steps
    ledger = pm.get_ledger()
    assert len(ledger) >= 2