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

"""底仓、补仓记录与待确认订单的 JSON 存储。"""

from __future__ import annotations

import json
import logging as log
import math
from dataclasses import asdict, dataclass, replace
from pathlib import Path
from threading import RLock
from typing import Iterable

from sdk import OrderItem, PositionItem

TERMINAL_STATUSES = {"53", "54", "56", "57"}


@dataclass(frozen=True, slots=True)
class StateItem:
    code: str
    base_order_id: str = ""
    base_qty: int = 0
    base_cost: float = 0.0
    added_order_id: str = ""
    added_num: int = 0
    added_qty: int = 0
    added_cost: float = 0.0


@dataclass(frozen=True, slots=True)
class PendingOrder:
    order_id: str
    code: str
    kind: str
    pre_qty: int


class State:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.lock = RLock()
        self.items: dict[str, StateItem] = {}
        self.pending: dict[str, PendingOrder] = {}
        self.IsModify = False
        self._load()

    def _load(self) -> None:
        """启动时读取 JSON 中的持仓记录和待确认订单。"""
        if self.path.is_file():
            raw = json.loads(self.path.read_text(encoding="utf-8"))
            self.items = {code: StateItem(**item) for code, item in raw["items"].items()}
            self.pending = {code: PendingOrder(**item) for code, item in raw["pending"].items()}

    @classmethod
    def for_strategy(cls, data_dir: str | Path, strategy: str, account_id: str) -> State:
        return cls(Path(data_dir) / f"{strategy}_{account_id}_state.json")

    def get(self, code: str) -> StateItem:
        with self.lock:
            return self.items[code]

    def busy(self, code: str) -> bool:
        with self.lock:
            return code in self.pending

    def new_order(self, order: PendingOrder) -> bool:
        """提交前保存待确认订单，同一证券已有 pending 时跳过。"""
        with self.lock:
            if order.code in self.pending:
                return False
            was_modified = self.IsModify
            self.pending[order.code] = order
            self.IsModify = True
            try:
                self.save()
            except Exception:
                del self.pending[order.code]
                self.IsModify = was_modified
                raise
            return True

    def reconcile(self, positions: Iterable[PositionItem], orders: list[OrderItem]) -> None:
        """核对 pending 的订单结果，再导入/核对实际持仓。"""
        by_id: dict[str, dict[str, OrderItem]] = {}
        for order in orders:
            if order.local_order_id:
                by_id.setdefault(order.local_order_id, {})[order.id] = order

        with self.lock:
            for code, pending in list(self.pending.items()):
                rows = [order for order in by_id.get(pending.order_id, {}).values()
                        if order.code == code and order.side == "BUY"]
                if not rows:
                    del self.pending[code]
                    self.IsModify = True
                    log.info("[状态] 清理无对应订单的 pending，代码=%s，订单=%s", code, pending.order_id)
                    continue
                filled = _finished_fill(rows, pending.pre_qty)
                if filled is None:
                    continue
                qty, amount = filled
                if qty:
                    item = self.items.get(code, StateItem(code))
                    if pending.kind == "base":
                        item = replace(item, base_order_id=pending.order_id,
                                       base_qty=qty, base_cost=amount / qty)
                    else:
                        total_qty = item.added_qty + qty
                        total_amount = item.added_qty * item.added_cost + amount
                        item = replace(item, added_order_id=pending.order_id,
                                       added_num=item.added_num + 1,
                                       added_qty=total_qty, added_cost=total_amount / total_qty)
                    self.items[code] = item
                del self.pending[code]
                self.IsModify = True
                log.info("[状态] 对账结束，代码=%s，订单=%s，计划=%d，成交=%d",
                         code, pending.order_id, pending.pre_qty, qty)

            for position in positions:
                code = position.stock_code
                if not code or position.volume <= 0:
                    continue
                if code not in self.items and code not in self.pending:
                    self.items[code] = StateItem(code=code, base_order_id=position.trade_id,
                                                base_qty=position.volume, base_cost=position.open_price)
                    self.IsModify = True
                elif code in self.items:
                    item = self.items[code]
                    if item.base_qty + item.added_qty != position.volume:
                        log.info("[状态] 持仓数量差异，代码=%s，记录=%d，实际=%d",
                                 code, item.base_qty + item.added_qty, position.volume)
            self.save()

    def save(self) -> None:
        with self.lock:
            if not self.IsModify:
                return
            self.path.parent.mkdir(parents=True, exist_ok=True)
            temporary = self.path.with_suffix(self.path.suffix + ".tmp")
            payload = dict(items={code: asdict(item) for code, item in self.items.items()},
                           pending={code: asdict(item) for code, item in self.pending.items()})
            temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8")
            temporary.replace(self.path)
            self.IsModify = False


def _finished_fill(orders: list[OrderItem], pre_qty: int) -> tuple[int, float] | None:
    """订单全部结束且成交金额完整后一次记账，部分成交撤单也按实计入。"""
    if sum(order.volume for order in orders) != pre_qty:
        return None
    qty, amount = 0, 0.0
    for order in orders:
        if (order.status not in TERMINAL_STATUSES
                or (order.status == "56" and order.traded_volume != order.volume)):
            return None
        if order.traded_volume:
            value = order.trade_amount if order.trade_amount > 0 else order.trade_price * order.traded_volume
            if not math.isfinite(value) or value <= 0:
                return None
            qty += order.traded_volume
            amount += value
    return qty, amount