from __future__ import annotations

from typing import Dict, List, Any

from .core import LocalArbProblem, SharedSignals, PlanDelta
from .governance import GraphOfContracts
from .bridge import EnergiBridge


class CentralCoordinator:
    """A lightweight, async-ADMM-like coordinator with deterministic replay.

    This MVP keeps a registry of latest SharedSignals per venue and merges inputs
    deterministically to produce PlanDelta actions.
    """

    def __init__(self) -> None:
        self.shared_signals_by_venue: Dict[str, SharedSignals] = {}
        self.version: int = 0
        self.contracts = GraphOfContracts()
        # Each element is a list of delta action dicts produced by a venue
        self.pending_delta_actions: List[List[Dict[str, Any]]] = []
        self.last_plan: PlanDelta | None = None
        # Optional bridge for canonical IR interoperability (can be used by adapters)
        self.bridge = EnergiBridge()
        self.last_catopt: Dict[str, Any] | None = None
        self.replay_log: List[str] = []

    def _other_venue(self, venue: str) -> str:
        venues = sorted(self.shared_signals_by_venue)
        if len(venues) < 2:
            return "VenueB" if venue != "VenueB" else "VenueA"
        for candidate in venues:
            if candidate != venue:
                return candidate
        return venue

    def ingest_local(self, risk_source: LocalArbProblem, signals: SharedSignals) -> PlanDelta | None:
        self.version += 1
        self.shared_signals_by_venue[risk_source.venue] = signals
        # Produce a canonical IR representation for interoperability
        try:
            self.last_catopt = self.bridge.to_catopt(risk_source, signals)
        except Exception:
            # Non-critical for MVP; retain existing behavior on bridge failure
            self.last_catopt = None
        delta = []
        for asset in sorted(signals.price_delta_by_asset):
            delta_price = signals.price_delta_by_asset[asset]
            if abs(delta_price) > 0.0005:
                action = {
                    "from_venue": risk_source.venue,
                    "to_venue": self._other_venue(risk_source.venue),
                    "asset": asset,
                    "size": min(risk_source.max_exposure, 10_000.0 * abs(delta_price)),
                    "time": float(self.version),
                }
                delta.append(action)
        if delta:
            contract_id = risk_source.fingerprint()
            plan = PlanDelta(
                delta_actions=delta,
                timestamp=float(self.version),
                contract_id=contract_id,
            )
            self.last_plan = plan
            self.pending_delta_actions.append(delta)
            self.contracts.register(contract_id, adapter_version=str(signals.version), schema_hash=plan.signature)
            self.replay_log.append(plan.signature)
            return plan
        return None

    def reconcile(self) -> PlanDelta | None:
        if not self.pending_delta_actions:
            return None
        merged = []
        for d in self.pending_delta_actions:
            merged.extend(d)
        merged.sort(key=lambda a: (a.get("from_venue"), a.get("to_venue"), a.get("asset")))
        contract_id = "merged:" + ",".join(
            sorted({action.get("asset", "") for action in merged})
        )
        plan = PlanDelta(
            delta_actions=merged,
            timestamp=float(self.version),
            contract_id=contract_id,
        )
        self.pending_delta_actions.clear()
        self.last_plan = plan
        self.contracts.register(contract_id, adapter_version="reconciled", schema_hash=plan.signature)
        self.replay_log.append(plan.signature)
        return plan