import unittest

import pandas as pd

import trader
from strategies.permanent import (
    ETF_UNIVERSE,
    GLOBAL_ETF_UNIVERSE,
    HK_ETF_UNIVERSE,
    TREND_RIDER_V4_UNIVERSE,
    TrendRiderV3,
    TrendRiderV4,
)


class TrendRiderTraderIntegrationTests(unittest.TestCase):
    def test_trend_rider_strategies_are_registered(self):
        self.assertIsInstance(trader.STRATEGY_REGISTRY["trend_rider_v3_us"](), TrendRiderV3)
        self.assertIsInstance(trader.STRATEGY_REGISTRY["trend_rider_v3_global"](), TrendRiderV3)
        self.assertIsInstance(trader.STRATEGY_REGISTRY["trend_rider_v3_hk"](), TrendRiderV3)
        self.assertIsInstance(trader.STRATEGY_REGISTRY["trend_rider_v4"](), TrendRiderV4)

    def test_strategy_universe_uses_etfs_for_trend_rider(self):
        tickers, benchmark = trader.strategy_universe("us", "trend_rider_v3_us")
        self.assertEqual(tickers, sorted(ETF_UNIVERSE))
        self.assertEqual(benchmark, "SPY")
        self.assertEqual(trader.strategy_data_market("us", "trend_rider_v3_us"), "etfs")

        global_tickers, global_benchmark = trader.strategy_universe("us", "trend_rider_v3_global")
        self.assertEqual(global_tickers, sorted(set(GLOBAL_ETF_UNIVERSE)))
        self.assertEqual(global_benchmark, "SPY")

        hk_tickers, hk_benchmark = trader.strategy_universe("us", "trend_rider_v3_hk")
        self.assertEqual(hk_tickers, sorted(set(HK_ETF_UNIVERSE)))
        self.assertEqual(hk_benchmark, "SPY")

        v4_tickers, v4_benchmark = trader.strategy_universe("us", "trend_rider_v4")
        self.assertEqual(v4_tickers, sorted(set(TREND_RIDER_V4_UNIVERSE)))
        self.assertEqual(v4_benchmark, "SPY")

    def test_filter_tradable_columns_preserves_strategy_assets(self):
        close_data = pd.DataFrame(columns=["SPY", "TQQQ", "UPRO", "GLD", "DBC", "SHY"])
        tickers = trader.filter_tradable_tickers(close_data, ["SPY", "TQQQ", "MISSING"])
        self.assertEqual(tickers, ["SPY", "TQQQ"])

    def test_stock_strategies_keep_market_cache(self):
        self.assertEqual(trader.strategy_data_market("us", "recovery_mom_top10"), "us")


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