"""因子缓存存储模块（HDF5 版）。

提供单个因子的 HDF5 读写、日期范围查询、以及新旧数据校验功能。
每个因子对应一个独立的 HDF5 文件，实现轻量级、可增量更新的本地缓存。
"""

from pathlib import Path
from typing import Any, Dict, Optional, Tuple, cast

import pandas as pd
import polars as pl

from src.config.settings import settings


class FactorStorage:
    """因子缓存存储类。

    每个因子对应一个 HDF5 文件，内部以 `factor_values` 为 key 存储 DataFrame，
    包含 `trade_date`, `ts_code`, `{factor_name}` 三列。

    Attributes:
        base_dir: 缓存文件存放目录
        _HDF_KEY: HDF5 内部存储 key，固定为 "factor_values"
    """

    _HDF_KEY = "factor_values"

    def __init__(self, base_dir: Optional[Path] = None) -> None:
        """初始化存储器。

        Args:
            base_dir: 缓存根目录。若未传入，使用 `settings.data_path_resolved / "factor"`。
        """
        if base_dir is None:
            self.base_dir = settings.data_path_resolved / "factor"
        else:
            self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def _file_path(self, factor_name: str) -> Path:
        """返回单个因子的 HDF5 文件路径。"""
        return self.base_dir / f"{factor_name}.h5"

    def exists(self, factor_name: str) -> bool:
        """检查因子的缓存文件是否存在。"""
        return self._file_path(factor_name).exists()

    def _load_pandas(self, factor_name: str) -> pd.DataFrame:
        """内部方法：以 pandas DataFrame 形式读取缓存。"""
        path = self._file_path(factor_name)
        if not path.exists():
            raise FileNotFoundError(f"因子缓存文件不存在: {path}")
        df = pd.read_hdf(path, key=self._HDF_KEY)  # type: ignore[assignment]
        if not isinstance(df, pd.DataFrame):
            raise ValueError("读取 HDF5 返回的不是 DataFrame")
        return df

    def save(self, factor_name: str, df: pl.DataFrame) -> None:
        """保存因子数据到 HDF5。

        若文件已存在，会读取旧数据并与新数据按 `trade_date` + `ts_code` 合并，
        新数据优先覆盖重叠部分，保留旧数据中未覆盖的行，最后整体覆写文件。

        Args:
            factor_name: 因子名称
            df: 包含 `trade_date`, `ts_code`, `{factor_name}` 的 Polars DataFrame
        """
        path = self._file_path(factor_name)
        path.parent.mkdir(parents=True, exist_ok=True)

        new_pdf = df.to_pandas()
        merge_cols = ["trade_date", "ts_code"]
        val_col = factor_name

        if path.exists():
            old_pdf = self._load_pandas(factor_name)
            combined = pd.merge(
                old_pdf, new_pdf, on=merge_cols, how="outer", suffixes=("_old", "_new")
            )
            # 使用新列覆盖旧列，若新列为 NaN 则保留旧列
            new_s = cast(pd.Series, combined[f"{val_col}_new"])
            old_s = cast(pd.Series, combined[f"{val_col}_old"])
            combined[val_col] = new_s.combine_first(old_s)
            # 删除中间列
            drop_cols = [
                c for c in combined.columns if c.endswith("_old") or c.endswith("_new")
            ]
            combined = combined.drop(columns=drop_cols)
            # 恢复列顺序
            combined = combined[merge_cols + [val_col]]
            combined.to_hdf(
                path, key=self._HDF_KEY, mode="w", format="table", data_columns=True
            )
        else:
            new_pdf.to_hdf(
                path, key=self._HDF_KEY, mode="w", format="table", data_columns=True
            )

    def load(
        self,
        factor_name: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> pl.DataFrame:
        """从 HDF5 读取因子数据。

        Args:
            factor_name: 因子名称
            start_date: 可选，起始日期过滤 (YYYYMMDD)
            end_date: 可选，结束日期过滤 (YYYYMMDD)

        Returns:
            Polars DataFrame
        """
        pdf = self._load_pandas(factor_name)
        if start_date is not None:
            pdf = pdf[pdf["trade_date"] >= start_date]
        if end_date is not None:
            pdf = pdf[pdf["trade_date"] <= end_date]
        return pl.DataFrame(pdf)

    def get_date_range(self, factor_name: str) -> Optional[Tuple[str, str]]:
        """获取已有缓存的日期范围。

        Args:
            factor_name: 因子名称

        Returns:
            (min_date, max_date) 元组，不存在时返回 None
        """
        if not self.exists(factor_name):
            return None
        pdf = self._load_pandas(factor_name)
        min_date = str(pdf["trade_date"].min())
        max_date = str(pdf["trade_date"].max())
        return min_date, max_date

    def validate(
        self,
        factor_name: str,
        new_df: pl.DataFrame,
        tolerance: float = 1e-6,
    ) -> Tuple[bool, Dict[str, Any]]:
        """校验新计算的数据与本地缓存的重叠部分是否一致。

        Args:
            factor_name: 因子名称
            new_df: 新计算的全量数据
            tolerance: 最大绝对差异容忍度

        Returns:
            (passed, stats)
            passed: True 表示校验通过（差异在容忍度内或无重叠）
            stats: 包含 max_abs_diff, mean_abs_diff, matched_rows 的字典
        """
        if not self.exists(factor_name):
            return True, {}

        old_pdf = self._load_pandas(factor_name)
        new_pdf = new_df.to_pandas()
        merge_cols = ["trade_date", "ts_code"]
        merged = old_pdf.merge(
            new_pdf, on=merge_cols, how="inner", suffixes=("_old", "_new")
        )

        if len(merged) == 0:
            return True, {}

        val_col = factor_name
        diff_series = (merged[f"{val_col}_new"] - merged[f"{val_col}_old"]).abs()
        valid_diff = diff_series.dropna()

        matched_rows = len(valid_diff)
        if matched_rows == 0:
            return True, {"matched_rows": 0, "max_abs_diff": 0.0, "mean_abs_diff": 0.0}

        stats = {
            "matched_rows": matched_rows,
            "max_abs_diff": float(valid_diff.max()),
            "mean_abs_diff": float(valid_diff.mean()),
        }

        if stats["max_abs_diff"] > tolerance:
            return False, stats
        return True, stats
