import tempfile
import unittest
from pathlib import Path
from unittest import mock

import pandas as pd

from research import fetch_historical


class FetchHistoricalTests(unittest.TestCase):
    def test_fetch_all_historical_ohlcv_writes_field_specific_csvs(self):
        dates = pd.to_datetime(["2024-01-02", "2024-01-03"])
        raw = pd.DataFrame(
            {
                ("Close", "AAA"): [10.0, 11.0],
                ("Close", "BBB"): [20.0, 21.0],
                ("High", "AAA"): [10.5, 11.5],
                ("High", "BBB"): [20.5, 21.5],
                ("Low", "AAA"): [9.5, 10.5],
                ("Low", "BBB"): [19.5, 20.5],
                ("Volume", "AAA"): [1000.0, 1100.0],
                ("Volume", "BBB"): [2000.0, 2100.0],
            },
            index=dates,
        )
        raw.columns = pd.MultiIndex.from_tuples(raw.columns)

        with tempfile.TemporaryDirectory() as tmpdir:
            with mock.patch.object(fetch_historical, "DATA_DIR", tmpdir):
                with mock.patch.object(fetch_historical, "OUT_PATH", str(Path(tmpdir) / "us_pit.csv")):
                    with mock.patch("research.fetch_historical.uh.load_sp500_history", return_value={"AAA": [[None, None]], "BBB": [[None, None]]}):
                        with mock.patch("research.fetch_historical.uh.all_tickers_ever", return_value=["AAA", "BBB"]):
                            with mock.patch("research.fetch_historical.yf.download", return_value=raw):
                                panels = fetch_historical.fetch_all_historical_ohlcv(force=True)

            self.assertEqual(set(panels.keys()), {"close", "high", "low", "volume"})
            self.assertTrue((Path(tmpdir) / "us_pit.csv").exists())
            self.assertTrue((Path(tmpdir) / "us_pit_close.csv").exists())
            self.assertTrue((Path(tmpdir) / "us_pit_high.csv").exists())
            self.assertTrue((Path(tmpdir) / "us_pit_low.csv").exists())
            self.assertTrue((Path(tmpdir) / "us_pit_volume.csv").exists())


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