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

"""简单成交账本：订单明确结束后一次记账，JSON 原子保存。"""

from __future__ import annotations

import json
import logging as log
import math
import shutil
import time
from dataclasses import asdict, dataclass, field, 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(slots=True)
class StateItem:
    code: str
    base_order_id: str = ""
    base_pre_qty: int =0
    base_qty: int = 0
    base_cost: float = 0.0
    added_order_id: str = ""
    added_pre_qty: int =0
    added_num: int = 0
    added_qty: int = 0
    added_amount: float = 0.0

    @property
    def added_cost(self) -> float:
        return self.added_amount / self.added_qty if self.added_qty else 0.0


@dataclass(slots=True)
class PendingOrder:
    order_id: str
    kind: str
    expected_qty: int
    submitted_at: float = field(default_factory=time.time)


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._dirty = False
        self._next_warning = 0.0
        self._load()

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

    @property
    def codes(self) -> list[str]:
        with self.lock:
            return list(self.items)

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

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

    def begin(self, order_id: str, code: str, kind: str, expected_qty: int) -> bool:
        if kind not in {"base", "add"} or type(expected_qty) is not int or expected_qty <= 0:
            raise ValueError("订单类型或预开仓数量无效")
        with self.lock:
            if code in self.pending or (kind == "base" and code in self.items):
                return False
            if kind == "add" and code not in self.items:
                raise ValueError("缺少底仓记录")
            dirty = self._dirty
            self.pending[code] = PendingOrder(order_id, kind, expected_qty)
            self._dirty = True
            try:
                self.save()  # 必须成功落盘后才允许发送请求
            except Exception:
                self.pending.pop(code)
                self._dirty = dirty
                raise
            return True

    def reject(self, order_id: str) -> None:
        """仅用于已明确未受理的请求。"""
        with self.lock:
            for code, pending in self.pending.items():
                if pending.order_id == order_id:
                    del self.pending[code]
                    self._dirty = True
                    self.save()
                    return

    def reconcile(self, positions: Iterable[PositionItem], orders: list[OrderItem]) -> None:
        by_id: dict[str, dict[str, OrderItem]] = {}
        for order in orders:
            if order.local_order_id and order.id:
                by_id.setdefault(order.local_order_id, {})[order.id] = order
        with self.lock:
            for position in positions:
                code = position.stock_code
                if (code and code not in self.items and code not in self.pending
                        and position.volume > 0 and math.isfinite(position.open_price)
                        and position.open_price > 0):
                    self.items[code] = StateItem(code, position.volume, position.open_price)
                    self._dirty = True

            now = time.time()
            for code, pending in list(self.pending.items()):
                rows = [o for o in by_id.get(pending.order_id, {}).values()
                        if o.code == code and o.side == "BUY"]
                filled = _finished_fill(rows, pending.expected_qty)
                if filled is None:
                    if now >= self._next_warning and now - pending.submitted_at >= 180:
                        log.warning("[状态] 订单待核查，代码=%s，订单=%s，计划=%d；保留防重",
                                    code, pending.order_id, pending.expected_qty)
                    continue
                qty, amount = filled
                if qty:
                    item = self.items.setdefault(code, StateItem(code))
                    if pending.kind == "base":
                        item.base_qty, item.base_cost = qty, amount / qty
                    else:
                        item.added_num += 1
                        item.added_qty += qty
                        item.added_amount += amount
                del self.pending[code]
                self._dirty = True
                log.info("[状态] 对账结束，代码=%s，订单=%s，计划=%d，成交=%d",
                         code, pending.order_id, pending.expected_qty, qty)
            if now >= self._next_warning:
                self._next_warning = now + 180
            self.save()  # 成交记账与移除 pending 在同一文件中原子提交

    def save(self) -> None:
        with self.lock:
            if not self._dirty:
                return
            self.path.parent.mkdir(parents=True, exist_ok=True)
            temporary = self.path.with_suffix(self.path.suffix + ".tmp")
            payload = dict(version=3, items={k: asdict(v) for k, v in self.items.items()},
                           pending={k: asdict(v) for k, v 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._dirty = False

    def _load(self) -> None:
        try:
            raw = json.loads(self.path.read_text(encoding="utf-8"))
        except FileNotFoundError:
            return
        if raw.get("version") not in {2, 3}:
            raise ValueError("仅支持 v2/v3 状态文件；旧格式请先核实转换，原文件未修改")
        self.items = {k: StateItem(**v) for k, v in raw["items"].items()}
        if raw["version"] == 3:
            self.pending = {k: PendingOrder(**v) for k, v in raw["pending"].items()}
            return
        # v2 的未决订单可能已经增量记账，先撤回这部分，结束时再完整记一次。
        for old in raw["pending"].values():
            code = old["code"]
            if code in self.pending:
                raise ValueError(f"{code} 存在多个旧未决订单，请核查；原文件未修改")
            item = self.items[code]
            qty, amount = old.get("applied_qty", 0), old.get("applied_amount", 0.0)
            if old["kind"] == "add":
                item.added_qty -= qty
                item.added_amount -= amount
                item.added_num -= int(old.get("counted", False))
            else:
                base_amount = item.base_qty * item.base_cost - amount
                item.base_qty -= qty
                item.base_cost = base_amount / item.base_qty if item.base_qty else 0.0
                if not item.base_qty and not item.added_qty:
                    del self.items[code]
            self.pending[code] = PendingOrder(old["order_id"], old["kind"], old["expected_qty"],
                                              old.get("submitted_at", time.time()))
        backup = self.path.with_suffix(self.path.suffix + ".v2.bak")
        if not backup.exists():
            shutil.copy2(self.path, backup)
        self._dirty = True
        self.save()


def _finished_fill(orders: list[OrderItem], expected_qty: int) -> tuple[int, float] | None:
    """仅完整终态快照可记账；缺项、未成交完或金额未知均继续等待。"""
    if not orders or expected_qty <= 0 or sum(o.volume for o in orders) != expected_qty:
        return None
    qty, amount = 0, 0.0
    for order in orders:
        if (order.status not in TERMINAL_STATUSES or not 0 <= order.traded_volume <= order.volume
                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