#!/usr/bin/env python3
from typing import Any, Dict, List, Optional
import logging
import time

from core.modules.data_contract.cache import ContractCacheManager
from core.modules.strategy.engines.shared.data_classes.strategy_settings.dict_view_settings import (
    StrategySettingsView,
)
from core.modules.market_profile import get_market_profile
from core.modules.strategy.engines.shared.helpers.market_profile_id import (
    resolve_market_profile_id,
)
from core.modules.strategy.engines.shared.helpers.backtest_calendar_context import (
    BacktestCalendarContext,
)
from core.modules.strategy.engines.shared.helpers.skip_investment_when import (
    stamp_stock_status_at_trigger,
)
from core.modules.strategy.engines.shared.helpers.stock_status_risk_context import (
    build_stock_status_risk_runtime_context,
)
from core.modules.strategy.engines.shared.helpers.tradability import stamp_buy_tradability
from core.modules.strategy.engines.shared.helpers.tradability_stats import (
    tradability_bundle_from_opportunities,
)
from core.modules.strategy.engines.shared.helpers.simulation_day_execution import (
    execute_pending_exits_on_active,
    fill_pending_buys,
    queue_deferred_buy,
    resolve_pending_buys_at_end,
    resolve_pending_exits_on_active_at_end,
)
from core.modules.strategy.engines.shared.helpers.simulation_pricing import (
    apply_buy_slippage,
    trade_price_defers_to_next_session,
    trade_theoretical_price_same_day,
)
from core.modules.strategy.engines.shared.helpers.strategy_runtime import resolve_worker_class
from core.modules.indicator import IndicatorService
from core.modules.strategy.engines.shared.performance_profiler import PerformanceProfiler
from core.modules.strategy.enums import ExecutionMode
from core.modules.strategy.engines.shared.data_classes.investment_state import (
    InvestmentLifecycle,
    ScanSignalPhase,
)
from core.modules.strategy.services.data import StrategyDataInjectionService
from core.modules.strategy.services.data.injection.job_contract_batch import (
    StrategyJobContractBatch,
)
from core.modules.strategy.services.data.output import EnumeratorOutputWriterService
from core.utils.date.date_utils import DateUtils

logger = logging.getLogger(__name__)

MAX_LOOKBACK_DAYS = 60

# 多进程 spawn 下每个子进程独立；首只股票 job 时预热 pandas-ta，后续 job 复用。
_indicator_runtime_warmed = False


def _warmup_indicator_runtime_once() -> None:
    global _indicator_runtime_warmed
    if _indicator_runtime_warmed:
        return
    IndicatorService.warmup()
    _indicator_runtime_warmed = True


def enumeration_actual_start_date(start_date: str, min_required_records: int) -> str:
    lookback_days = min(min_required_records, MAX_LOOKBACK_DAYS)
    try:
        adjusted_days = int(lookback_days * 1.5)
        return DateUtils.sub_days(start_date, adjusted_days)
    except Exception:
        return start_date


