"""Tests for the lot- and cash-constrained rebalancer (Task 4)."""

from __future__ import annotations

import datetime as dt
import unittest

from app.portfolio_ledger import PortfolioLedger
from app.portfolio_rebalancer import (
    PortfolioRebalancer,
    RebalanceResult,
    build_candidates,
)
from app.simulation import Candidate


def make_series(symbols: list[str], start: str, days: int, price: float = 100.0) -> dict:
    """Flat daily price series for every symbol at a fixed price."""
    s = dt.date.fromisoformat(start)
    bars = [
        {"date": (s + dt.timedelta(days=i)).isoformat(), "adjusted_close": price}
        for i in range(days)
    ]
    return {sym: {"bars": list(bars)} for sym in symbols}


def scorer(*, score: float = 1.0, is_div: bool = True, yield_pct: float = 3.0):
    """Build a frozen score map {sym: meta} with the given attributes."""
    def _build(symbols: list[str]) -> dict:
        return {
            sym: {
                "combined": score, "is_dividend": is_div,
                "dividend_yield": yield_pct,
            }
            for sym in symbols
        }
    return _build


DATE = dt.date(2026, 1, 5)


class RebalancerTest(unittest.TestCase):
    def setUp(self):
        self.series = make_series(["A", "B", "C"], "2026-01-01", 30)

    def _rebalance(self, ledger, symbols, scores):
        cands = build_candidates(scores(symbols), self.series, DATE)
        rb = PortfolioRebalancer(ledger, self.series, cands)
        return rb.rebalance(date=DATE, signal_date=dt.date(2026, 1, 1)), cands

    def test_builds_initial_lot_positions(self):
        ledger = PortfolioLedger(1_000_000)
        score_fn = scorer(score=1.0, is_div=True, yield_pct=3.0)
        res, _ = self._rebalance(ledger, ["A", "B", "C"], score_fn)
        self.assertGreater(res.trade_count(), 0)
        for pos in ledger.positions():
            self.assertEqual(pos.qty % 100, 0)  # every position a 100-lot
        # equity reconciliation holds
        r = ledger.reconcile({"A": 100.0, "B": 100.0, "C": 100.0})
        self.assertTrue(r["balanced"])

    def test_unchanged_target_produces_no_trade(self):
        ledger = PortfolioLedger(1_000_000)
        score_fn = scorer(score=1.0, is_div=True, yield_pct=3.0)
        res1, cands = self._rebalance(ledger, ["A", "B"], score_fn)
        self.assertGreater(res1.trade_count(), 0)
        # same scores/regime -> target unchanged -> zero trades on re-rebalance
        rb = PortfolioRebalancer(ledger, self.series, cands)
        res2 = rb.rebalance(date=DATE, signal_date=dt.date(2026, 1, 1))
        self.assertEqual(res2.trade_count(), 0)
        self.assertFalse(res2.target_changed)

    def test_sale_profit_funds_next_purchase(self):
        # There must be enough proceeds from a profitable sale to afford a new
        # 100-lot, and the buy must actually happen.
        ledger = PortfolioLedger(1_000_000)
        # Buy A only at first: A dividend payer score 1
        score_fn = scorer(score=1.0, is_div=True)
        self._rebalance(ledger, ["A", "B"], score_fn)
        # Now target shifts to B (A exits). A is sold at same price -> no profit
        # here but proceeds fund B; test the buy occurs and reconciliation holds.
        score_b = scorer(score=2.0, is_div=True)  # B outranks A
        cands = build_candidates(
            {"B": {"combined": 2.0, "is_dividend": True, "dividend_yield": 3.0},
             "A": {"combined": 0.1, "is_dividend": True, "dividend_yield": 3.0}},
            self.series, DATE)
        rb = PortfolioRebalancer(ledger, self.series, cands)
        res = rb.rebalance(date=DATE, signal_date=dt.date(2026, 2, 1))
        self.assertGreater(res.trade_count(), 0)
        # B is held, in a 100-lot
        pos = ledger.position("B")
        assert pos is not None
        self.assertEqual(pos.qty % 100, 0)
        r = ledger.reconcile({"A": 100.0, "B": 100.0})
        self.assertTrue(r["balanced"])

    def test_cash_constraint_keeps_cash_and_skips_odd_lot(self):
        # capital that lets bucket 1 (50%) afford exactly 500 shares @100; the
        # cash-and-lot constraint must still hold and never go negative.
        ledger = PortfolioLedger(100_000)
        series = make_series(["A", "B"], "2026-01-01", 30, price=100.0)
        cands = build_candidates(
            {"A": {"combined": 1.0, "is_dividend": True, "dividend_yield": 3.0},
             "B": {"combined": 0.5, "is_dividend": False, "dividend_yield": 0.0}},
            series, DATE)
        rb = PortfolioRebalancer(ledger, series, cands)
        rb.rebalance(date=DATE, signal_date=dt.date(2026, 1, 1))
        # every position is a 100-lot, cash never negative
        for pos in ledger.positions():
            self.assertEqual(pos.qty % 100, 0)
        self.assertGreaterEqual(ledger.state.cash, 0)
        self.assertTrue(ledger.reconcile({"A": 100.0, "B": 100.0})["balanced"])

    def test_reconcile_after_paid_dividend_funds_next_buy(self):
        ledger = PortfolioLedger(1_000_000)
        score_fn = scorer(score=1.0, is_div=True)
        self._rebalance(ledger, ["A", "B"], score_fn)
        # record a dividend on A's holding, then pay it (ex+30)
        ledger.record_dividend_entitlement("A", "2026-01-10", 2.0)
        ledger.pay_due_dividends("2026-02-09")
        # dividend cash now in ledger; reconciliation stays balanced
        r = ledger.reconcile({"A": 100.0, "B": 100.0})
        self.assertTrue(r["balanced"])
        self.assertGreater(ledger.state.dividend_cash_received, 0)


if __name__ == "__main__":
    unittest.main()