import sqlite3
import tempfile
import unittest
from pathlib import Path

from database import ReviewDatabase
from screener import (
    CURATED_STRATEGIES,
    FACTOR_FIELDS,
    FACTOR_GROUPS,
    ScreenerEngine,
    _quarter_periods,
)


class CuratedScreenerTests(unittest.TestCase):
    def test_first_batch_contains_ten_distinct_curated_strategies(self):
        self.assertEqual(10, len(CURATED_STRATEGIES))
        self.assertEqual(10, len({item["name"] for item in CURATED_STRATEGIES}))
        self.assertTrue(
            all(item["formula"]["meta"]["library"] == "curated" for item in CURATED_STRATEGIES)
        )

    def test_every_curated_formula_uses_supported_factors(self):
        with tempfile.TemporaryDirectory() as root:
            database = ReviewDatabase(Path(root) / "review.db")
            engine = ScreenerEngine(database)
            for strategy in CURATED_STRATEGIES:
                formula = engine.validate_formula(strategy["formula"])
                fields = {
                    item["field"]
                    for item in formula["filters"] + formula["score"]
                }
                self.assertTrue(fields.issubset(FACTOR_FIELDS), strategy["name"])

    def test_factor_groups_cover_every_quant_factor(self):
        grouped = [field for fields in FACTOR_GROUPS.values() for field in fields]
        self.assertEqual(set(FACTOR_FIELDS), set(grouped))
        self.assertEqual(len(grouped), len(set(grouped)))

    def test_database_migrates_valuation_and_fundamental_columns(self):
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "review.db"
            ReviewDatabase(path)
            connection = sqlite3.connect(path)
            try:
                indicator_columns = {
                    row[1] for row in connection.execute("PRAGMA table_info(daily_indicators)")
                }
                tables = {
                    row[0] for row in connection.execute(
                        "SELECT name FROM sqlite_master WHERE type='table'"
                    )
                }
            finally:
                connection.close()
            self.assertTrue({"pe_ttm", "pb", "ps_ttm", "dv_ttm"}.issubset(indicator_columns))
            self.assertIn("fundamental_indicators", tables)

    def test_quarter_periods_stop_at_selected_date(self):
        periods = _quarter_periods("20260722", 5)
        self.assertEqual(
            ["20250630", "20250930", "20251231", "20260331", "20260630"],
            periods,
        )

    def test_factor_health_summary_uses_availability_counts(self):
        with tempfile.TemporaryDirectory() as root:
            database = ReviewDatabase(Path(root) / "review.db")
            with database.connect() as connection:
                connection.execute(
                    "INSERT INTO daily_bars (trade_date, ts_code) VALUES (?, ?)",
                    ("20260722", "600000.SH"),
                )
                connection.execute(
                    """
                    INSERT INTO daily_indicators
                        (trade_date, ts_code, pe_ttm)
                    VALUES (?, ?, ?)
                    """,
                    ("20260722", "600000.SH", 8.5),
                )
                connection.executemany(
                    "INSERT INTO daily_indicators (trade_date, ts_code) VALUES (?, ?)",
                    [(f"{year}1231", f"{year % 100:02d}0000.SZ") for year in range(2022, 2026)],
                )
                connection.execute(
                    "INSERT INTO auction_factors (trade_date, ts_code) VALUES (?, ?)",
                    ("20260722", "600000.SH"),
                )
                connection.executemany(
                    """
                    INSERT INTO fundamental_indicators (end_date, ann_date, ts_code, roe)
                    VALUES (?, ?, ?, ?)
                    """,
                    [
                        ("20251231", "20260430", f"{index:06d}.SZ", 10.0)
                        for index in range(100)
                    ],
                )

            health = database.factor_health_summary("20260722")
            self.assertTrue(health["market"])
            self.assertTrue(health["auction"])
            self.assertTrue(health["valuation"])
            self.assertTrue(health["fundamental"])
            self.assertTrue(health["dividend_history"])
            self.assertEqual(health["valuation_rows"], 1)
            self.assertEqual(health["fundamental_rows"], 100)
            self.assertEqual(health["dividend_years"], 5)


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