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

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

from __future__ import annotations

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

from sdk import OrderItem, PositionItem
from .order import BUSY_STATUSES, COMPLETED_STATUSES


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


@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)

    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:
            return self.items.pop(code, None) is not None


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

        无证券代码、无持仓数量或成本无效的记录会被忽略。同步结束后
        立即保存，确保首次接管的持仓在程序重启后仍可恢复。
        """
        known_codes = set(self.codes)
        imported = 0
        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(
                    base_order_id=position.trade_id,
                    code=position.stock_code,
                    base_qty=position.volume,
                    base_cost=round(position.open_price, 2),
                    base_status=STATUS_OK,
                )
            )
            known_codes.add(position.stock_code)
            imported += 1

        self.save()
        if imported:
            log.info("[状态] 导入持仓=%d，状态总数=%d", imported, len(known_codes))

    def reconcile(
        self,
        positions: Iterable[PositionItem],
        orders: list[OrderItem],
    ) -> None:
        """用真实持仓和委托恢复本地状态；拆分订单全部完成才算完成。"""
        position_list = list(positions)
        self.sync_positions(position_list)
        position_codes = {
            item.stock_code for item in position_list if item.volume > 0
        }
        orders_by_local_id: dict[str, list[OrderItem]] = {}
        for order in orders:
            if order.local_order_id:
                orders_by_local_id.setdefault(order.local_order_id, []).append(order)

        for code in list(self.codes):
            item = self.get(code)
            for order_id_attr, status_attr in (
                ("base_order_id", "base_status"),
                ("added_order_id", "added_status"),
            ):
                local_order_id = getattr(item, order_id_attr)
                current_status = getattr(item, status_attr)
                if current_status != STATUS_ING or not local_order_id:
                    continue

                matching_orders = orders_by_local_id.get(local_order_id)
                status = _order_status(matching_orders)
                if status != current_status:
                    log.info("[状态] %s 订单=%s，状态=%s->%s", code, local_order_id, current_status, status)
                setattr(item, status_attr, status)
            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:
                item = self.get(code)
                if any(
                    order_id and order_id in orders_by_local_id
                    for order_id in (item.base_order_id, item.added_order_id)
                ):
                    continue
                if self.delete(code):
                    log.info("[状态] 已移除持仓状态，代码=%s", 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 _order_status(orders: list[OrderItem] | None) -> str:
    """将柜台订单简化为无状态、处理中或已成交。"""
    if not orders:
        return STATUS_NONE
    statuses = {order.status for order in orders}
    if statuses <= COMPLETED_STATUSES:
        return STATUS_OK
    if statuses <= BUSY_STATUSES | COMPLETED_STATUSES:
        return STATUS_ING
    return STATUS_NONE