"""Rebalancing: the pure planning rules, then the refresh step end to end."""

from datetime import timedelta
from decimal import Decimal

import pytest
from sqlalchemy import select

from factories import make_account, make_event, make_instrument, make_price, refresh
from fintracker.analytics import FINDINGS, today_local
from fintracker.analytics.rebalance import (
    Position,
    Target,
    build_plan,
    compute_rebalance,
    rebuild_rebalance,
)
from fintracker.db import get_sessionmaker
from fintracker.models import (
    AccountKind,
    AccountRole,
    AllocationDimension,
    AssetClass,
    EventKind,
    Instrument,
    MetricAllocation,
    MetricRebalance,
    Portfolio,
    PortfolioAccount,
    PortfolioTarget,
)

D = Decimal
DIM = AllocationDimension


def position(
    instrument_id: int = 1,
    *,
    ticker: str = "SBER",
    lot: int = 10,
    qty: str = "100",
    unit: str = "275",
) -> Position:
    return Position(
        instrument_id=instrument_id,
        ticker=ticker,
        name=ticker,
        lot=lot,
        qty=D(qty),
        unit_value_rub=D(unit),
        price=D(unit),
        price_currency="RUB",
    )


def plan(
    *,
    bucket_values: dict[str, Decimal],
    positions: dict[str, list[Position]],
    targets: dict[str, Target],
    cash: str = "1000000",
):
    total = sum((v for v in bucket_values.values() if v > 0), start=D(0))
    return build_plan(
        portfolio_id=1,
        dimension=DIM.asset_class,
        as_of=today_local(),
        total_value_rub=total,
        bucket_values=bucket_values,
        positions=positions,
        targets=targets,
        cash_available_rub=D(cash),
    )


def bucket(result, name: str):
    return next(b for b in result.buckets if b.bucket == name)


# --------------------------------------------------------------------------- lots and cash


def test_a_buy_is_whole_lots_even_when_the_money_would_stretch_further():
    # 100 lots' worth of money, a lot of 10 at 275 => 2750 a lot
    result = plan(
        bucket_values={"share": D("27500"), "cash": D("22500")},
        positions={"share": [position(qty="100", unit="275", lot=10)]},
        targets={"share": Target(D("0.8")), "cash": Target(D("0.2"))},
        cash="22500",
    )
    trade = bucket(result, "share").trades[0]
    assert trade.action == "buy"
    # 0.8 * 50000 - 27500 = 12500 -> 45.45 units -> 4 lots = 40, never 45
    assert trade.qty == D(40)
    assert trade.qty % trade.lot == 0
    assert trade.amount_rub == D(40) * D("275")


def test_a_buy_is_cut_to_the_cash_on_hand_and_says_so():
    result = plan(
        bucket_values={"share": D("27500"), "cash": D("22500")},
        positions={"share": [position(qty="100", unit="275", lot=10)]},
        targets={"share": Target(D("0.8")), "cash": Target(D("0.2"))},
        cash="6000",
    )
    trade = bucket(result, "share").trades[0]
    # 6000 buys two lots (5500), not the 4 the target asks for
    assert trade.qty == D(20)
    assert trade.blocked_by_cash is True
    assert trade.amount_rub <= D("6000")


def test_no_cash_at_all_still_reports_the_blocked_buy_rather_than_hiding_it():
    result = plan(
        bucket_values={"share": D("27500"), "cash": D("22500")},
        positions={"share": [position(qty="100", unit="275", lot=10)]},
        targets={"share": Target(D("0.8")), "cash": Target(D("0.2"))},
        cash="0",
    )
    trade = bucket(result, "share").trades[0]
    assert trade.qty == D(0)
    assert trade.blocked_by_cash is True


