"""Arbitrage scanning and execution for cross-exchange trading.

Provides pipeline-compatible arb scanning and one-shot sweep functions
that build on the Engine's scan_arbitrage() and execute_arbitrage().

Usage:
    # Pipeline mode (auto-scan each cycle):
    hz.run(
        pipeline=[
            hz.arb_scanner("market_1", ["paper", "polymarket"],
                           feed_map={"paper": "paper_feed", "polymarket": "poly_feed"},
                           min_edge=0.02, auto_execute=True),
        ],
        ...
    )

    # One-shot mode:
    result = hz.arb_sweep(engine, "market_1", feed_map, min_edge=0.01)
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Callable

from horizon._horizon import Quote
from horizon.context import Context

logger = logging.getLogger(__name__)


@dataclass
class ArbResult:
    """Result of an arbitrage execution."""

    market_id: str
    buy_exchange: str
    sell_exchange: str
    buy_price: float
    sell_price: float
    size: float
    raw_edge: float
    net_edge: float
    buy_order_id: str = ""
    sell_order_id: str = ""
    timestamp: float = field(default_factory=time.time)


def arb_scanner(
    market_id: str,
    exchanges: list[str],
    feed_map: dict[str, str],
    min_edge: float = 0.01,
    max_size: float = 50.0,
    fee_rates: dict[str, float] | None = None,
    auto_execute: bool = False,
    cooldown: float = 5.0,
) -> Callable[[Context], list[Quote] | None]:
    """Create a pipeline function that scans for arbitrage opportunities.

    Args:
        market_id: Market to scan.
        exchanges: List of exchange names to compare.
        feed_map: Mapping of exchange_name -> feed_name.
        min_edge: Minimum net edge to consider actionable.
        max_size: Maximum trade size.
        fee_rates: Per-exchange fee rates {exchange_name: rate}. Default 0.002.
        auto_execute: If True, execute arb via engine.execute_arbitrage().
        cooldown: Seconds between executions.

    Returns:
        Pipeline function: (Context) -> list[Quote] | None
    """
    fees = fee_rates or {}
    last_exec_time = 0.0

    def _scanner(ctx: Context) -> list[Quote] | None:
        nonlocal last_exec_time

        engine = ctx.params.get("engine")
        if engine is None:
            return None

        # Build market_feed_map for scan_arbitrage
        market_feed_map = []
        for exch in exchanges:
            fname = feed_map.get(exch, exch)
            fee = fees.get(exch, 0.002)
            market_feed_map.append((fname, exch, fee))

        opps = engine.scan_arbitrage(market_id, market_feed_map, max_size)

        if not opps:
            return None

        # Filter by min_edge
        best = None
        for opp in opps:
            if opp.net_edge >= min_edge:
                if best is None or opp.net_edge > best.net_edge:
                    best = opp

        if best is None:
            return None

        # Auto-execute if enabled and cooldown elapsed
        if auto_execute:
            now = time.monotonic()
            if now - last_exec_time >= cooldown:
                try:
                    buy_id, sell_id = engine.execute_arbitrage(
                        market_id,
                        best.buy_exchange,
                        best.sell_exchange,
                        best.buy_price,
                        best.sell_price,
                        min(best.max_size, max_size),
                    )
                    last_exec_time = now
                    # Store result in params for downstream inspection
                    ctx.params["last_arb"] = ArbResult(
                        market_id=market_id,
                        buy_exchange=best.buy_exchange,
                        sell_exchange=best.sell_exchange,
                        buy_price=best.buy_price,
                        sell_price=best.sell_price,
                        size=min(best.max_size, max_size),
                        raw_edge=best.raw_edge,
                        net_edge=best.net_edge,
                        buy_order_id=buy_id,
                        sell_order_id=sell_id,
                    )
                except Exception as e:
                    logger.warning("Arb execution failed for %s: %s", market_id, e)

        return None  # Scanner doesn't generate quotes directly

    _scanner.__name__ = "arb_scanner"
    return _scanner


def arb_sweep(
    engine,
    market_id: str,
    feed_map: dict[str, str],
    min_edge: float = 0.01,
    max_size: float = 50.0,
    fee_rates: dict[str, float] | None = None,
) -> ArbResult | None:
    """One-shot arbitrage sweep: scan and execute the best opportunity.

    Args:
        engine: Horizon Engine instance.
        market_id: Market to scan.
        feed_map: Mapping of exchange_name -> feed_name.
        min_edge: Minimum net edge.
        max_size: Maximum trade size.
        fee_rates: Per-exchange fee rates.

    Returns:
        ArbResult if executed, None if no opportunity found.
    """
    fees = fee_rates or {}

    market_feed_map = []
    for exch, fname in feed_map.items():
        fee = fees.get(exch, 0.002)
        market_feed_map.append((fname, exch, fee))

    opps = engine.scan_arbitrage(market_id, market_feed_map, max_size)

    best = None
    for opp in opps:
        if opp.net_edge >= min_edge:
            if best is None or opp.net_edge > best.net_edge:
                best = opp

    if best is None:
        return None

    trade_size = min(best.max_size, max_size)

    try:
        buy_id, sell_id = engine.execute_arbitrage(
            market_id,
            best.buy_exchange,
            best.sell_exchange,
            best.buy_price,
            best.sell_price,
            trade_size,
        )
    except Exception as e:
        logger.warning("Arb sweep execution failed for %s: %s", market_id, e)
        return None

    return ArbResult(
        market_id=market_id,
        buy_exchange=best.buy_exchange,
        sell_exchange=best.sell_exchange,
        buy_price=best.buy_price,
        sell_price=best.sell_price,
        size=trade_size,
        raw_edge=best.raw_edge,
        net_edge=best.net_edge,
        buy_order_id=buy_id,
        sell_order_id=sell_id,
    )
