big-qmt/py-client/strategy/trend/state.py

"""趋势策略持仓状态的内存管理与 JSON 持久化。"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from threading import Lock
from typing import Iterable

from sdk import OrderItem, PositionItem


# 委托状态：无操作、处理中、已完成。
STATUS_NONE = ""
STATUS_ING = "ING"
STATUS_OK = "OK"
STATUS_FAILED = "FAILED"
STATUS_CANCELED = "CANCELED"
STATUS_UNKNOWN = "UNKNOWN"


@dataclass(slots=True)
class StateItem:
    """单只证券的底仓和补仓状态。"""

    # 证券代码。
    code: str

    # 底仓订单、数量、成本和处理状态。
    base_order_id: str = ""
    base_qty: int = 0
    base_cost: float = 0.0
    base_status: str = STATUS_NONE

    # 补仓订单、补仓次数、数量、成本和处理状态。
    added_order_id: str = ""
    added_num: int = 0
    added_qty: int = 0
    added_cost: float = 0.0
    added_status: str = STATUS_NONE


class State:
    """线程安全的策略状态存储。

    状态以内存字典提供快速访问，并通过临时文件替换的方式写入 JSON，
    防止程序在写入过程中退出而破坏原状态文件。
    """

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.lock = Lock()
        self.items = self._load()

    @classmethod
    def for_strategy(
        cls,
        data_dir: str | Path,
        strategy: str,
        account_id: str,
    ) -> "State":
        """根据数据目录、策略名称和账户生成独立状态文件。"""
        state_path = Path(data_dir) / f"{strategy}_{account_id}_state.json"
        return cls(state_path)

    @property
    def codes(self) -> list[str]:
        """返回当前已经接管的全部证券代码快照。"""
        with self.lock:
            return list(self.items)

    @property
    def unresolved_codes(self) -> list[str]:
        """返回存在处理中或未知订单状态的证券代码快照。"""
        with self.lock:
            return [
                code
                for code, item in self.items.items()
                if _has_unresolved_order(item)
            ]

    def get(self, code: str) -> StateItem:
        """获取指定证券的状态；不存在时抛出 KeyError。"""
        with self.lock:
            return self.items[code]

    def set(self, item: StateItem) -> None:
        """新增或覆盖一只证券的状态。"""
        with self.lock:
            self.items[item.code] = item

    def delete(self, code: str) -> bool:
        """删除已终结的证券状态，并返回是否实际删除。"""
        with self.lock:
            item = self.items.get(code)
            if item is not None and _has_unresolved_order(item):
                return False
            return self.items.pop(code, None) is not None

    def has_unresolved_order(self, code: str) -> bool:
        """判断证券是否存在必须阻止自动下单的未决订单。"""
        with self.lock:
            item = self.items.get(code)
            return item is not None and _has_unresolved_order(item)

    def sync_positions(self, positions: Iterable[PositionItem]) -> None:
        """把尚未接管的真实持仓初始化为已完成底仓。

        无证券代码、无持仓数量或成本无效的记录会被忽略。同步结束后
        立即保存，确保首次接管的持仓在程序重启后仍可恢复。
        """
        known_codes = set(self.codes)
        for position in positions:
            if (
                not position.stock_code
                or position.volume <= 0
                or position.open_price <= 0
                or position.stock_code in known_codes
            ):
                continue

            self.set(
                StateItem(
                    code=position.stock_code,
                    base_qty=position.volume,
                    base_cost=position.open_price,
                    base_status=STATUS_OK,
                )
            )
            known_codes.add(position.stock_code)

        self.save()

    def reconcile(
        self,
        positions: Iterable[PositionItem],
        orders: list[OrderItem],
        deals: list[dict[str, str]],
    ) -> None:
        """用真实持仓、委托和成交恢复本地状态，不增加持久化字段。"""
        position_list = list(positions)
        self.sync_positions(position_list)
        position_codes = {
            item.stock_code for item in position_list if item.volume > 0
        }
        for code in list(self.codes):
            item = self.get(code)
            item.base_status = _reconcile_leg(
                item.base_order_id, item.base_status, item.base_qty, orders, deals
            )
            item.added_status = _reconcile_leg(
                item.added_order_id, item.added_status, item.added_qty, orders, deals
            )
            self.set(item)

        # Opening orders normally have no position until their first fill.  Order
        # reconciliation must therefore happen before stale state is removed.
        for code in list(self.codes):
            if code not in position_codes:
                self.delete(code)
        self.save()

    def save(self) -> None:
        """将内存状态格式化写入 JSON，并原子替换正式文件。"""
        with self.lock:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            temporary_path = self.path.with_suffix(self.path.suffix + ".tmp")
            payload = {
                code: asdict(item)
                for code, item in self.items.items()
            }
            temporary_path.write_text(
                json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
            temporary_path.replace(self.path)

    def _load(self) -> dict[str, StateItem]:
        """读取已有状态文件；文件不存在时从空状态开始。"""
        try:
            raw = json.loads(self.path.read_text(encoding="utf-8"))
        except FileNotFoundError:
            return {}
        except (OSError, json.JSONDecodeError) as exc:
            raise ValueError(f"[状态] 读取或解析失败: {exc}") from exc

        if not isinstance(raw, dict):
            raise ValueError("[状态] 状态文件根节点必须是 JSON 对象")

        try:
            return {
                code: StateItem(**item)
                for code, item in raw.items()
            }
        except (TypeError, ValueError) as exc:
            raise ValueError(f"[状态] 状态字段无效: {exc}") from exc


def _reconcile_leg(
    local_order_id: str,
    current_status: str,
    expected_qty: int,
    orders: list[OrderItem],
    deals: list[dict[str, str]],
) -> str:
    if current_status not in {STATUS_ING, STATUS_UNKNOWN} or not local_order_id:
        return current_status
    matching_orders = [
        order for order in orders if order.local_order_id == local_order_id
    ]
    order = matching_orders[-1] if matching_orders else None
    if order is None:
        matching_deals = [
            row for row in deals if _matches_local_order(row, local_order_id)
        ]
        dealt = sum(_deal_volume(row) for row in matching_deals)
        if expected_qty > 0 and dealt >= expected_qty:
            return STATUS_OK
        if expected_qty <= 0 and matching_deals:
            return STATUS_OK
        return STATUS_UNKNOWN

    system_order_id = order.id.strip()
    matching_deals = [
        row
        for row in deals
        if (
            system_order_id
            and str(row.get("m_strOrderSysID") or "").strip() == system_order_id
        )
        or (not system_order_id and _matches_local_order(row, local_order_id))
    ]
    dealt = sum(_deal_volume(row) for row in matching_deals)
    traded = max(order.traded_volume, dealt)
    ordered = order.volume or expected_qty
    status = order.status
    if ordered > 0 and traded >= ordered:
        return STATUS_OK
    if status in {"48", "49", "50", "51", "52", "55"}:
        return STATUS_ING
    if status in {"54", "56"}:
        return STATUS_UNKNOWN if traded > 0 else STATUS_CANCELED
    if status in {"57", "58"}:
        return STATUS_UNKNOWN if traded > 0 else STATUS_FAILED
    return STATUS_UNKNOWN


def _has_unresolved_order(item: StateItem) -> bool:
    return item.base_status in {STATUS_ING, STATUS_UNKNOWN} or item.added_status in {
        STATUS_ING,
        STATUS_UNKNOWN,
    }


def _matches_local_order(row: dict[str, str], local_order_id: str) -> bool:
    remark = str(row.get("m_strRemark") or "")
    return remark.split("|", 1)[0] == local_order_id


def _deal_volume(deal: dict[str, str]) -> int:
    for key in ("m_nVolume", "m_nTradeVolume", "m_nVolumeTraded"):
        volume = _as_int(deal.get(key))
        if volume > 0:
            return volume
    return 0


def _as_int(value: object) -> int:
    try:
        return int(value or 0)
    except (TypeError, ValueError):
        return 0