"""LeadLagStrategy — trade the follower when the leader moves.

For temporal relations where A leads B by N steps, when A makes a
significant move, B is expected to follow. We trade B in A's direction
before B catches up.

Entry: A moves by > entry_threshold from its recent mean.
  - A moved up → buy B (B will follow up)
  - A moved down → sell B / buy B's NO (B will follow down)
Exit: B has caught up (B's move matches A's), or timeout.

Usage:
    coinjure engine run \\
      --exchange polymarket --mode paper \\
      --strategy-ref coinjure/strategy/builtin/lead_lag_strategy.py:LeadLagStrategy \\
      --strategy-kwargs-json '{"relation_id": "610380-610379"}'
"""

from __future__ import annotations

import logging
from collections import deque
from decimal import Decimal

from coinjure.events import Event, PriceChangeEvent
from coinjure.strategy.relation_mixin import RelationArbMixin
from coinjure.strategy.strategy import Strategy
from coinjure.trading.sizing import compute_trade_size_with_llm
from coinjure.trading.trader import Trader
from coinjure.trading.types import TradeSide

logger = logging.getLogger(__name__)


class LeadLagStrategy(RelationArbMixin, Strategy):
    """Trade the follower when the leader moves.

    Parameters
    ----------
    relation_id:
        Relation ID from RelationStore. Must be a temporal type.
    trade_size:
        Max dollar amount per trade.
    entry_threshold:
        Minimum move in A (as fraction, e.g. 0.03 = 3 cents) to trigger.
    warmup:
        Number of price observations to build A's baseline before trading.
    exit_reversion:
        Fraction of A's move that B must match to trigger exit (0.0-1.0).
    max_hold:
        Maximum number of B price updates to hold before forced exit.
    kelly_fraction:
        Conservative Kelly multiplier for dynamic sizing.
    """

    name = 'lead_lag'
    version = '1.0.0'
    author = 'coinjure'

    def __init__(
        self,
        relation_id: str = '',
        trade_size: float = 10.0,
        entry_threshold: float = 0.03,
        warmup: int = 50,
        exit_reversion: float = 0.5,
        max_hold: int = 100,
        kelly_fraction: float = 0.1,
        llm_trade_sizing: bool = False,
        llm_model: str | None = None,
        llm_portfolio_review: bool = False,
    ) -> None:
        super().__init__()
        self.relation_id = relation_id
        self.max_trade_size = Decimal(str(trade_size))
        self.entry_threshold = Decimal(str(entry_threshold))
        self.kelly_fraction = Decimal(str(kelly_fraction))
        self._warmup_size = warmup
        self._exit_reversion = exit_reversion
        self._max_hold = max_hold
        self.kelly_fraction = Decimal(str(kelly_fraction))
        self.llm_trade_sizing = llm_trade_sizing
        self.llm_model = llm_model
        self.llm_portfolio_review = llm_portfolio_review

        self._init_from_relation(relation_id)

        # Determine which is leader and which is follower.
        id_a = self._ids[0] if self._ids else ''
        id_b = self._ids[1] if len(self._ids) > 1 else ''
        if self._relation:
            lag = self._relation.lead_lag or 0
            if lag >= 0:
                self._leader_slot = 0
                self._follower_slot = 1
                self._leader_id = id_a
                self._follower_id = id_b
            else:
                self._leader_slot = 1
                self._follower_slot = 0
                self._leader_id = id_b
                self._follower_id = id_a
        else:
            self._leader_slot = 0
            self._follower_slot = 1
            self._leader_id = id_a
            self._follower_id = id_b

        # Price tracking
        self._leader_prices: deque[float] = deque(maxlen=warmup)
        self._leader_price: Decimal | None = None
        self._follower_price: Decimal | None = None

        # Position state
        self._position_state = 'flat'  # flat | long_follower | short_follower
        self._entry_leader_price: float = 0.0
        self._entry_follower_price: float = 0.0
        self._hold_count = 0

    def reset_live_state(self) -> None:
        self._leader_price = None
        self._follower_price = None
        self._position_state = 'flat'
        self._entry_leader_price = 0.0
        self._entry_follower_price = 0.0
        self._hold_count = 0
        self._owned_symbols = set()

    async def process_event(self, event: Event, trader: Trader) -> None:
        if self.is_paused() or not isinstance(event, PriceChangeEvent):
            return

        ticker = event.ticker
        if ticker.side == 'no':
            return

        is_leader = self._slot_matches(ticker, self._leader_slot)
        is_follower = self._slot_matches(ticker, self._follower_slot)

        if is_leader:
            self._leader_price = event.price
            self._leader_prices.append(float(event.price))
        elif is_follower:
            self._follower_price = event.price
            if self._position_state != 'flat':
                self._hold_count += 1
        else:
            return

        # Need warmup for leader baseline
        if len(self._leader_prices) < self._warmup_size:
            return
        if self._leader_price is None or self._follower_price is None:
            return

        leader_mean = sum(self._leader_prices) / len(self._leader_prices)
        leader_move = float(self._leader_price) - leader_mean

        if self._position_state == 'flat':
            if is_leader and abs(leader_move) > float(self.entry_threshold):
                if leader_move > 0:
                    await self._enter_long(trader, leader_move)
                else:
                    await self._enter_short(trader, leader_move)
            elif is_leader:
                self.record_decision(
                    ticker_name=f'lag({self.relation_id[:20]})',
                    action='HOLD',
                    executed=False,
                    reasoning=(
                        f'leader_move={leader_move:.4f} '
                        f'< threshold={float(self.entry_threshold):.4f}'
                    ),
                    signal_values={
                        'leader': float(self._leader_price),
                        'follower': float(self._follower_price),
                        'leader_mean': leader_mean,
                        'leader_move': leader_move,
                    },
                )
        else:
            if is_follower:
                await self._check_exit(trader, leader_move)

    async def _enter_long(self, trader: Trader, leader_move: float) -> None:
        """Leader moved up → buy follower YES."""
        size = await compute_trade_size_with_llm(
            trader.position_manager,
            Decimal(str(abs(leader_move))),
            strategy_id=self.relation_id or self.name,
            strategy_type=self.name,
            relation_type='temporal',
            llm_trade_sizing=self.llm_trade_sizing,
            llm_model=self.llm_model,
            kelly_fraction=self.kelly_fraction,
            max_size=self.max_trade_size,
            leg_count=1,
            leg_prices=[self._follower_price or Decimal('0')],
        )
        ticker_b = self._find_ticker(trader, self._follower_id, side='yes')
        if not ticker_b or not self._follower_price:
            return

        result = await trader.place_order(
            side=TradeSide.BUY,
            ticker=ticker_b,
            limit_price=self._follower_price,
            quantity=size,
        )
        if result.failure_reason:
            return
        self._owned_symbols.add(ticker_b.symbol)

        self._position_state = 'long_follower'
        self._entry_leader_price = float(self._leader_price or 0)
        self._entry_follower_price = float(self._follower_price or 0)
        self._hold_count = 0

        self.record_decision(
            ticker_name=f'lag({self.relation_id[:20]})',
            action='BUY_FOLLOWER',
            executed=True,
            reasoning=f'Leader moved up {leader_move:.4f}, buying follower',
            signal_values={
                'leader': self._entry_leader_price,
                'follower': self._entry_follower_price,
                'leader_move': leader_move,
            },
        )
        logger.info('ENTER long follower: leader_move=%.4f', leader_move)

    async def _enter_short(self, trader: Trader, leader_move: float) -> None:
        """Leader moved down → sell follower (buy NO)."""
        size = await compute_trade_size_with_llm(
            trader.position_manager,
            Decimal(str(abs(leader_move))),
            strategy_id=self.relation_id or self.name,
            strategy_type=self.name,
            relation_type='temporal',
            llm_trade_sizing=self.llm_trade_sizing,
            llm_model=self.llm_model,
            kelly_fraction=self.kelly_fraction,
            max_size=self.max_trade_size,
            leg_count=1,
            leg_prices=[Decimal('1') - (self._follower_price or Decimal('0'))],
        )
        ticker_b_no = self._find_ticker(trader, self._follower_id, side='no')
        if not ticker_b_no or not self._follower_price:
            return

        no_price = Decimal('1') - self._follower_price
        result = await trader.place_order(
            side=TradeSide.BUY,
            ticker=ticker_b_no,
            limit_price=no_price,
            quantity=size,
        )
        if result.failure_reason:
            return
        self._owned_symbols.add(ticker_b_no.symbol)

        self._position_state = 'short_follower'
        self._entry_leader_price = float(self._leader_price or 0)
        self._entry_follower_price = float(self._follower_price or 0)
        self._hold_count = 0

        self.record_decision(
            ticker_name=f'lag({self.relation_id[:20]})',
            action='SELL_FOLLOWER',
            executed=True,
            reasoning=f'Leader moved down {leader_move:.4f}, selling follower',
            signal_values={
                'leader': self._entry_leader_price,
                'follower': self._entry_follower_price,
                'leader_move': leader_move,
            },
        )
        logger.info('ENTER short follower: leader_move=%.4f', leader_move)

    async def _check_exit(self, trader: Trader, leader_move: float) -> None:
        """Exit if follower caught up or max hold exceeded."""
        follower_move = float(self._follower_price or 0) - self._entry_follower_price
        leader_at_entry_move = self._entry_leader_price - float(
            sum(self._leader_prices) / len(self._leader_prices)
        )

        if abs(leader_at_entry_move) > 1e-6:
            catchup_ratio = follower_move / leader_at_entry_move
        else:
            catchup_ratio = 0.0

        should_exit = (
            catchup_ratio >= self._exit_reversion or self._hold_count >= self._max_hold
        )

        if should_exit:
            await self._close_owned(trader)

            reason = 'catchup' if catchup_ratio >= self._exit_reversion else 'timeout'
            self.record_decision(
                ticker_name=f'lag({self.relation_id[:20]})',
                action='EXIT',
                executed=True,
                reasoning=(
                    f'{reason}: follower_move={follower_move:.4f} '
                    f'catchup={catchup_ratio:.2f} hold={self._hold_count}'
                ),
                signal_values={
                    'follower_move': follower_move,
                    'catchup_ratio': catchup_ratio,
                    'hold_count': self._hold_count,
                },
            )
            logger.info(
                'EXIT %s: %s catchup=%.2f', self._position_state, reason, catchup_ratio
            )
            self._position_state = 'flat'
