"""并表 JoinResolver — 根据 BSON required_datasets 执行附表加载、pre-join DSL、Join、列重命名

columns 为附表源列名，column_suffix 由调用方（ms-quant）主动告知，
JoinResolver 用 source + suffix 构造最终别名。"""
from typing import Any

import polars as pl

from ldnode import logger
from .computed_factors import planning_computed_factors
from .join_rules import JOIN_RULES, SCOPE_TO_FILE, JoinRule, TimeAlignStrategy


def resolve_joins(
    main_lf: pl.LazyFrame,
    strategy_scope: int,
    required_datasets: list[dict[str, Any]],
    builtin_factors: list[str],
    aux_data_overrides: dict[int, pl.LazyFrame] | None = None,
) -> tuple[pl.LazyFrame, list[str]]:
    """执行所有并表操作，返回 (合并后 LazyFrame, 更新后 builtin_factors)。

    aux_data_overrides: 按 scope 覆盖默认附表数据源，
    用于盘中模式传入 CONCAT(历史, 实时) 的合并数据。
    """
    builtin_factors = list(builtin_factors)

    for rd in required_datasets:
        target_scope = rd["security_type"]
        columns: list[str] = rd.get("columns", [])
        pre_join_dsl: list[list[dict]] = rd.get("pre_join_dsl", [])
        # column_suffix 由调用方主动告知，fallback 到 JoinRule 硬编码（向后兼容）
        suffix: str = rd.get("column_suffix") or ""

        rule_key = (strategy_scope, target_scope)
        rule = JOIN_RULES.get(rule_key)
        if rule is None:
            raise ValueError(f"JoinResolver: 未找到 Join 规则 {rule_key}，无法执行并表")

        if not suffix:
            suffix = rule.column_suffix

        # 优先使用 handler 传入的覆盖数据源（盘中模式 concat 实时+历史）
        if aux_data_overrides and target_scope in aux_data_overrides:
            aux_lf = aux_data_overrides[target_scope]
        else:
            aux_file = SCOPE_TO_FILE.get(target_scope)
            if aux_file is None:
                raise ValueError(f"JoinResolver: 未找到 Scope {target_scope} 的数据集映射")
            aux_lf = aux_file.output_lazyframe()

        # pre-join DSL（在目标 Scope 数据集上执行，使用原始列名）
        if pre_join_dsl:
            aux_builtin = list(aux_lf.collect_schema().names())
            aux_lf = planning_computed_factors(
                frame=aux_lf,
                builtin_factors=aux_builtin,
                computed_factors=pre_join_dsl,
                exclude_suspension=False,
            )

        # columns 为源列名，加 suffix 构造最终别名
        aux_tk = rule.aux_time_key or rule.time_key
        source_to_alias = {col: (col + suffix if suffix else col) for col in columns}
        select_cols = [rule.aux_join_key, aux_tk, *source_to_alias.keys()]
        aux_schema = aux_lf.collect_schema().names()
        missing = [c for c in source_to_alias if c not in aux_schema and c not in (rule.aux_join_key, aux_tk)]
        if missing:
            raise ValueError(f"JoinResolver: 附表 Scope {target_scope} 中找不到列 {missing}，"
                             f"请确认 pre_join_dsl 输出与 columns 映射一致")
        aux_lf = aux_lf.select([c for c in select_cols if c in aux_schema])

        # 列重命名：源列名 → 别名（加后缀）
        rename_map = {src: alias for src, alias in source_to_alias.items() if src != alias}
        if rename_map:
            aux_lf = aux_lf.rename(rename_map)

        # 冲突处理：从主表移除与别名冲突的列（附表优先）
        alias_cols = list(source_to_alias.values())
        main_cols = set(main_lf.collect_schema().names())
        conflict_cols = [c for c in alias_cols if c in main_cols]
        if conflict_cols:
            main_lf = main_lf.drop(conflict_cols)

        # Join
        main_lf = _execute_join(main_lf, aux_lf, rule)

        existing = set(builtin_factors)
        new_cols = [c for c in alias_cols if c not in existing]
        if len(new_cols) < len(alias_cols):
            dupes = [c for c in alias_cols if c in existing]
            logger.calculator.warning(f"JoinResolver: builtin_factors 中已存在列 {dupes}，跳过重复添加")
        builtin_factors.extend(new_cols)

    return main_lf, builtin_factors


def _execute_join(
    main_lf: pl.LazyFrame,
    aux_lf: pl.LazyFrame,
    rule: JoinRule,
) -> pl.LazyFrame:
    aux_tk = rule.aux_time_key or rule.time_key
    if rule.time_align == TimeAlignStrategy.POINT_IN_TIME:
        return main_lf.sort(rule.time_key).join_asof(
            aux_lf.sort(aux_tk),
            left_on=rule.time_key,
            right_on=aux_tk,
            by_left=rule.main_join_key,
            by_right=rule.aux_join_key,
            strategy="backward",
        )
    return main_lf.join(
        aux_lf,
        left_on=[rule.main_join_key, rule.time_key],
        right_on=[rule.aux_join_key, aux_tk],
        how="left",
    )


def _collect_required_datasets(parameters: dict[str, Any]) -> list[dict[str, Any]] | None:
    """从不同参数结构中提取 required_datasets。

    - backtest/evalution/binning: strategy_config.required_datasets
    - screen/glance: strategy_configs[*].required_datasets（合并去重）
    - fallback: 顶层 required_datasets
    """
    # backtest/evalution/binning
    strategy_config = parameters.get("strategy_config")
    if strategy_config:
        rd = strategy_config.get("required_datasets")
        if rd:
            return rd
    # screen/glance: 合并所有策略配置中的 required_datasets
    strategy_configs = parameters.get("strategy_configs")
    if strategy_configs:
        merged: dict[tuple[int, str], dict[str, Any]] = {}
        for cfg in strategy_configs.values():
            for item in cfg.get("required_datasets") or []:
                key = (item["security_type"], item.get("column_suffix", ""))
                if key not in merged:
                    merged[key] = {
                        "security_type": item["security_type"],
                        "columns": list(item.get("columns", [])),
                        "column_suffix": item.get("column_suffix", ""),
                        "pre_join_dsl": list(item.get("pre_join_dsl", [])),
                    }
                else:
                    existing = merged[key]
                    existing_cols = set(existing["columns"])
                    for c in item.get("columns", []):
                        if c not in existing_cols:
                            existing["columns"].append(c)
                    existing["pre_join_dsl"].extend(item.get("pre_join_dsl", []))
        return list(merged.values()) if merged else None
    # fallback: 顶层
    return parameters.get("required_datasets")


def maybe_resolve_joins(
    lf: pl.LazyFrame,
    strategy_scope: int,
    parameters: dict[str, Any],
    default_builtin_factors: list[str],
    aux_data_overrides: dict[int, pl.LazyFrame] | None = None,
) -> tuple[pl.LazyFrame, list[str]]:
    """Handler 统一入口：有 required_datasets 则并表，否则透传。

    aux_data_overrides: 按 scope 覆盖默认附表数据源，
    用于盘中模式传入 CONCAT(历史, 实时) 的合并数据。
    """
    required_datasets = _collect_required_datasets(parameters)
    if not required_datasets:
        return lf, list(default_builtin_factors)
    return resolve_joins(lf, strategy_scope, required_datasets, default_builtin_factors, aux_data_overrides)
