"""财务数据加载与清洗模块。

提供财务数据的加载、清洗和与行情数据拼接功能。
"""

from datetime import datetime, timedelta
from typing import Optional, List

import polars as pl

from src.data.storage import Storage


class FinancialLoader:
    """财务数据加载器。

    负责财务数据的清洗、去重，以及与行情数据的 as-of join。

    Attributes:
        storage: DuckDB 存储实例
    """

    def __init__(self):
        self.storage = Storage()

    def load_financial_data(
        self,
        table_name: str,
        columns: List[str],
        start_date: str,
        end_date: str,
        ts_code: Optional[str] = None,
    ) -> pl.DataFrame:
        """加载并清洗财务数据。

        数据清洗流程：
        1. 仅保留 report_type == '1'（合并报表）
        2. 按 (ts_code, f_ann_date) 分组，按 update_flag 降序去重
        3. 转换为 Date 类型，按 ts_code 和 f_ann_date 排序

        Args:
            table_name: 财务表名（如 'financial_income'）
            columns: 需要的字段列表（必须包含核心字段）
            start_date: 数据开始日期（YYYYMMDD）
            end_date: 数据结束日期（YYYYMMDD）
            ts_code: 可选，单个股票代码过滤

        Returns:
            清洗后的 Polars DataFrame，已排序，f_ann_date 为 pl.Date 类型
        """
        # 确保包含必要字段
        required_cols = {
            "ts_code",
            "f_ann_date",
            "report_type",
            "update_flag",
            "end_date",
        }
        query_cols = list(set(columns) | required_cols)

        # 从数据库加载原始数据
        df = self._load_from_db(table_name, query_cols, start_date, end_date, ts_code)

        if df.is_empty():
            return df

        # 步骤1: 仅保留合并报表 (report_type 可能是字符串或整数)
        df = df.filter(pl.col("report_type") == 1)

        # 步骤2: 添加辅助列用于排序和过滤
        df = df.with_columns(
            [
                # 将报告期 20231231 转为整数
                pl.col("end_date").cast(pl.Int32).alias("end_date_int"),
                # 将更新标识转为整数，容错处理 null 和空字符串
                pl.col("update_flag")
                .fill_null("0")
                .cast(pl.Int32, strict=False)
                .fill_null(0)
                .alias("update_flag_int"),
            ]
        )

        # 步骤3: 绝对确定性排序
        # 按照: 股票代码(升序) -> 公告日(升序) -> 报告期(升序) -> 更新标识(升序)
        # 这样排序后，同一天发布的多份报表中，end_date最大、update_flag最新的，
        # 必然排在这一天的最后面
        df = df.sort(["ts_code", "f_ann_date", "end_date_int", "update_flag_int"])

        # 步骤4: 核心算法 - 计算截至每一个公告日，市场"见过的"最新报告期（最高水位线）
        df = df.with_columns(
            pl.col("end_date_int").cum_max().over("ts_code").alias("max_end_date_seen")
        )

        # 步骤5: 剔除"历史包袱"（解决2026年发2021年财报的问题）
        # 如果今天发布的财报，它的 end_date 小于我们之前已经见过的最大值，
        # 说明它是陈旧的追溯调整，直接抛弃！
        df = df.filter(pl.col("end_date_int") == pl.col("max_end_date_seen"))

        # 步骤6: 唯一化处理：满足 join_asof 的前置要求
        # 经过上述处理后，同一个 f_ann_date 的最后一行，必定是 end_date 最大的那一份
        # 我们只保留这最后一行，确保每个 f_ann_date 只有唯一的一条记录暴露给行情去 join
        df = df.unique(subset=["ts_code", "f_ann_date"], keep="last")

        # 步骤7: 清理临时辅助列
        df = df.drop(["end_date_int", "update_flag_int", "max_end_date_seen"])

        # 步骤8: 确保 f_ann_date 是 Date 类型并排序（join_asof 要求）
        if df["f_ann_date"].dtype != pl.Date:
            raise TypeError(
                f"f_ann_date 必须是 Date 类型，实际类型为 {df['f_ann_date'].dtype}. "
                f"请检查数据库表结构，确保日期字段为 DATE 类型。"
            )

        df = df.sort(["ts_code", "f_ann_date"])

        return df

    def merge_financial_with_price(
        self,
        df_price: pl.DataFrame,
        df_financial: pl.DataFrame,
        financial_cols: List[str],
    ) -> pl.DataFrame:
        """将财务数据拼接到行情数据。

        使用 join_asof 向后匹配：对于每个交易日，找到最近的历史公告数据。

        注意：输入的 df_price 的 trade_date 必须是 pl.Date 类型且已排序。

        Args:
            df_price: 行情数据 DataFrame，必须包含 ts_code, trade_date（Date 类型）
            df_financial: 财务数据 DataFrame（已通过 load_financial_data 清洗，Date 类型）
            financial_cols: 需要从财务表保留的字段列表

        Returns:
            拼接后的 DataFrame（trade_date 仍为 Date 类型）
        """
        if df_financial.is_empty():
            # 财务数据为空，返回行情数据（财务列为空）
            for col in financial_cols:
                if col not in df_price.columns:
                    df_price = df_price.with_columns([pl.lit(None).alias(col)])
            return df_price

        # 执行 asof join: 向后寻找最近的历史数据
        # strategy='backward': 对于每个 trade_date，找 f_ann_date <= trade_date 的最新记录
        merged = df_price.join_asof(
            df_financial.select(["ts_code", "f_ann_date"] + financial_cols),
            left_on="trade_date",
            right_on="f_ann_date",
            by="ts_code",
            strategy="backward",
        )

        # 删除 f_ann_date_right 列（如果有），避免多次 asof join 时列名冲突
        if "f_ann_date_right" in merged.columns:
            merged = merged.drop("f_ann_date_right")

        return merged

    def _load_from_db(
        self,
        table_name: str,
        columns: List[str],
        start_date: str,
        end_date: str,
        ts_code: Optional[str] = None,
    ) -> pl.DataFrame:
        """从数据库加载财务数据（内部方法）。"""
        conn = self.storage._connection

        cols_str = ", ".join(f'"{c}"' for c in columns)

        start_dt = datetime.strptime(start_date, "%Y%m%d").date()
        end_dt = datetime.strptime(end_date, "%Y%m%d").date()

        conditions = [f"f_ann_date BETWEEN '{start_dt}' AND '{end_dt}'"]
        if ts_code:
            conditions.append(f"ts_code = '{ts_code}'")

        where_clause = " AND ".join(conditions)
        query = f"SELECT {cols_str} FROM {table_name} WHERE {where_clause} ORDER BY ts_code, f_ann_date"

        try:
            df = conn.sql(query).pl()
            return df
        except Exception as e:
            print(f"[FinancialLoader] 加载 {table_name} 失败: {e}")
            return pl.DataFrame()

    def get_date_range_with_lookback(
        self,
        start_date: str,
        end_date: str,
        lookback_years: int = 2,
    ) -> tuple[str, str]:
        """计算包含回看期的日期范围。

        Args:
            start_date: 原始开始日期（YYYYMMDD）
            end_date: 原始结束日期（YYYYMMDD）
            lookback_years: 回看年数（默认2年，防止ST/停牌公司财报缺失）

        Returns:
            (扩展后的开始日期, 结束日期)
        """
        start_dt = datetime.strptime(start_date, "%Y%m%d")
        adjusted_start = start_dt - timedelta(days=365 * lookback_years)
        return adjusted_start.strftime("%Y%m%d"), end_date