def test_cash_is_spent_once_across_buckets():
    result = plan(
        bucket_values={"share": D("1000"), "bond": D("1000"), "cash": D("8000")},
        positions={
            "share": [position(1, ticker="SBER", qty="10", unit="100", lot=1)],
            "bond": [position(2, ticker="OFZ", qty="10", unit="100", lot=1)],
        },
        targets={"share": Target(D("0.45")), "bond": Target(D("0.45")), "cash": Target(D("0.1"))},
        cash="1000",
    )
    spent = sum(t.amount_rub for b in result.buckets for t in b.trades if t.action == "buy")
    assert spent <= D("1000")


# --------------------------------------------------------------------------- the band


def test_a_drift_inside_the_band_proposes_nothing():
    result = plan(
        bucket_values={"share": D("6200"), "bond": D("3800")},
        positions={"share": [position(qty="62", unit="100", lot=1)]},
        targets={"share": Target(D("0.6"), D("0.05")), "bond": Target(D("0.4"), D("0.05"))},
    )
    share = bucket(result, "share")
    assert share.drift == D("0.02")
    assert share.within_band is True
    assert share.trades == []
    assert share.delta_value_rub == D(0)


def test_the_same_drift_outside_the_band_proposes_a_trade():
    result = plan(
        bucket_values={"share": D("6200"), "bond": D("3800")},
        positions={
            "share": [position(qty="62", unit="100", lot=1)],
            "bond": [position(2, ticker="OFZ", qty="38", unit="100", lot=1)],
        },
        targets={"share": Target(D("0.6"), D("0.01")), "bond": Target(D("0.4"), D("0.01"))},
    )
    share = bucket(result, "share")
    assert share.within_band is False
    assert share.trades[0].action == "sell"
    assert share.trades[0].qty == D(2)


# --------------------------------------------------------------------------- sells


def test_a_sell_never_exceeds_the_position_and_never_goes_short():
    # the bucket must shrink by more than it holds: the target moved to zero
    result = plan(
        bucket_values={"share": D("1000"), "bond": D("9000")},
        positions={"share": [position(qty="10", unit="100", lot=1)]},
        targets={"share": Target(D("0")), "bond": Target(D("1"))},
    )
    trade = bucket(result, "share").trades[0]
    assert trade.action == "sell"
    assert trade.qty == D(10)
    assert trade.qty <= D(10)


def test_a_sell_is_capped_to_whole_lots_of_what_is_held():
    # 25 units of a 10-lot paper: at most two lots can be sold
    result = plan(
        bucket_values={"share": D("2500"), "bond": D("7500")},
        positions={"share": [position(qty="25", unit="100", lot=10)]},
        targets={"share": Target(D("0")), "bond": Target(D("1"))},
    )
    trade = bucket(result, "share").trades[0]
    assert trade.qty == D(20)


def test_a_bucket_is_trimmed_proportionally_not_from_one_paper():
    result = plan(
        bucket_values={"share": D("10000"), "bond": D("0")},
        positions={
            "share": [
                position(1, ticker="BIG", qty="75", unit="100", lot=1),
                position(2, ticker="SMALL", qty="25", unit="100", lot=1),
            ]
        },
        targets={"share": Target(D("0.5")), "bond": Target(D("0.5"))},
    )
    by_ticker = {t.ticker: t.qty for t in bucket(result, "share").trades}
    # 5000 to raise, split 75/25 by value: 37 and 12 units (floored to whole lots)
    assert by_ticker == {"BIG": D(37), "SMALL": D(12)}


def test_a_bucket_with_nothing_priced_in_it_warns_instead_of_inventing_a_trade():
    result = plan(
        bucket_values={"share": D("10000"), "bond": D("0")},
        positions={},
        targets={"share": Target(D("0.5")), "bond": Target(D("0.5"))},
    )
    assert bucket(result, "share").trades == []
    assert any("share" in w for w in result.warnings)


def test_the_cash_bucket_needs_no_trades_and_produces_no_warning():
    result = plan(
        bucket_values={"share": D("5000"), "cash": D("5000")},
        positions={"share": [position(qty="50", unit="100", lot=1)]},
        targets={"share": Target(D("0.9")), "cash": Target(D("0.1"))},
        cash="5000",
    )
    assert bucket(result, "cash").trades == []
    assert not any("cash" in w for w in result.warnings)


