import datetime as dt

import pandas as pd
import pytest

from ibkr_quant.strategy import select_with_buffer, sessions_until_rebalance

SCORE = pd.Series({"A": 0.9, "B": 0.8, "C": 0.7, "D": 0.6, "E": 0.5, "F": None})


def test_picks_top_n_equal_weight_without_holdings():
    assert select_with_buffer(SCORE, set(), top_n=2, buffer=0.0) == {"A": 0.5, "B": 0.5}


def test_keeps_held_names_within_the_buffer_and_fills_from_the_top():
    # top_n * (1 + buffer) = 4: D (rank 3) stays, E (rank 4) and unscored F go.
    assert select_with_buffer(SCORE, {"D", "E", "F"}, top_n=2, buffer=1.0) == {"D": 0.5, "A": 0.5}


def test_refuses_too_few_scored_names():
    with pytest.raises(ValueError, match="only 5 names"):
        select_with_buffer(SCORE, set(), top_n=6, buffer=0.0)


def test_rebalances_on_the_anchor_session_and_every_n_after():
    sessions = pd.bdate_range("2026-09-24", periods=12)  # anchor 09-29 is the 4th
    wait = [sessions_until_rebalance(sessions[: i + 1], dt.date(2026, 9, 29), 5) for i in range(len(sessions))]
    assert wait == [None, None, None, 0, 4, 3, 2, 1, 0, 4, 3, 2]