sanguo_vnpy_v2/sanguo_api/routes.py

"""
FastAPI routes for Sanguo Quant API
"""
import os
from fastapi import APIRouter, HTTPException, Depends, WebSocket, Query, Header
from fastapi.responses import FileResponse
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, get_strategy_class
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"""
    cls = get_strategy_class(req.strategy)
    if cls is None:
        raise HTTPException(status_code=400, detail=f"未知策略: {req.strategy}")
    tid = await get_orchestrator().submit_cta(
        strategy_class=cls,
        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"""
    cls = get_strategy_class(req.strategy)
    if cls is None:
        raise HTTPException(status_code=400, detail=f"未知策略: {req.strategy}")
    tid = await get_orchestrator().submit_optimize(
        strategy_class=cls,
        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")
    pool = get_orchestrator().pool
    stage = pool.get_stage(task_id)
    task = pool.get_task(task_id)
    return {
        "task_id": task_id,
        "status": s.value if hasattr(s, "value") else str(s),
        "stage": stage or "",
        "error_msg": task.error_msg if task else None,
    }


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


# ===== Factor (投研) endpoints (S2) =====

@router.get("/factor/list", dependencies=[Depends(verify_token)])
def factor_list():
    """List registered factors for the UI dropdown."""
    from sanguo_factor.registry import list_factors
    items = [{"name": f["name"], "category": f.get("category", "")} for f in list_factors()]
    return {"factors": items}


@router.get("/task/{task_id}/ic-summary", dependencies=[Depends(verify_token)])
def ic_summary(task_id: str):
    """Factor IC summary (mean/std/icir/t_stat per period)."""
    r = get_orchestrator().get_raw_result(task_id)
    if r is None:
        raise HTTPException(status_code=404, detail="result not ready")
    ic = getattr(r, "ic_summary", None)
    if ic is None:
        raise HTTPException(status_code=404, detail="no ic_summary (not a factor result?)")
    return {"task_id": task_id, "ic_summary": ic}


@router.get("/task/{task_id}/report/{factor}")
def factor_report(task_id: str, factor: str, token: str = Query(...)):
    """Serve the alphalens tears HTML report (token via query for iframe use)."""
    try:
        verify_token_impl(token)
    except Exception:
        raise HTTPException(status_code=401, detail="invalid token")
    r = get_orchestrator().get_raw_result(task_id)
    if r is None:
        raise HTTPException(status_code=404, detail="result not ready")
    paths = getattr(r, "report_paths", {}) or {}
    path = paths.get(factor)
    if not path or not os.path.exists(path):
        raise HTTPException(status_code=404, detail=f"report for {factor} not found")
    return FileResponse(path)