zhixing-system/zhixing-server/src/zhixing_server/modules/selection/application/run.py

"""Application orchestration for persisted whole-universe B1 runs."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import date
from typing import Literal, Protocol

from ..domain.models import SelectionEvaluation
from ..domain.runs import (
    SelectionExecutionSource,
    SelectionRerunRequired,
    SelectionResultQuery,
    SelectionRun,
    SelectionRunInProgress,
    SelectionRunItem,
    SelectionRunStatus,
    SelectionRunStore,
    SelectionUniverseReader,
)
from .evaluate import EvaluateZhixingB1

logger = logging.getLogger(__name__)
StrategyName = Literal["zhixing_b1"]
_FAILURE_STATUSES = {"insufficient_history", "missing_target_bar", "data_error"}


class SelectionEvaluator(Protocol):
    """Minimal single-stock evaluator required by the batch orchestrator."""

    def execute(self, ts_code: str, target_trade_date: date) -> SelectionEvaluation: ...


@dataclass(frozen=True, slots=True)
class PreparedSelectionRun:
    """A claimed run and its immutable market-data source snapshot."""

    run: SelectionRun
    source: SelectionExecutionSource


class RunZhixingB1:
    """Prepare, execute, and query persisted Zhixing B1 result batches."""

    def __init__(
        self,
        reader: SelectionUniverseReader,
        store: SelectionRunStore,
        evaluator: SelectionEvaluator | None = None,
    ) -> None:
        """Inject storage ports and optionally a test evaluator."""

        self.reader = reader
        self.store = store
        self.evaluator = evaluator or EvaluateZhixingB1(reader)

    def prepare(
        self,
        strategy: StrategyName,
        target_trade_date: date,
        *,
        rerun: bool,
    ) -> PreparedSelectionRun:
        """Validate source eligibility before claiming the rerunnable key."""

        source = self.reader.load_execution_source(strategy, target_trade_date)
        run = self.store.prepare_run(
            strategy,
            target_trade_date,
            source,
            rerun=rerun,
        )
        return PreparedSelectionRun(run=run, source=source)

    def execute(self, prepared: PreparedSelectionRun) -> None:
        """Evaluate every eligible stock and converge the persisted run status.

        This method is the boundary used by FastAPI's in-process background
        task.  An unexpected batch-level error is recorded before the worker
        returns so the UI never mistakes a lost worker exception for success.
        """

        evaluated_count = 0
        selected_stock_count = 0
        signal_count = 0
        failed_count = 0
        try:
            for stock in prepared.source.stocks:
                try:
                    evaluation = self.evaluator.execute(
                        stock.ts_code,
                        prepared.source.target_trade_date,
                    )
                except Exception as exc:  # noqa: BLE001 - isolate one stock from the batch
                    logger.exception(
                        "selection_item_failed run_id=%s ts_code=%s",
                        prepared.run.id,
                        stock.ts_code,
                    )
                    evaluation = SelectionEvaluation(
                        ts_code=stock.ts_code,
                        target_trade_date=prepared.source.target_trade_date,
                        status="data_error",
                        reason=_safe_item_error(exc),
                    )
                item = _to_item(stock.ts_code, stock.name, evaluation)
                self.store.record_item(prepared.run.id, item)
                evaluated_count += 1
                selected_stock_count += evaluation.status == "selected"
                signal_count += len(evaluation.signals)
                failed_count += evaluation.status in _FAILURE_STATUSES

            status = _run_status(evaluated_count, failed_count)
            self.store.finish_run(
                prepared.run.id,
                status,
                evaluated_count=evaluated_count,
                selected_stock_count=selected_stock_count,
                signal_count=signal_count,
                failed_count=failed_count,
            )
        except Exception as exc:  # noqa: BLE001 - worker boundary must persist failure state
            logger.exception("selection_run_failed run_id=%s", prepared.run.id)
            try:
                self.store.finish_run(
                    prepared.run.id,
                    "failed",
                    evaluated_count=evaluated_count,
                    selected_stock_count=selected_stock_count,
                    signal_count=signal_count,
                    failed_count=max(failed_count, 1),
                    error_type="batch_error",
                    error_message=str(exc),
                )
            except Exception:  # noqa: BLE001 - preserve the original worker failure
                logger.exception("selection_run_failure_persist_failed run_id=%s", prepared.run.id)

    def get_run(
        self,
        run_id: str,
        *,
        query: SelectionResultQuery | None = None,
    ) -> SelectionRun | None:
        """Read one persisted run for polling."""

        return self.store.get_run(run_id, query=query)

    def get_latest(
        self,
        strategy: StrategyName,
        target_trade_date: date | None = None,
        *,
        query: SelectionResultQuery | None = None,
    ) -> SelectionRun | None:
        """Read the current result by date or the latest result for a strategy."""

        return self.store.get_latest_run(strategy, target_trade_date, query=query)


def _to_item(ts_code: str, name: str, evaluation: SelectionEvaluation) -> SelectionRunItem:
    """Translate a single-stock domain result into a stored item."""

    return SelectionRunItem(
        ts_code=ts_code,
        name=name or (evaluation.signals[0].name if evaluation.signals else ""),
        status=evaluation.status,
        signal_count=len(evaluation.signals),
        reason=evaluation.reason,
        signals=evaluation.signals,
    )


def _run_status(evaluated_count: int, failed_count: int) -> SelectionRunStatus:
    """Map per-stock outcomes into a visible batch status."""

    if failed_count == 0:
        return "success"
    if evaluated_count == 0 or failed_count >= evaluated_count:
        return "failed"
    return "partial_success"


def _safe_item_error(error: Exception) -> str:
    """Keep per-stock failure context readable without persisting tracebacks."""

    return " ".join(str(error).split())[:500] or error.__class__.__name__


__all__ = [
    "PreparedSelectionRun",
    "RunZhixingB1",
    "SelectionRerunRequired",
    "SelectionRunInProgress",
]