from typing import Dict, List, Optional
from .primitives import LocalArbProblem, SharedSignals, DualVariables, AuditLog, PrivacyBudget


def map_to_ir(lp: LocalArbProblem) -> Dict:
    # Simple IR mapping that preserves essential fields in a canonical form
    return {
        "type": "LocalArbProblem",
        "asset_pair": [lp.asset_pair[0], lp.asset_pair[1]],
        "target_mispricing": lp.target_mispricing,
        "liquidity_budget": lp.liquidity_budget,
        "latency_budget": lp.latency_budget,
    }


def map_to_full_ir(
    lp: LocalArbProblem,
    shared: Optional[SharedSignals] = None,
    dual: Optional[DualVariables] = None,
    privacy: Optional[PrivacyBudget] = None,
    audit: Optional[AuditLog] = None,
) -> Dict:
    """Map ArbSphere primitives to a richer, vendor-agnostic IR.

    All arguments are optional to allow incremental construction of the IR as data
    becomes available (e.g., after federation steps or governance events).
    """
    ir: Dict = map_to_ir(lp)
    # Attach optional components if provided
    if shared is not None:
        ir["shared_signals"] = {
            "deltas": list(shared.deltas),
            "cross_venue_corr": shared.cross_venue_corr,
            "liquidity_availability": dict(shared.liquidity_availability),
            "latency_proxy": shared.latency_proxy,
        }
    if dual is not None:
        ir["dual_variables"] = {"shadow_prices": dict(dual.shadow_prices)}
    if privacy is not None:
        ir["privacy_budget"] = {"budgets": dict(privacy.budgets)}
    if audit is not None:
        # Flatten audit log entries minimally for the IR; keep structure lightweight
        ir["audit_log"] = {
            "entries": [
                {
                    "ts": e.ts,
                    "event": e.event,
                    "details": dict(e.details),
                    "signature": e.signature,
                }
                for e in audit.entries
            ]
        }
    return ir