"""
Extended Counterparty Ecology — 10+ diverse adversarial agents.

Real markets have: market makers, HFT, institutional, retail, liquidators,
momentum traders, mean reversion traders, stale quote attackers, inventory MMs.
"""
from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Optional, Tuple

from malkhut.state import MarketWorldState, TradePathState, Side
from malkhut.actions import ActionKind, CounterpartyAction, AgentRole
from malkhut.counterparties import CounterpartyPolicy


@dataclass(frozen=True, slots=True)
class MomentumTakerPolicy:
    """Buys on upward momentum, sells on downward."""
    role: AgentRole = AgentRole.MOMENTUM_TAKER
    threshold: float = 0.3

    def legal_actions(self, state: MarketWorldState, params=None) -> Tuple[CounterpartyAction, ...]:
        return (
            CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.BUY, 0, 0.15, toxicity=0.4),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.15, toxicity=0.4),
        )

    def rollout_action(self, state: MarketWorldState, rng: random.Random) -> CounterpartyAction:
        path = state.trade_path
        momentum = path.pnl_bps if path else 0.0
        if abs(momentum) > self.threshold * 100:
            side = Side.BUY if momentum > 0 else Side.SELL
            return CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, side, 0, 0.15, toxicity=0.4)
        return CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0)


@dataclass(frozen=True, slots=True)
class MeanReversionTakerPolicy:
    """Buys on downward moves, sells on upward moves (mean reversion)."""
    role: AgentRole = AgentRole.MEAN_REVERSION_TAKER
    threshold: float = 0.5

    def legal_actions(self, state: MarketWorldState, params=None) -> Tuple[CounterpartyAction, ...]:
        return (
            CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.BUY, 0, 0.1, toxicity=0.3),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.1, toxicity=0.3),
        )

    def rollout_action(self, state: MarketWorldState, rng: random.Random) -> CounterpartyAction:
        path = state.trade_path
        momentum = path.pnl_bps if path else 0.0
        if abs(momentum) > self.threshold * 100:
            # Mean reversion: buy when price dropped, sell when price rose
            side = Side.SELL if momentum > 0 else Side.BUY
            return CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, side, 0, 0.1, toxicity=0.3)
        return CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0)


@dataclass(frozen=True, slots=True)
class InventoryMarketMakerPolicy:
    """Market maker that manages inventory levels."""
    role: AgentRole = AgentRole.INVENTORY_MM
    target_inventory: float = 0.0
    max_inventory: float = 0.1

    def legal_actions(self, state: MarketWorldState, params=None) -> Tuple[CounterpartyAction, ...]:
        return (
            CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0),
            CounterpartyAction(self.role, ActionKind.PLACE, Side.BUY, 0, 0.2),
            CounterpartyAction(self.role, ActionKind.PLACE, Side.SELL, 0, 0.2),
            CounterpartyAction(self.role, ActionKind.CANCEL, None, 0, 0.0),
        )

    def rollout_action(self, state: MarketWorldState, rng: random.Random) -> CounterpartyAction:
        pos = state.account.positions.get(state.venue.symbol)
        inv = pos.qty if pos else 0.0
        if inv > self.max_inventory:
            return CounterpartyAction(self.role, ActionKind.PLACE, Side.SELL, 0, 0.2)
        elif inv < -self.max_inventory:
            return CounterpartyAction(self.role, ActionKind.PLACE, Side.BUY, 0, 0.2)
        return CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0)


@dataclass(frozen=True, slots=True)
class LiquidationFlowPolicy:
    """Simulates forced liquidation during price drops."""
    role: AgentRole = AgentRole.LIQUIDATION_FLOW
    trigger_bps: float = 50.0  # LOWER threshold = more aggressive

    def legal_actions(self, state: MarketWorldState, params=None) -> Tuple[CounterpartyAction, ...]:
        return (
            CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.3, toxicity=0.9),
        )

    def rollout_action(self, state: MarketWorldState, rng: random.Random) -> CounterpartyAction:
        path = state.trade_path
        if path and path.mae_bps < -self.trigger_bps:
            return CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.3, toxicity=0.9)
        return CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0)


@dataclass(frozen=True, slots=True)
class StaleQuoteAttackerPolicy:
    """Attacks stale quotes that haven't been updated."""
    role: AgentRole = AgentRole.STALE_QUOTE_ATTACKER
    stale_threshold_s: float = 5.0

    def legal_actions(self, state: MarketWorldState, params=None) -> Tuple[CounterpartyAction, ...]:
        return (
            CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.BUY, 0, 0.2, toxicity=0.7),
            CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.2, toxicity=0.7),
        )

    def rollout_action(self, state: MarketWorldState, rng: random.Random) -> CounterpartyAction:
        # Attack if there are open orders that look stale
        if state.open_orders:
            return CounterpartyAction(self.role, ActionKind.CROSS_SPREAD, Side.SELL, 0, 0.2, toxicity=0.7)
        return CounterpartyAction(self.role, ActionKind.NOOP, None, 0, 0.0)


def extended_counterparty_ecology() -> Tuple[CounterpartyPolicy, ...]:
    """Full ecology with 10 diverse agents."""
    from malkhut.counterparties import (
        ToxicTakerPolicy, PassiveMakerPolicy, LatencyArbPolicy, NoiseTraderPolicy,
    )
    return (
        ToxicTakerPolicy(),
        PassiveMakerPolicy(),
        LatencyArbPolicy(),
        NoiseTraderPolicy(),
        MomentumTakerPolicy(),
        MeanReversionTakerPolicy(),
        InventoryMarketMakerPolicy(),
        LiquidationFlowPolicy(),
        StaleQuoteAttackerPolicy(),
    )