sanguo_vnpy_v2/sanguo_api/routes.py

"""
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,
        "symbol": r.symbol,
        "start": r.start,
        "end": r.end,
        "strategy": r.strategy,
        "params": r.params,
        "status": r.status,
    }


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