class OpportunityEnumeratorWorker:
    def __init__(
        self,
        job_payload: Dict[str, Any],
        *,
        stock_id: Optional[str] = None,
        contract_cache: Optional[ContractCacheManager] = None,
        fresh_strategy_cache: bool = True,
        job_contract_batch: Optional[StrategyJobContractBatch] = None,
    ):
        self.job_payload = job_payload
        self.stock_id = str(stock_id or job_payload.get("stock_id") or "").strip()
        if not self.stock_id:
            raise ValueError("OpportunityEnumeratorWorker 缺少 stock_id")
        self._fresh_strategy_cache = bool(fresh_strategy_cache)
        self._job_contract_batch = job_contract_batch
        self.strategy_name = job_payload["strategy_name"]
        self.start_date = job_payload["start_date"]
        self.end_date = job_payload["end_date"]
        self.profiler = PerformanceProfiler(self.stock_id)
        self.settings = StrategySettingsView.from_dict(job_payload["settings"])
        self.settings_dict = self.settings.to_dict()
        self.stock_info = self._load_stock_info()
        self.contract_cache = contract_cache or ContractCacheManager()
        self.data_manager = StrategyDataInjectionService(
            stock_id=self.stock_id,
            settings=self.settings,
            contract_cache=self.contract_cache,
            global_extra_cache=self.job_payload.get("global_extra_cache"),
        )
        self.data_manager.attach_load_profiler(self.profiler)
        self.opportunity_counter = 0
        self.simulation = self.settings.simulation_settings
        profile_id = resolve_market_profile_id(
            self.job_payload,
            settings_market_profile=self.settings.market_profile,
        )
        self.market_profile = get_market_profile(profile_id)
        self.backtest_calendar = BacktestCalendarContext.from_dict(
            job_payload.get("backtest_calendar")
        )
        self.stock_status_risk = build_stock_status_risk_runtime_context(
            stock_meta=self.stock_info,
            settings=self.settings.stock_status_risk,
            stock_id=self.stock_id,
            period_start=self.start_date,
            period_end=self.end_date,
        )
        self._load_user_strategy()

    def _load_stock_info(self) -> Dict[str, Any]:
        # 编排层元数据：不在 settings.data 声明范围内，直调 list 服务（决策 11）。
        fallback = {
            "id": self.stock_id,
            "name": self.stock_id,
            "industry": "",
            "type": "",
            "exchange_center": "",
        }
        try:
            from core.modules.data_manager import DataManager

            row = DataManager().stock.list.load_meta(self.stock_id)
            if isinstance(row, dict) and row.get("id"):
                return {
                    **fallback,
                    **row,
                }
        except Exception as exc:
            if str(self.stock_id).upper() != "DUMMY":
                logger.warning(
                    "加载股票元数据失败: %s, error=%s",
                    self.stock_id,
                    exc,
                )
        return fallback

    def _load_user_strategy(self):
        strategy_class = resolve_worker_class(
            self.strategy_name,
            worker_module_path=self.job_payload.get("worker_module_path"),
            worker_class_name=self.job_payload.get("worker_class_name"),
            worker_file_path=str(self.job_payload.get("worker_file_path") or ""),
        )
        self.strategy_instance = strategy_class(
            {
                "stock_id": self.stock_id,
                "execution_mode": ExecutionMode.SCAN.value,
                "strategy_name": self.strategy_name,
                "settings": self.settings_dict,
            }
        )
        if hasattr(self.strategy_instance, "stock_info"):
            self.strategy_instance.stock_info = self.stock_info

    def run(self) -> Dict[str, Any]:
        _warmup_indicator_runtime_once()
        self.profiler.start_timer("total")
        try:
            # step1: prepare query time range
            actual_start_date = self._prepare_actual_start_date()
            # step2: load declared data into injection service
            self._load_runtime_data(actual_start_date)
            # step3: enumerate opportunities over kline timeline
            all_klines = self.data_manager.get_klines()
            if not all_klines or len(all_klines) < self.settings.min_required_records:
                return self._finish_with_metrics(
                    success=True,
                    opportunities_dict=[],
                    all_klines=all_klines or [],
                    skipped_short_data=True,
                )
            tracker = self._enumerate_opportunities(all_klines)
            # step4: serialize opportunity objects
            opportunities_dict = self._serialize_opportunities(tracker)
            # step5: persist outputs and finalize metrics
            return self._finalize_success_result(all_klines, opportunities_dict)
        except Exception as exc:
            logger.error("enumeration failed: stock_id=%s, error=%s", self.stock_id, exc, exc_info=True)
            return self._finish_with_metrics(
                success=False,
                opportunities_dict=[],
                all_klines=[],
                error=str(exc),
            )

    def _prepare_actual_start_date(self) -> str:
        return enumeration_actual_start_date(
            self.start_date,
            self.settings.min_required_records,
        )

    def _load_runtime_data(self, actual_start_date: str) -> None:
        self.profiler.start_timer("load_data")
        dm = self.data_manager
        if self._job_contract_batch is not None:
            self.profiler.metrics.load_path = "contract_batch"
            t0 = time.perf_counter()
            dm.adopt_job_batch(self._job_contract_batch)
            self.profiler.metrics.time_load_contracts += time.perf_counter() - t0
        else:
            self.profiler.metrics.load_path = "contract"
            dm._current_data = {}
            dm._slot_contracts = {}
            t0 = time.perf_counter()
            dm.hydrate_row_slots(
                actual_start_date,
                self.end_date,
                fresh_strategy_cache=self._fresh_strategy_cache,
            )
            self.profiler.metrics.time_load_contracts += time.perf_counter() - t0
        t0 = time.perf_counter()
        dm.apply_indicators()
        self.profiler.metrics.time_calculate_indicators += time.perf_counter() - t0
        t0 = time.perf_counter()
        dm.rebuild_data_cursor()
        self.profiler.metrics.time_build_cursor += time.perf_counter() - t0
        self.profiler.metrics.time_load_data = self.profiler.end_timer("load_data")

    def _enumerate_opportunities(self, all_klines: List[Dict[str, Any]]) -> Dict[str, Any]:
        tracker = {
            "stock_id": self.stock_id,
            "passed_dates": [],
            "pending_buys": [],
            "active_opportunities": [],
            "all_opportunities": [],
        }
        self.profiler.start_timer("enumerate")
        last_kline = None
        for idx, current_kline in enumerate(all_klines):
            virtual_date_of_today = current_kline["date"]
            tracker["passed_dates"].append(virtual_date_of_today)
            if len(tracker["passed_dates"]) < self.settings.min_required_records:
                continue
            data_of_today = self.data_manager.get_data_until(virtual_date_of_today)
            prev_kline = all_klines[idx - 1] if idx > 0 else None
            self._enumerate_single_day(tracker, current_kline, data_of_today, prev_kline=prev_kline)
            last_kline = current_kline
        self.profiler.metrics.time_enumerate = self.profiler.end_timer("enumerate")
        if last_kline is not None:
            resolve_pending_buys_at_end(
                tracker["pending_buys"],
                tracker["active_opportunities"],
                tracker["all_opportunities"],
                sim=self.simulation,
                market_profile=self.market_profile,
            )
            resolve_pending_exits_on_active_at_end(
                tracker["active_opportunities"],
                last_bar=last_kline,
                sim=self.simulation,
                market_profile=self.market_profile,
            )
            if tracker["active_opportunities"]:
                self._close_all_open_opportunities(tracker, last_kline)
        return tracker

    def _serialize_opportunities(self, tracker: Dict[str, Any]) -> List[Dict[str, Any]]:
        self.profiler.start_timer("serialize")
        opportunities_dict = [opportunity.to_dict() for opportunity in tracker["all_opportunities"]]
        self.profiler.metrics.time_serialize = self.profiler.end_timer("serialize")
        return opportunities_dict

    def _finalize_success_result(
        self,
        all_klines: List[Dict[str, Any]],
        opportunities_dict: List[Dict[str, Any]],
    ) -> Dict[str, Any]:
        if (
            self.job_payload.get("output_dir")
            and opportunities_dict
            and not self.job_payload.get("_dispatch_probe")
        ):
            self.profiler.start_timer("save_csv")
            self._save_stock_results(self.job_payload["output_dir"], opportunities_dict)
            self.profiler.metrics.time_save_csv = self.profiler.end_timer("save_csv")
        return self._finish_with_metrics(
            success=True,
            opportunities_dict=opportunities_dict,
            all_klines=all_klines,
        )

    def _finish_with_metrics(
        self,
        *,
        success: bool,
        opportunities_dict: List[Dict[str, Any]],
        all_klines: List[Dict[str, Any]],
        skipped_short_data: bool = False,
        error: Optional[str] = None,
    ) -> Dict[str, Any]:
        self.profiler.metrics.kline_count = len(all_klines or [])
        self.profiler.metrics.opportunity_count = len(opportunities_dict)
        self.profiler.metrics.target_count = sum(
            len(opportunity.get("completed_targets", []) or [])
            for opportunity in opportunities_dict
        )
        self.profiler.metrics.skipped_short_data = skipped_short_data
        completed_count = 0
        unfinished_count = 0
        for opportunity in opportunities_dict:
            sell_reason = str(opportunity.get("sell_reason", "") or "").lower()
            if sell_reason in {"enumeration_end", "backtest_end"}:
                unfinished_count += 1
            else:
                lifecycle = str(opportunity.get("lifecycle", "") or "").lower()
                phase = str(opportunity.get("signal_phase", "") or "").lower()
                if (
                    lifecycle == InvestmentLifecycle.OPEN.value
                    or phase in {ScanSignalPhase.ACTIVE.value, ScanSignalPhase.TESTING.value}
                ):
                    unfinished_count += 1
                else:
                    completed_count += 1
        self.profiler.metrics.time_total = self.profiler.end_timer("total")
        metrics = self.profiler.finalize()
        n = len(opportunities_dict)
        completion_rate = round((completed_count / n) * 100.0, 1) if n else 0.0
        avg_gap = self._avg_trigger_gap_days(opportunities_dict)
        display_name = str((self.stock_info or {}).get("name") or self.stock_id)
        enum_bundle = (
            self._build_enumeration_report_bundle(opportunities_dict)
            if opportunities_dict
            else self._empty_enumeration_report_bundle()
        )
        payload: Dict[str, Any] = {
            "success": success,
            "stock_id": self.stock_id,
            "stock_name": display_name,
            "opportunity_count": n,
            "completed_count": completed_count,
            "unfinished_count": unfinished_count,
            "completion_rate": completion_rate,
            "avg_opportunity_interval_days": avg_gap,
            "performance_metrics": metrics.to_dict(),
            "enumeration_report_bundle": enum_bundle,
        }
        if error:
            payload["error"] = error
        return payload

    def _empty_enumeration_report_bundle(self) -> Dict[str, Any]:
        display_name = str((self.stock_info or {}).get("name") or self.stock_id)
        return {
            "stock_name": display_name,
            "opportunity_count": 0,
            "report_completed_count": 0,
            "report_unfinished_count": 0,
            "status_completed_count": 0,
            "trigger_gap_days": [],
            "holding_duration_days": [],
            "buy_tradability_sample_count": 0,
            "buy_at_limit_up_count": 0,
            "sell_tradability_sample_count": 0,
            "sell_at_limit_down_count": 0,
            "limit_up_buy_ratio": 0.0,
            "limit_down_sell_ratio": 0.0,
        }

    def _build_enumeration_report_bundle(self, opportunities_dict: List[Dict[str, Any]]) -> Dict[str, Any]:
        """供主进程汇总 ``enumMetrics``，避免再读该股 CSV。字段语义与 ``EnumeratorReport.from_opportunities`` 一致。"""
        display_name = str((self.stock_info or {}).get("name") or self.stock_id)
        if not opportunities_dict:
            return {
                **self._empty_enumeration_report_bundle(),
                "stock_name": display_name,
            }
        report_completed = 0
        report_unfinished = 0
        for opportunity in opportunities_dict:
            sell_reason = str(opportunity.get("sell_reason", "") or "").lower()
            if sell_reason in {"enumeration_end", "backtest_end"}:
                report_unfinished += 1
            else:
                lifecycle = str(opportunity.get("lifecycle", "") or "").lower()
                phase = str(opportunity.get("signal_phase", "") or "").lower()
                if (
                    lifecycle == InvestmentLifecycle.OPEN.value
                    or phase in {ScanSignalPhase.ACTIVE.value, ScanSignalPhase.TESTING.value}
                ):
                    report_unfinished += 1
                else:
                    report_completed += 1
        status_completed = sum(
            1
            for o in opportunities_dict
            if str((o or {}).get("lifecycle") or "").lower() == InvestmentLifecycle.COMPLETE.value
        )
        sorted_rows = sorted(
            opportunities_dict,
            key=lambda r: str((r or {}).get("trigger_date") or ""),
        )
        trigger_dates: List[str] = []
        for r in sorted_rows:
            d = DateUtils.normalize_str(str((r or {}).get("trigger_date") or ""))
            if isinstance(d, str) and d:
                trigger_dates.append(d)
        gaps: List[float] = []
        for idx in range(1, len(trigger_dates)):
            gaps.append(float(DateUtils.diff_days(trigger_dates[idx - 1], trigger_dates[idx])))
        durations: List[float] = []
        for r in sorted_rows:
            d0 = DateUtils.normalize_str(str((r or {}).get("trigger_date") or ""))
            d1 = DateUtils.normalize_str(str((r or {}).get("sell_date") or ""))
            if d0 and d1:
                durations.append(float(DateUtils.diff_days(d0, d1)))
        return {
            "stock_name": display_name,
            "opportunity_count": len(opportunities_dict),
            "report_completed_count": report_completed,
            "report_unfinished_count": report_unfinished,
            "status_completed_count": status_completed,
            "trigger_gap_days": gaps,
            "holding_duration_days": durations,
            **tradability_bundle_from_opportunities(opportunities_dict),
        }

    @staticmethod
    def _avg_trigger_gap_days(opportunities_dict: List[Dict[str, Any]]) -> float:
        """相邻两次机会触发日（trigger_date）之间的间隔天数均值；不足 2 个触发则 0。"""
        if not opportunities_dict or len(opportunities_dict) < 2:
            return 0.0
        sorted_rows = sorted(
            opportunities_dict,
            key=lambda r: str((r or {}).get("trigger_date") or ""),
        )
        trigger_dates: List[str] = []
        for r in sorted_rows:
            d = DateUtils.normalize_str(str((r or {}).get("trigger_date") or ""))
            if isinstance(d, str) and d:
                trigger_dates.append(d)
        if len(trigger_dates) < 2:
            return 0.0
        gaps: List[float] = []
        for idx in range(1, len(trigger_dates)):
            gaps.append(float(DateUtils.diff_days(trigger_dates[idx - 1], trigger_dates[idx])))
        return round(sum(gaps) / len(gaps), 1) if gaps else 0.0

    def _get_date_before(self, date_str: str, days: int) -> str:
        try:
            return DateUtils.sub_days(date_str, days)
        except Exception:
            return date_str

    def _enumerate_single_day(
        self,
        tracker: Dict[str, Any],
        current_kline: Dict[str, Any],
        data_of_today: Dict[str, Any],
        *,
        prev_kline: Optional[Dict[str, Any]] = None,
    ):
        fill_pending_buys(
            tracker["pending_buys"],
            tracker["active_opportunities"],
            bar=current_kline,
            sim=self.simulation,
            market_profile=self.market_profile,
            prev_bar=prev_kline,
            stock_status_risk=self.stock_status_risk,
        )
        exit_indices = execute_pending_exits_on_active(
            tracker["active_opportunities"],
            bar=current_kline,
            sim=self.simulation,
            market_profile=self.market_profile,
            prev_bar=prev_kline,
            stock_status_risk=self.stock_status_risk,
        )
        for idx in reversed(exit_indices):
            tracker["active_opportunities"].pop(idx)

        completed_indices = []
        for idx, opportunity in enumerate(tracker["active_opportunities"]):
            if opportunity.check_targets(
                self.simulation,
                current_kline=current_kline,
                goal_config=self.settings.goal,
                market_profile=self.market_profile,
                prev_bar=prev_kline,
                backtest_calendar=self.backtest_calendar,
                stock_status_risk=self.stock_status_risk,
            ):
                completed_indices.append(idx)
        for idx in reversed(completed_indices):
            tracker["active_opportunities"].pop(idx)

        opportunity = self._scan_opportunity_with_data(data_of_today)
        if not opportunity:
            return
        if opportunity.stock:
            opportunity.stock = {**self.stock_info, **opportunity.stock}
        else:
            opportunity.stock = self.stock_info
        opportunity.stock_id = self.stock_id
        opportunity.strategy_name = self.strategy_name
        opportunity.trigger_date = current_kline["date"]
        opportunity.trigger_price = float(current_kline.get("close") or 0.0)
        stamp_stock_status_at_trigger(
            opportunity,
            trade_date=opportunity.trigger_date,
            tier_periods=self.stock_status_risk.tier_periods,
        )
        opportunity.completed_targets = []
        self.opportunity_counter += 1
        opportunity.opportunity_id = str(self.opportunity_counter)
        opportunity.enrich_from_framework(
            strategy_name=self.strategy_name,
            strategy_version="1.0",
            opportunity_id=opportunity.opportunity_id,
        )
        tracker["all_opportunities"].append(opportunity)

        if trade_price_defers_to_next_session(self.simulation.buy_price_model):
            queue_deferred_buy(opportunity, signal_bar=current_kline)
            tracker["pending_buys"].append(opportunity)
            return

        buy_raw = trade_theoretical_price_same_day(
            self.simulation.buy_price_model,
            side="buy",
            bar=current_kline,
            no_next_bar=self.simulation.edges_no_next_bar,
        )
        if buy_raw is None:
            tracker["all_opportunities"].pop()
            return
        buy_price = apply_buy_slippage(buy_raw, self.simulation.slippage_buy_bps)
        opportunity.buy_price = buy_price
        opportunity.buy_date = str(current_kline.get("date") or "")
        stamp_buy_tradability(
            opportunity,
            self.market_profile,
            self.stock_id,
            prev_kline,
            buy_price,
            stock_status_risk=self.stock_status_risk,
            trade_date=opportunity.buy_date,
            exec_bar=current_kline,
        )
        opportunity.signal_phase = ScanSignalPhase.ACTIVE.value
        opportunity.lifecycle = InvestmentLifecycle.OPEN.value
        tracker["active_opportunities"].append(opportunity)

    def _scan_opportunity_with_data(self, data: Dict[str, Any]):
        return self.strategy_instance.scan_opportunity(data, self.settings_dict)

    def _close_all_open_opportunities(
        self,
        tracker: Dict[str, Any],
        last_kline: Dict[str, Any],
    ):
        for opportunity in tracker["active_opportunities"]:
            opportunity.settle(
                self.simulation,
                last_kline=last_kline,
                reason="enumeration_end",
                market_profile=self.market_profile,
                stock_status_risk=self.stock_status_risk,
            )
        tracker["active_opportunities"].clear()

    def _save_stock_results(self, output_dir: str, opportunities: List[Dict[str, Any]]):
        from pathlib import Path

        output_path = Path(output_dir)
        opportunity_rows, target_rows = EnumeratorOutputWriterService.build_stock_rows(
            opportunities=opportunities
        )
        EnumeratorOutputWriterService.write_stock_csv(
            output_dir=output_path,
            stock_id=self.stock_id,
            opportunity_rows=opportunity_rows,
            target_rows=target_rows,
        )


