from __future__ import annotations

from datetime import date, timedelta
import unittest

from dragon_bridge_predictive_break import (
    allow_predictive_b1_break_long_exit,
    allow_predictive_b1_break_short_exit,
    allow_predictive_error_reentry,
)
from dragon_strategy_config import StrategyConfig


class TestPredictiveBridgeModule(unittest.TestCase):
    def test_predictive_error_reentry_gate(self) -> None:
        d = date(2026, 4, 9)
        self.assertTrue(
            allow_predictive_error_reentry(
                enabled=True,
                last_exit_predictive_break=True,
                last_real_sell_date=d - timedelta(days=2),
                row_date=d,
                a1=-0.005,
                b1=-0.10,
                c1=55.0,
            )
        )
        self.assertFalse(
            allow_predictive_error_reentry(
                enabled=True,
                last_exit_predictive_break=False,
                last_real_sell_date=d - timedelta(days=2),
                row_date=d,
                a1=-0.005,
                b1=-0.10,
                c1=55.0,
            )
        )

    def test_predictive_short_exit_gate(self) -> None:
        cfg = StrategyConfig()
        self.assertTrue(
            allow_predictive_b1_break_short_exit(
                enabled=True,
                has_sell_signal=False,
                entry_is_glued=True,
                holding_days=1,
                a1=-0.01,
                b1=-0.14,
                c1=55.0,
                config=cfg,
            )
        )
        self.assertFalse(
            allow_predictive_b1_break_short_exit(
                enabled=True,
                has_sell_signal=True,
                entry_is_glued=True,
                holding_days=1,
                a1=-0.01,
                b1=-0.14,
                c1=55.0,
                config=cfg,
            )
        )

    def test_predictive_long_exit_gate(self) -> None:
        cfg = StrategyConfig()
        d = date(2026, 4, 9)
        self.assertTrue(
            allow_predictive_b1_break_long_exit(
                enabled=True,
                has_sell_signal=False,
                entry_is_glued=True,
                holding_days=45,
                max_c1_since_entry=70.0,
                max_a1_since_entry=0.2,
                max_b1_since_entry=0.35,
                last_ql_sell_date=d - timedelta(days=3),
                row_date=d,
                a1=-0.01,
                b1=-0.13,
                c1=62.0,
                config=cfg,
            )
        )
        self.assertFalse(
            allow_predictive_b1_break_long_exit(
                enabled=True,
                has_sell_signal=False,
                entry_is_glued=True,
                holding_days=45,
                max_c1_since_entry=70.0,
                max_a1_since_entry=0.2,
                max_b1_since_entry=0.35,
                last_ql_sell_date=d - timedelta(days=10),
                row_date=d,
                a1=-0.01,
                b1=-0.13,
                c1=62.0,
                config=cfg,
            )
        )


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