big-qmt/docs/arch/state.py

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

from __future__ import annotations

import json
import logging as log
import time
from dataclasses import asdict, dataclass, field, replace
from pathlib import Path
from threading import RLock
from typing import Any

from sdk import OrderItem, PositionItem

PENDING_TIME_OUT = 3600


@dataclass(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(slots=True)
class PendingOrder:
    order_id: str
    code: str
    kind: str
    pre_qty: int
    submit_at: int = field(default_factory=lambda: int(time.time()))


class State:
    def __init__(self, data_dir: str | Path, strategy: str, account_id: str) -> None:
        self.path = Path(data_dir) / f"{strategy}_{account_id}_state.json"
        self.lock = RLock()
        self.items: dict[str, StateItem] = {}
        self.pending: list[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 = [PendingOrder(**item) for item in raw["pending"]]
            now = int(time.time())
            for pending in list(self.pending):
                if now - pending.submit_at >= PENDING_TIME_OUT:
                    self.pending.remove(pending)
                    self.IsModify = True
                    log.info("[状态] 清理超时 pending，代码=%s，订单=%s", pending.code, pending.order_id)
            self.save()


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

    def new_order(self, order: PendingOrder) -> None:
        """提交前保存待确认订单，同一证券已有 pending 时跳过。"""
        with self.lock:
            if order:
                self.pending.append(order)
                self.IsModify = True
                self.save()

    def merged_order(self, orders: list[OrderItem]) -> dict[str, dict[str, Any]]:
        """按证券代码、本地订单号合并，返回成交数量、金额、均价和状态。"""
        merged: dict[str, dict[str, Any]] = {}
        for order in orders:
            if not order.local_order_id:
                continue

            if not order.local_order_id in merged.keys():
                strStatus = "ING"
                if order.status == "56":
                    strStatus = "OK"
                merged[order.local_order_id] = {
                    "code":order.code,
                    "qty": order.traded_volume,
                    "cost": order.trade_price,
                    "status": order.status,
                    "merged_status":strStatus
                }
                continue

            old = merged[order.local_order_id]
            strStatus = "ING"
            if old["status"] == order.status == "56":
                strStatus = "OK"

            totalQty = old["qty"]+order.traded_volume
            # 合计成交数量为零，跳过，避免除零。
            if totalQty == 0:
                continue
            # 有成交数量但缺少成交金额，跳过，避免拉低成本。
            if order.traded_volume > 0 and not order.trade_amount:
                continue
            cost = ((old["qty"]*old["cost"])+order.trade_amount) / totalQty
            merged[order.local_order_id]["qty"] = totalQty
            merged[order.local_order_id]["cost"] = cost
            merged[order.local_order_id]["merged_status"] = strStatus
            
        return merged

    def reconcile(self, positions: list[PositionItem], orders: list[OrderItem]) -> None:
        """合并订单 → 对齐 pending 和持仓 → 保存 JSON。"""
        # 1. 合并同一本地订单的成交数据。
        merged = self.merged_order(orders)

        with self.lock:
            # 2. 按本地订单号查找合并结果，核对证券代码后写入底仓或补仓。
            # 只留下尚未完成确认的订单，避免循环中反复查找、删除列表元素。
            remaining: list[PendingOrder] = []
            for pending in self.pending:
                code = pending.code
                result = merged.get(pending.order_id)
                if result is None or result["code"] != code:
                    # 按现有规则：本轮快照中没有对应订单，就清理 pending。
                    self.IsModify = True
                    log.info("[状态] 清理无对应订单的 pending，代码=%s，订单=%s", code, pending.order_id)
                    continue
                if result["merged_status"] != "OK" or result["qty"] != pending.pre_qty:
                    # 未全部成功，或拆单快照的数量尚未齐全，留到下轮确认。
                    remaining.append(pending)
                    continue
                qty, cost = result["qty"], result["cost"]
                # 只记录实际成交。
                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=cost)
                    else:
                        # 补仓：次数加一，数量累加，成本按成交数量加权。
                        total_qty = item.added_qty + qty
                        total_amount = item.added_qty * item.added_cost + qty * cost
                        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
                self.IsModify = True
                log.info("[状态] 对账结束，代码=%s，订单=%s，计划=%d，成交=%d",
                         code, pending.order_id, pending.pre_qty, qty)

            self.pending = remaining

            # 未记录且没有待确认订单的持仓，作为首次接管的底仓导入。
            pending_codes = {pending.code for pending in remaining}
            for position in positions:
                code = position.stock_code
                if not code or position.volume <= 0:
                    continue
                item = self.items.get(code)
                if item is None:
                    if code not in pending_codes:
                        self.items[code] = StateItem(code=code, base_order_id=position.trade_id,
                                                    base_qty=position.volume, base_cost=position.open_price)
                        self.IsModify = True
                    continue

                # 已有记录只报告数量差异，不用总持仓覆盖底仓/补仓的划分。
                recorded_qty = item.base_qty + item.added_qty
                if recorded_qty != position.volume:
                    log.info("[状态] 持仓数量差异，代码=%s，记录=%d，实际=%d",
                             code, recorded_qty, position.volume)

            # 3. 本轮统一保存；没有修改时 save() 不写文件。
            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=[asdict(item) for item in self.pending])
            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