""" FastAPI routes for Sanguo Quant API """ from fastapi import APIRouter, HTTPException, Depends, WebSocket, Query, Header from pydantic import BaseModel from .schemas import CtaBacktestRequest, OptimizeRequest, FactorAnalysisRequest from .auth import verify_token as verify_token_impl, verify_password, create_token from .ws import manager from .strategy_registry import list_strategies, strategy_params from .kline import load_kline router = APIRouter() _orchestrator = None _auth_config = {"username": "admin", "password_hash": "", "jwt_secret": "x", "expire_minutes": 60} def set_orchestrator(orch): """Set the global orchestrator instance""" global _orchestrator _orchestrator = orch def get_orchestrator(): """Get the global orchestrator instance""" return _orchestrator def set_auth_config(cfg): """Set authentication configuration""" _auth_config.update(cfg) class LoginRequest(BaseModel): """Login request schema""" username: str password: str async def verify_token(authorization: str | None = Header(None)): """Dependency to verify JWT token from Authorization header""" if authorization is None: raise HTTPException(status_code=401, detail="Missing authorization header") if not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Invalid authorization header format") token = authorization.split(" ")[1] return verify_token_impl(token) @router.post("/auth/login") def login(req: LoginRequest): """Authenticate user and return JWT token""" if req.username != _auth_config["username"] or not verify_password(req.password, _auth_config["password_hash"]): raise HTTPException(status_code=401, detail="用户名或密码错误") return {"token": create_token(req.username)} @router.post("/backtest/cta", dependencies=[Depends(verify_token)]) async def submit_cta(req: CtaBacktestRequest): """Submit CTA backtest task""" tid = await get_orchestrator().submit_cta( strategy_class=req.strategy, symbol=req.symbol, params=req.params, start=req.start, end=req.end, cfg=None ) return {"task_id": tid} @router.post("/backtest/optimize", dependencies=[Depends(verify_token)]) async def submit_optimize(req: OptimizeRequest): """Submit optimization task""" tid = await get_orchestrator().submit_optimize( strategy_class=req.strategy, symbol=req.symbol, grid=req.grid, start=req.start, end=req.end, cfg=None ) return {"task_id": tid} @router.post("/factor/analyze", dependencies=[Depends(verify_token)]) async def submit_factor(req: FactorAnalysisRequest): """Submit factor analysis task""" tid = await get_orchestrator().submit_factor( symbols=req.symbols, factor_names=req.factor_names, start=req.start, end=req.end, cfg=None, output_dir="/tmp/factor" ) return {"task_id": tid} @router.get("/task/{task_id}", dependencies=[Depends(verify_token)]) def get_status(task_id: str): """Get task status""" s = get_orchestrator().get_status(task_id) if s is None: raise HTTPException(status_code=404, detail="task not found") stage = get_orchestrator().pool.get_stage(task_id) return { "task_id": task_id, "status": s.value if hasattr(s, "value") else str(s), "stage": stage or "" } @router.get("/task/{task_id}/result", dependencies=[Depends(verify_token)]) def get_result(task_id: str): """Get task result""" r = get_orchestrator().get_result(task_id) if r is None: raise HTTPException(status_code=404, detail="result not ready") return {"task_id": task_id, "statistics": r.statistics} @router.websocket("/ws/task/{task_id}") async def task_ws(websocket: WebSocket, task_id: str, token: str = Query(...)): """WebSocket endpoint for task status updates""" try: verify_token(token) except Exception: await websocket.close(code=4401) return await websocket.accept() manager.connect(task_id, websocket) try: while True: await websocket.receive_text() # Keep connection alive except Exception: pass finally: manager.disconnect(task_id, websocket) # ===== Backtest UI support endpoints (S1.4 / S1.5) ===== def _df_to_records(df) -> list[dict]: """DataFrame → list[dict] (empty-safe).""" if df is None: return [] if hasattr(df, "empty") and df.empty: return [] if hasattr(df, "to_dict"): return df.to_dict(orient="records") return list(df) @router.get("/strategy/list", dependencies=[Depends(verify_token)]) def strategy_list(): """List available CTA strategies for the UI dropdown.""" return {"strategies": list_strategies()} @router.get("/strategy/{name}/params", dependencies=[Depends(verify_token)]) def strategy_params_route(name: str): """Strategy parameters + defaults for the dynamic form.""" return strategy_params(name) @router.get("/task/{task_id}/equity-curve", dependencies=[Depends(verify_token)]) def equity_curve(task_id: str): r = get_orchestrator().get_result(task_id) if r is None: raise HTTPException(status_code=404, detail="result not ready") return {"task_id": task_id, "equity_curve": _df_to_records(r.equity_curve)} @router.get("/task/{task_id}/daily-pnl", dependencies=[Depends(verify_token)]) def daily_pnl(task_id: str): r = get_orchestrator().get_result(task_id) if r is None: raise HTTPException(status_code=404, detail="result not ready") ec = r.equity_curve if ec is None or (hasattr(ec, "empty") and ec.empty) or "balance" not in ec.columns: return {"task_id": task_id, "daily_pnl": []} import pandas as pd bal = pd.to_numeric(ec["balance"], errors="coerce") pnl = bal.diff().fillna(0.0) return { "task_id": task_id, "daily_pnl": [{"date": str(d), "pnl": float(p)} for d, p in zip(ec["date"], pnl)], } @router.get("/task/{task_id}/trades", dependencies=[Depends(verify_token)]) def trades_route(task_id: str): r = get_orchestrator().get_result(task_id) if r is None: raise HTTPException(status_code=404, detail="result not ready") return {"task_id": task_id, "trades": _df_to_records(r.trades)} @router.get("/kline", dependencies=[Depends(verify_token)]) def kline(symbol: str, start: str, end: str): """Historical daily K-line for the backtest chart.""" try: return {"symbol": symbol, "kline": load_kline(symbol, start, end)} except Exception as e: raise HTTPException(status_code=500, detail=f"kline load failed: {type(e).__name__}: {e}")