from src.portfolio.runner import rebalance_allocations
from src.portfolio.ledger import PortfolioLedger


class FakeWorker:
    def __init__(self, cap, in_position=False):
        self.capital = cap
        self.in_position = in_position


def test_rebalance_resizes_flat_workers(tmp_path):
    L = PortfolioLedger("PX", total_capital=1000.0, data_dir=tmp_path)
    workers = {"a": FakeWorker(700.0), "b": FakeWorker(500.0)}   # equity 1200, both flat
    rebalance_allocations(L, workers, {"a": 0.5, "b": 0.5})
    assert L.total_capital == 1200.0
    assert workers["a"].capital == 600.0 and workers["b"].capital == 600.0


def test_rebalance_in_position_conserves_equity(tmp_path):
    # FIX 2026-06-13: il worker in posizione trattiene il suo capitale (deployato);
    # i flat si dividono il RESTO per peso rinormalizzato -> Σcapital CONSERVATO.
    # Prima (bug): i flat prendevano peso×total includendo il capitale dell'in-pos
    # -> somma gonfiata (qui sarebbe stata 1300, +100 fantasma).
    L = PortfolioLedger("PX", total_capital=1000.0, data_dir=tmp_path)
    workers = {"a": FakeWorker(700.0, in_position=True),
               "b": FakeWorker(300.0), "c": FakeWorker(200.0)}   # equity 1200
    rebalance_allocations(L, workers, {"a": 0.5, "b": 0.25, "c": 0.25})
    assert L.total_capital == 1200.0
    assert workers["a"].capital == 700.0       # invariato: posizione aperta (reserved)
    # b e c si dividono 1200-700=500 per peso 0.25/0.25 rinormalizzato -> 250 ciascuno
    assert workers["b"].capital == 250.0 and workers["c"].capital == 250.0
    total = sum(w.capital for w in workers.values())
    assert total == 1200.0                     # equity CONSERVATA dal ribilancio