"""Tiny ADMM-lite solver for AidMesh.
This is a toy, deterministic solver to illustrate cross-organization coordination.
"""
from __future__ import annotations

from typing import Dict
import time

from .core import LocalProblem, SharedSignals, DualVariables, PlanDelta


def admm_step(
    contract_id: str,
    local_problem: LocalProblem,
    shared: SharedSignals,
    dual: DualVariables,
    rho: float = 1.0,
) -> PlanDelta:
    """Perform a single, toy ADMM step and return a PlanDelta.

    The toy logic simply pushes a delta that moves assets towards the forecasted needs
    weighted by current multipliers. This is not production-grade optimization.
    """
    # Simple heuristic: adjust each asset towards forecast needs, bounded by available capacity
    delta: Dict[str, float] = {}
    for asset, need in shared.forecast.items():
        current = local_problem.assets.get(asset, 0.0)
        cap = shared.capacity_proxies.get(asset, 0.0)
        # target: min(need, cap) but respect current
        target = max(0.0, min(need - current, cap))
        # apply a small portion dictated by rho and a multiplier if available
        mult = dual.multipliers.get(asset, 1.0) if isinstance(dual, DualVariables) else 1.0
        delta[asset] = target * (0.5 * rho) * mult

    # Build a PlanDelta with a simplistic signature and timestamp
    pd = PlanDelta(
        delta=delta,
        timestamp=time.time(),
        author="admm-lite-solver",
        contract_id=contract_id,
        signature="toy-signature",
    )
    return pd