def test_a_bucket_without_a_target_is_reported_but_never_traded():
    result = plan(
        bucket_values={"share": D("5000"), "etf": D("5000")},
        positions={"etf": [position(2, ticker="TMOS", qty="50", unit="100", lot=1)]},
        targets={"share": Target(D("1"))},
    )
    etf = bucket(result, "etf")
    assert etf.target_weight is None
    assert etf.drift is None
    assert etf.trades == []


def test_every_number_in_the_plan_is_a_decimal():
    result = plan(
        bucket_values={"share": D("6200"), "bond": D("3800")},
        positions={"share": [position(qty="62", unit="100", lot=1)]},
        targets={"share": Target(D("0.5")), "bond": Target(D("0.5"))},
    )
    for b in result.buckets:
        for value in (b.current_value_rub, b.current_weight, b.delta_value_rub):
            assert isinstance(value, Decimal)
        for t in b.trades:
            for value in (t.qty, t.price, t.amount_rub):
                assert isinstance(value, Decimal)


# --------------------------------------------------------------------------- database


async def _portfolio_with(*, unpriced: bool) -> dict[str, int]:
    """A broker account in a portfolio: 500 SBER (lot 10), 20 OFZ, the rest in cash."""
    t = today_local()
    bought = t - timedelta(days=40)
    account = await make_account(
        name="Брокерский",
        kind=AccountKind.broker,
        role=AccountRole.investment,
        balance=None,
        include_in_net_worth=False,
        source="tinvest",
    )
    sber = await make_instrument(ticker="SBER", name="Сбербанк", asset_class=AssetClass.share)
    ofz = await make_instrument(ticker="OFZ", name="ОФЗ", asset_class=AssetClass.bond)

    async with get_sessionmaker()() as session:
        instrument = await session.get(Instrument, sber)
        assert instrument is not None
        instrument.lot = 10
        portfolio = Portfolio(name="Основной")
        session.add(portfolio)
        await session.flush()
        session.add(PortfolioAccount(portfolio_id=portfolio.id, account_id=account))
        portfolio_id = portfolio.id
        await session.commit()

    await make_event(bought, account_id=account, kind=EventKind.deposit, amount="100000")
    await make_event(
        bought,
        account_id=account,
        kind=EventKind.buy,
        instrument_id=sber,
        quantity="500",
        price="100",
        amount="-50000",
    )
    await make_event(
        bought,
        account_id=account,
        kind=EventKind.buy,
        instrument_id=ofz,
        quantity="20",
        price="1000",
        amount="-20000",
    )
    ids = {"account": account, "portfolio": portfolio_id, "sber": sber, "ofz": ofz}

    if unpriced:
        silent = await make_instrument(
            ticker="SIBN6P4", name="Без цены", asset_class=AssetClass.share, board="SPBRUBND"
        )
        await make_event(
            bought,
            account_id=account,
            kind=EventKind.buy,
            instrument_id=silent,
            quantity="5",
            price="1000",
            amount="-5000",
        )
        ids["silent"] = silent

    d = bought
    while d <= t:
        await make_price(d, instrument_id=sber, close="100")
        await make_price(d, instrument_id=ofz, close="1000")
        d += timedelta(days=1)
    await refresh()
    return ids


async def _set_targets(portfolio_id: int, rows: list[tuple[str, str, str]]) -> None:
    async with get_sessionmaker()() as session:
        for bucket_name, weight, band in rows:
            session.add(
                PortfolioTarget(
                    portfolio_id=portfolio_id,
                    dimension=DIM.asset_class,
                    bucket=bucket_name,
                    target_weight=D(weight),
                    band=D(band),
                )
            )
        await session.commit()


@pytest.fixture
async def portfolio(app) -> dict[str, int]:
    ids = await _portfolio_with(unpriced=False)
    await _set_targets(
        ids["portfolio"],
        [("share", "0.6", "0.01"), ("bond", "0.2", "0.01"), ("cash", "0.2", "0.01")],
    )
    return ids


