import unittest
from datetime import datetime, timedelta
from unittest.mock import patch

import pandas as pd
import yfinance as yf
from dateutil.relativedelta import relativedelta

import error_printer

error_printer.configure_pretty_errors()
error_printer.configure_icecream()


class StockDataDownloader:
    def __init__(self, months: int = 1):
        """
        初始化 StockDataDownloader 類的實例。

        參數:
        months (int): 設定 start 日期為 end 日期之前的月數。預設值為 1。
        """
        self.months = months
        self.start_date, self.end_date = self.set_date_range()

    def set_date_range(self) -> tuple[datetime, datetime]:
        """
        設定日期範圍，用於下載股票資料。

        返回:
        tuple[datetime, datetime]: 包含 start 和 end 日期的元組。
        """
        end = datetime.now()
        start = end - relativedelta(months=self.months)
        print("Start date:", start)
        print("End date:", end)
        return start, end

    def download_stock_data(self, stock_symbol: str) -> pd.DataFrame:
        """
        下載指定股票程式碼在指定日期範圍內的股票資料。

        參數:
        stock_symbol (str): 要下載的股票程式碼。

        返回:
        pd.DataFrame: 包含股票資料的 DataFrame。
        """
        return yf.download(stock_symbol, start=self.start_date, end=self.end_date)


class TestStockDataDownloader(unittest.TestCase):
    def test_initialization(self):
        """測試類初始化是否正確設定月份和日期範圍。"""
        downloader = StockDataDownloader(3)
        self.assertEqual(downloader.months, 3)

    @patch(__name__ + ".datetime")
    def test_date_range(self, mock_datetime):
        """測試日期範圍是否正確計算。"""
        # 設定模擬的現在時間
        mock_now = datetime(2024, 5, 6, 9, 27, 48, 439831)
        mock_datetime.now.return_value = mock_now

        downloader = StockDataDownloader(2)
        expected_start_date = mock_now - timedelta(days=60)  # 假設每月30天
        expected_end_date = mock_now

        actual_start_date, actual_end_date = downloader.set_date_range()
        self.assertEqual(actual_start_date, expected_start_date)
        self.assertEqual(actual_end_date, expected_end_date)

    @patch("yfinance.download")
    def test_download_stock_data(self, mock_download):
        """測試是否正確下載股票資料。"""
        # 設定模擬的現在時間
        mock_now = datetime(2024, 5, 6, 9, 27, 48, 439831)
        mock_start_date = mock_now - timedelta(days=30)
        mock_end_date = mock_now

        # 設定模擬的股票資料
        mock_data = pd.DataFrame({"Date": pd.date_range(start=mock_start_date, end=mock_end_date, freq="D"), "Open": [100.0] * 31, "High": [101.0] * 31, "Low": [99.0] * 31, "Close": [100.0] * 31, "Volume": [1000] * 31})
        mock_data.set_index("Date", inplace=True)
        mock_download.return_value = mock_data

        downloader = StockDataDownloader(1)
        stock_data = downloader.download_stock_data("00929.TW")

        self.assertTrue(stock_data.equals(mock_data))


# 檢查目前程式是否被作為主程式執行
if __name__ == "__main__":
    unittest.main()

    # 建立 StockDataDownloader 實例
    # downloader = StockDataDownloader(months=1)

    # 下載股票資料
    # stock_data = downloader.download_stock_data("00929.TW")

    # 進一步處理股票資料
    # ...
    # ic(stock_data)