"""
数据缓存管理模块
"""

import os
import pandas as pd
from pathlib import Path
from typing import Optional


class DataCache:
    """CSV文件缓存管理器"""

    def __init__(self, cache_dir: str = "data_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)

    def _get_cache_path(self, code: str, start_date: str, end_date: str) -> Path:
        """生成缓存文件路径"""
        # 统一日期格式为 YYYYMMDD
        sd = start_date.replace("-", "")
        ed = end_date.replace("-", "")
        safe_code = code.replace(".", "_")
        return self.cache_dir / f"{safe_code}_{sd}_{ed}.csv"

    def get(self, code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        从缓存读取数据

        Returns:
            DataFrame or None（缓存不存在）
        """
        cache_path = self._get_cache_path(code, start_date, end_date)
        if cache_path.exists():
            df = pd.read_csv(cache_path)
            df["date"] = pd.to_datetime(df["date"])
            return df
        return None

    def set(self, code: str, start_date: str, end_date: str, df: pd.DataFrame) -> None:
        """保存数据到缓存"""
        cache_path = self._get_cache_path(code, start_date, end_date)
        df.to_csv(cache_path, index=False)

    def clear(self, code: str = None) -> None:
        """清除缓存"""
        if code:
            # 清除指定代码的缓存
            for f in self.cache_dir.glob(f"{code.replace('.', '_')}*.csv"):
                f.unlink()
        else:
            # 清除所有缓存
            for f in self.cache_dir.glob("*.csv"):
                f.unlink()