async def test_the_step_fills_the_target_columns_of_metric_allocation(portfolio):
    async with get_sessionmaker()() as session:
        await rebuild_rebalance(session)
        await session.commit()
        rows = (
            (
                await session.execute(
                    select(MetricAllocation).where(
                        MetricAllocation.scope == f"portfolio:{portfolio['portfolio']}",
                        MetricAllocation.dimension == DIM.asset_class,
                    )
                )
            )
            .scalars()
            .all()
        )
    by_bucket = {r.bucket: r for r in rows}
    assert by_bucket["share"].target_weight == D("0.6")
    assert by_bucket["share"].weight == D("0.5")
    assert by_bucket["share"].drift == D("-0.1")
    assert by_bucket["cash"].target_weight == D("0.2")
    assert by_bucket["cash"].drift == D("0.1")


async def test_metric_rebalance_agrees_with_metric_allocation(portfolio):
    async with get_sessionmaker()() as session:
        await rebuild_rebalance(session)
        await session.commit()
        allocation = {
            r.bucket: r
            for r in (
                (
                    await session.execute(
                        select(MetricAllocation).where(
                            MetricAllocation.scope == f"portfolio:{portfolio['portfolio']}",
                            MetricAllocation.dimension == DIM.asset_class,
                        )
                    )
                )
                .scalars()
                .all()
            )
        }
        summaries = {
            r.bucket: r
            for r in (
                (
                    await session.execute(
                        select(MetricRebalance).where(MetricRebalance.instrument_id.is_(None))
                    )
                )
                .scalars()
                .all()
            )
        }
        trades = (
            (
                await session.execute(
                    select(MetricRebalance).where(MetricRebalance.instrument_id.is_not(None))
                )
            )
            .scalars()
            .all()
        )

    for name, row in summaries.items():
        assert row.current_weight == allocation[name].weight
        assert row.target_weight == allocation[name].target_weight
        assert row.current_value_rub == allocation[name].value_rub

    # 0.6 of 100 000 is 60 000 against 50 000 held: 100 more shares at 100, lot 10
    buy = next(t for t in trades if t.instrument_id == portfolio["sber"])
    assert buy.suggested_qty == D(100)
    assert buy.suggested_qty is not None
    assert buy.lot is not None
    assert buy.suggested_qty % buy.lot == 0
    assert buy.blocked_by_cash is False
    # the bond bucket sits exactly on its target and proposes nothing
    assert summaries["bond"].within_band is True
    assert not [t for t in trades if t.instrument_id == portfolio["ofz"]]


async def test_an_instrument_without_a_price_is_left_out_but_reported(app):
    ids = await _portfolio_with(unpriced=True)
    await _set_targets(
        ids["portfolio"],
        [("share", "0.6", "0.01"), ("bond", "0.2", "0.01"), ("cash", "0.2", "0.01")],
    )
    FINDINGS.reset()
    async with get_sessionmaker()() as session:
        await rebuild_rebalance(session)
        await session.commit()
        trades = (
            (
                await session.execute(
                    select(MetricRebalance).where(MetricRebalance.instrument_id.is_not(None))
                )
            )
            .scalars()
            .all()
        )
    assert ids["silent"] not in {t.instrument_id for t in trades}
    assert any(
        f.check_name == "rebalance_incomplete" and "SIBN6P4" in f.detail for f in FINDINGS.items
    )


async def test_the_what_if_cash_overrides_the_real_balance(portfolio):
    async with get_sessionmaker()() as session:
        real = await compute_rebalance(session, portfolio["portfolio"], DIM.asset_class)
        poor = await compute_rebalance(
            session, portfolio["portfolio"], DIM.asset_class, cash_available_rub=D("500")
        )
    assert real.cash_available_rub == D("30000")
    rich_trade = next(t for b in real.buckets for t in b.trades)
    poor_trade = next(t for b in poor.buckets for t in b.trades)
    assert poor_trade.qty < rich_trade.qty
    assert poor_trade.blocked_by_cash is True