from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta

from drift_quant.backtesting.engine import BacktestConfig, run_backtest
from drift_quant.nse_data.loader import LoadedNseDataset, NseBar


@dataclass
class _SignalSequence:
    values: list[bool | None]
    name: str = "test_strategy"
    version: str = "test"
    _index: int = 0

    def target(self, symbol: str, bar: NseBar) -> bool | None:
        value = self.values[self._index]
        self._index += 1
        return value


def _bar(index: int, open_price: float, close: float) -> NseBar:
    return NseBar(
        timestamp=datetime(2024, 1, 1, 9, 15) + timedelta(minutes=15 * index),
        symbol="AAA",
        open=open_price,
        high=max(open_price, close),
        low=min(open_price, close),
        close=close,
        volume=100,
        source_row=index + 1,
    )


def test_orders_fill_at_next_bar_open_and_trade_ledger_reconciles() -> None:
    dataset = LoadedNseDataset(
        archive_path="fixture",
        archive_sha256="a" * 64,
        symbols=("AAA",),
        bars_by_symbol={"AAA": (_bar(0, 100, 100), _bar(1, 110, 110), _bar(2, 120, 120))},
        rows_scanned=3,
        rows_selected=3,
        invalid_rows_excluded=0,
    )
    result = run_backtest(
        dataset,
        _SignalSequence([True, False, None]),
        BacktestConfig(
            initial_capital=1_000,
            max_gross_exposure=1.0,
            transaction_cost_bps=0,
            slippage_bps=0,
        ),
    )

    assert len(result.trades) == 1
    trade = result.trades[0]
    assert trade.entry_price == 110
    assert trade.exit_price == 120
    assert trade.quantity == 9
    assert trade.net_pnl == 90
    assert result.metrics["final_equity"] == 1_090


def test_same_bar_close_is_not_used_as_entry_fill() -> None:
    dataset = LoadedNseDataset(
        archive_path="fixture",
        archive_sha256="a" * 64,
        symbols=("AAA",),
        bars_by_symbol={"AAA": (_bar(0, 100, 150), _bar(1, 200, 200))},
        rows_scanned=2,
        rows_selected=2,
        invalid_rows_excluded=0,
    )
    result = run_backtest(
        dataset,
        _SignalSequence([True, None]),
        BacktestConfig(initial_capital=1_000, transaction_cost_bps=0, slippage_bps=0),
    )

    assert result.trades[0].entry_price == 200