from __future__ import annotations

"""Toy two-venue ArbSphere demo.

This module wires two toy adapters (price feed and broker) from two venues
and demonstrates a minimal federated compute step using the existing ADMM-lite
coordinator API.
"""

import time
from typing import List
from arbsphere.primitives import LocalArbProblem, SharedSignals
from arbsphere.coordinator import admm_lite_step
from arbsphere.adapters.price_feed_adapter import PriceFeedAdapter
from arbsphere.adapters.broker_adapter import BrokerAdapter


def run_demo() -> None:
    # Define two toy venues with simple, deterministic budgets
    lp1 = LocalArbProblem(asset_pair=("AAPL", "MSFT"), target_mispricing=0.5, liquidity_budget=1000.0, latency_budget=5.0)
    lp2 = LocalArbProblem(asset_pair=("GOOG", "AMZN"), target_mispricing=0.7, liquidity_budget=800.0, latency_budget=6.0)

    shared = SharedSignals(
        deltas=[0.1, -0.05],
        cross_venue_corr=0.25,
        liquidity_availability={"venue-1": 0.8, "venue-2": 0.75},
        latency_proxy=0.9,
    )

    # Step 1: compute a delta (toy) via admm_lite_step
    delta = admm_lite_step([lp1, lp2], shared)

    print("Generated PlanDelta:")
    print(delta)

    # Step 2: route using toy broker adapters (two venues)
    broker_a = BrokerAdapter(venue_name="venue-1")
    broker_b = BrokerAdapter(venue_name="venue-2")

    # Create simple per-venue plan deltas; in a real system this would be split by venue
    plan_for_venue_a = {"legs": delta.legs, "total_size": delta.total_size}
    plan_for_venue_b = {"legs": delta.legs, "total_size": delta.total_size}

    broker_a.route(plan_for_venue_a)
    broker_b.route(plan_for_venue_b)

    print("Venue-1 execution history:", broker_a.history())
    print("Venue-2 execution history:", broker_b.history())


def build_toy_mvp() -> dict:
    """Bootstrap a minimal, test-friendly two-venue MVP scaffold.

    Returns a lightweight description of the MVP wiring, useful for tests
    or quick bootstrapping without running the full demo.
    """
    # Define two toy LocalArbProblems
    lp1 = LocalArbProblem(asset_pair=("AAPL", "MSFT"), target_mispricing=0.5, liquidity_budget=1000.0, latency_budget=5.0)
    lp2 = LocalArbProblem(asset_pair=("GOOG", "AMZN"), target_mispricing=0.7, liquidity_budget=800.0, latency_budget=6.0)

    shared = SharedSignals(
        deltas=[0.1, -0.05],
        cross_venue_corr=0.25,
        liquidity_availability={"venue-1": 0.8, "venue-2": 0.75},
        latency_proxy=0.9,
    )

    delta = admm_lite_step([lp1, lp2], shared)

    # Lightweight descriptor of the MVP wiring for external tooling
    return {
        "lp_pairs": [
            {"asset_pair": lp1.asset_pair, "target_mispricing": lp1.target_mispricing, "liquidity_budget": lp1.liquidity_budget, "latency_budget": lp1.latency_budget},
            {"asset_pair": lp2.asset_pair, "target_mispricing": lp2.target_mispricing, "liquidity_budget": lp2.liquidity_budget, "latency_budget": lp2.latency_budget},
        ],
        "shared_signals": {
            "deltas": list(shared.deltas),
            "cross_venue_corr": shared.cross_venue_corr,
            "liquidity_availability": dict(shared.liquidity_availability),
            "latency_proxy": shared.latency_proxy,
        },
        "plan_delta": {
            "legs": delta.legs,
            "total_size": delta.total_size,
            "delta_id": delta.delta_id,
            "timestamp": delta.timestamp,
        },
    }

def main() -> None:
    print("Starting ArbSphere two-venue demo...")
    run_demo()


if __name__ == "__main__":
    main()