def run_enumeration_payload(job_payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    执行单股或多股 dispatch job。

    多股时在同一子进程内顺序处理，共享 ``ContractCacheManager``。
    """
    stock_ids = job_payload.get("stock_ids")
    if not isinstance(stock_ids, list) or len(stock_ids) <= 1:
        sid = str(job_payload.get("stock_id") or (stock_ids or [None])[0] or "").strip()
        single = dict(job_payload)
        single["stock_id"] = sid
        return OpportunityEnumeratorWorker(single).run()

    ids = [str(s).strip() for s in stock_ids if str(s).strip()]
    if not ids:
        return {"success": False, "bulk": True, "stock_results": [], "error": "empty stock_ids"}

    settings = StrategySettingsView.from_dict(job_payload["settings"])
    actual_start_date = enumeration_actual_start_date(
        str(job_payload["start_date"]),
        settings.min_required_records,
    )

    shared_cache = ContractCacheManager()
    job_batch = StrategyJobContractBatch.hydrate(
        entity_ids=ids,
        settings=settings,
        start=actual_start_date,
        end=str(job_payload["end_date"]),
        contract_cache=shared_cache,
        global_extra_cache=job_payload.get("global_extra_cache"),
        fresh_strategy_cache=True,
    )

    stock_results: List[Dict[str, Any]] = []
    for sid in ids:
        sub_payload = dict(job_payload)
        sub_payload["stock_id"] = sid
        worker: Optional[OpportunityEnumeratorWorker] = None
        try:
            worker = OpportunityEnumeratorWorker(
                sub_payload,
                stock_id=sid,
                contract_cache=shared_cache,
                fresh_strategy_cache=False,
                job_contract_batch=job_batch,
            )
            stock_results.append(worker.run())
        except Exception as exc:
            logger.error(
                "enumeration failed: stock_id=%s, error=%s", sid, exc, exc_info=True
            )
            stock_results.append(
                {
                    "success": False,
                    "stock_id": sid,
                    "stock_name": sid,
                    "opportunity_count": 0,
                    "completed_count": 0,
                    "unfinished_count": 0,
                    "completion_rate": 0.0,
                    "avg_opportunity_interval_days": 0.0,
                    "error": str(exc),
                }
            )
        finally:
            if worker is not None:
                worker.data_manager.clear_working_state()

    return {
        "success": all(bool(r.get("success")) for r in stock_results),
        "bulk": True,
        "stock_results": stock_results,
        "stock_ids": ids,
    }


__all__ = [
    "OpportunityEnumeratorWorker",
    "enumeration_actual_start_date",
    "run_enumeration_payload",
]
