from dataclasses import replace
import math

from okx_codex_trader.models import Candle
from scripts.run_bb_squeeze_t_gated_observer import EMPTY_STATE, frame_from_candles, signal_from_frame


def pair_frame(close_multiplier: float = 1.0, btc_last: float = 120.0):
    eth = []
    btc = []
    for index in range(1_200):
        eth_price = 100.0 + ((5.0 if index < 1_000 else 0.05) * math.sin(index / 3.0))
        open_price = eth_price
        high = eth_price
        low = eth_price
        close = eth_price
        if index == 1_199:
            close = 100.0 * close_multiplier
            high = max(open_price, close)
            low = min(open_price, close)
        btc_price = 100.0 + (index * 0.01)
        if index == 1_199:
            btc_price = btc_last
        ts = 1_700_000_000_000 + (index * 900_000)
        eth.append(Candle("ETH-USDT-SWAP", ts, open_price, high, low, close, 1_000.0))
        btc.append(Candle("BTC-USDT-SWAP", ts, btc_price, btc_price, btc_price, btc_price, 1_000.0))
    eth_frame = frame_from_candles(eth)
    btc_frame = frame_from_candles(btc)[["ts", "close"]].rename(columns={"close": "btc_close"})
    return eth_frame.merge(btc_frame, on="ts", how="inner").sort_values("ts").reset_index(drop=True)


def test_entry_long_requires_btc_up_filter() -> None:
    frame = pair_frame(close_multiplier=1.02, btc_last=120.0)

    next_state, signal = signal_from_frame(frame, EMPTY_STATE)

    assert signal["signal"] == "entry_long"
    assert signal["target_side"] == "long"
    assert next_state.active_side == "long"


def test_seen_candle_replays_without_new_signal() -> None:
    frame = pair_frame(close_multiplier=1.02, btc_last=120.0)
    state = replace(EMPTY_STATE, last_candle_ts=int(frame.iloc[-1]["ts"]))

    _, signal = signal_from_frame(frame, state)

    assert signal["signal"] == "state_replay"


def test_reentry_short_uses_btc_against_gate() -> None:
    frame = pair_frame(close_multiplier=1.0, btc_last=120.0)
    state = replace(
        EMPTY_STATE,
        last_candle_ts=int(frame.iloc[-2]["ts"]),
        reentry_side="short",
        reentry_anchor_price=99.0,
        reentry_until_ts=int(frame.iloc[-1]["ts"]) + 900_000,
    )

    next_state, signal = signal_from_frame(frame, state)

    assert signal["signal"] == "reentry_short"
    assert signal["target_side"] == "short"
    assert next_state.active_side == "short"