NewStock/qmt/api_server.py

# coding:utf-8
import os
import datetime
from typing import Optional, List, Dict, Any

from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

# 导入新的管理器类
from qmt_engine import MultiEngineManager, TerminalStatus


# ================= Pydantic模型 =================

class TerminalStatusModel(BaseModel):
    """单个终端状态模型"""
    qmt_id: str
    alias: str
    account_id: str
    is_connected: bool
    last_heartbeat: str

class StatusResponse(BaseModel):
    """全局状态响应模型"""
    running: bool
    start_time: str
    terminals: List[TerminalStatusModel]

class PositionsResponse(BaseModel):
    """持仓响应模型"""
    # 按 qmt_id 分组的实盘持仓
    real_positions: Dict[str, List[Dict[str, Any]]]
    # 按策略名分组的虚拟持仓
    virtual_positions: Dict[str, Dict[str, str]]

class LogsResponse(BaseModel):
    """日志响应模型"""
    logs: List[str]


# ================= FastAPI应用 =================

class QMTAPIServer:
    """多终端 QMT API服务器"""
    
    def __init__(self, manager: MultiEngineManager):
        self.app = FastAPI(title="QMT Multi-Terminal Monitor")
        self.manager = manager
        self._setup_middleware()
        self._setup_routes()
        
    def _setup_middleware(self):
        """设置中间件"""
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
    def _setup_routes(self):
        """设置路由"""
        
        @self.app.get("/", summary="仪表盘页面")
        async def read_root():
            """返回仪表盘HTML页面"""
            if os.path.exists("dashboard.html"):
                return FileResponse("dashboard.html")
            return {"error": "Dashboard not found"}
        
        @self.app.get("/api/status", response_model=StatusResponse, summary="获取所有终端状态")
        def get_status():
            """获取所有 QMT 终端的连接状态"""
            terminal_data = self.manager.get_all_status()
            
            terminals = [
                TerminalStatusModel(
                    qmt_id=t.qmt_id,
                    alias=t.alias,
                    account_id=t.account_id,
                    is_connected=t.is_connected,
                    last_heartbeat=t.last_heartbeat
                ) for t in terminal_data
            ]
            
            return StatusResponse(
                running=self.manager.is_running,
                start_time=self.manager.start_time,
                terminals=terminals
            )
        
        @self.app.get("/api/positions", response_model=PositionsResponse, summary="获取持仓信息")
        def get_positions():
            """汇总所有终端的实盘持仓和所有策略的虚拟持仓"""
            real_pos_data = {}
            virtual_pos_data = {}

            # 1. 遍历所有终端单元获取实盘持仓
            for qmt_id, unit in self.manager.units.items():
                positions = []
                if unit.callback and unit.callback.is_connected:
                    try:
                        xt_pos = unit.xt_trader.query_stock_positions(unit.acc_obj)
                        if xt_pos:
                            positions = [
                                {
                                    "code": p.stock_code,
                                    "volume": p.volume,
                                    "can_use": p.can_use_volume,
                                    "market_value": round(p.market_value, 2)
                                } for p in xt_pos if p.volume > 0
                            ]
                    except:
                        pass
                real_pos_data[qmt_id] = positions

            # 2. 遍历所有策略获取虚拟持仓
            for s_name in self.manager.config.get('strategies', {}).keys():
                v_data = self.manager.pos_manager.get_all_virtual_positions(s_name)
                virtual_pos_data[s_name] = v_data

            return PositionsResponse(
                real_positions=real_pos_data,
                virtual_positions=virtual_pos_data
            )
        
        @self.app.get("/api/logs", response_model=LogsResponse, summary="获取系统日志")
        def get_logs(lines: int = Query(50, ge=1, le=1000)):
            """获取最近的系统运行日志"""
            logs = self.manager.get_logs(lines)
            return LogsResponse(logs=logs)
        
        @self.app.get("/api/health", summary="健康检查")
        def health_check():
            """健康检查：只要有一个终端在线即视为正常"""
            terminal_data = self.manager.get_all_status()
            any_connected = any(t.is_connected for t in terminal_data)
            
            if self.manager.is_running and any_connected:
                return {"status": "healthy", "timestamp": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
            else:
                return {"status": "unhealthy", "reason": "No terminals connected", "timestamp": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
        
    def get_app(self) -> FastAPI:
        """获取FastAPI应用实例"""
        return self.app


# ================= 辅助函数 =================

def create_api_server(manager: MultiEngineManager) -> FastAPI:
    """创建API服务器入口"""
    server = QMTAPIServer(manager)
    return server.get_app()