MomentumLab/app/akshare/main_api.py

"""
AKShare HTTP API 服务
提供股票数据接口和数据同步功能
"""
import os
from datetime import datetime
from typing import Optional, List

import akshare as ak
import pandas as pd
from fastapi import FastAPI, HTTPException, Query, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware

from database import check_connection, init_db, get_db, Stock, StockKLine
from data_sync import (
    sync_all_stocks, sync_realtime_quotes, sync_stock_kline,
    sync_sectors, sync_sector_quotes, sync_market_indices,
    sync_all_klines, sync_all, sync_daily
)

app = FastAPI(
    title="AKShare HTTP API",
    description="AKShare 数据接口 HTTP 服务 - 支持数据同步到 MySQL",
    version="2.0.0"
)

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def dataframe_to_records(df: pd.DataFrame) -> List[dict]:
    """将 DataFrame 转换为可 JSON 序列化的记录列表"""
    if df is None or df.empty:
        return []
    df = df.replace({pd.NaT: None})
    df = df.where(pd.notnull(df), None)
    return df.to_dict('records')


# ==================== 健康检查 ====================

@app.get("/")
async def root():
    """健康检查"""
    db_status = "connected" if check_connection() else "disconnected"
    return {
        "status": "healthy",
        "service": "AKShare HTTP API",
        "version": "2.0.0",
        "database": db_status,
        "timestamp": datetime.now().isoformat()
    }


@app.get("/health")
async def health_check():
    """详细健康检查"""
    db_ok = check_connection()
    return {
        "status": "healthy" if db_ok else "unhealthy",
        "database": "connected" if db_ok else "disconnected",
        "timestamp": datetime.now().isoformat()
    }


# ==================== 原始 AKShare 接口 ====================

@app.get("/stock_zh_a_spot")
async def stock_zh_a_spot():
    """
    获取 A 股实时行情数据
    """
    try:
        df = ak.stock_zh_a_spot_em()
        records = dataframe_to_records(df)
        mapped_records = []
        for record in records:
            mapped_records.append({
                "code": record.get("代码"),
                "name": record.get("名称"),
                "price": record.get("最新价", 0),
                "change": record.get("涨跌额", 0),
                "change_percent": record.get("涨跌幅", 0),
                "volume": record.get("成交量", 0),
                "turnover": record.get("成交额", 0),
                "open": record.get("开盘价", 0),
                "high": record.get("最高价", 0),
                "low": record.get("最低价", 0),
                "pre_close": record.get("昨收", 0),
                "turnover_rate": record.get("换手率", 0),
                "amplitude": record.get("振幅", 0),
                "market_cap": record.get("总市值", 0),
                "pe": record.get("市盈率-动态", 0),
                "pb": record.get("市净率", 0),
                "industry": record.get("行业", ""),
            })
        return mapped_records
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取数据失败: {str(e)}")


@app.get("/stock_zh_a_hist")
async def stock_zh_a_hist(
    symbol: str = Query(..., description="股票代码，如 000001"),
    period: str = Query("daily", description="周期: daily/weekly/monthly"),
    start_date: Optional[str] = Query(None, description="开始日期 YYYYMMDD"),
    end_date: Optional[str] = Query(None, description="结束日期 YYYYMMDD"),
    adjust: str = Query("qfq", description="复权方式: qfq-前复权, hfq-后复权, 不复权")
):
    """
    获取 A 股历史 K 线数据
    """
    try:
        period_map = {
            "daily": "daily",
            "weekly": "weekly",
            "monthly": "monthly"
        }
        ak_period = period_map.get(period, "daily")
        
        adjust_map = {
            "qfq": "qfq",
            "hfq": "hfq",
            "": ""
        }
        ak_adjust = adjust_map.get(adjust, "qfq")
        
        df = ak.stock_zh_a_hist(
            symbol=symbol,
            period=ak_period,
            start_date=start_date or "19700101",
            end_date=end_date or "20500101",
            adjust=ak_adjust
        )
        records = dataframe_to_records(df)
        
        mapped_records = []
        for record in records:
            mapped_records.append({
                "date": record.get("日期"),
                "open": record.get("开盘", 0),
                "high": record.get("最高", 0),
                "low": record.get("最低", 0),
                "close": record.get("收盘", 0),
                "volume": record.get("成交量", 0),
                "turnover": record.get("成交额", 0),
                "amplitude": record.get("振幅", 0),
                "change": record.get("涨跌幅", 0),
                "change_amount": record.get("涨跌额", 0),
                "turnover_rate": record.get("换手率", 0),
            })
        return mapped_records
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取数据失败: {str(e)}")


@app.get("/stock_zh_index_spot")
async def stock_zh_index_spot():
    """
    获取股票指数实时行情
    """
    try:
        df = ak.index_zh_a_spot_em()
        records = dataframe_to_records(df)
        mapped_records = []
        for record in records:
            mapped_records.append({
                "code": record.get("代码"),
                "name": record.get("名称"),
                "price": record.get("最新价", 0),
                "change": record.get("涨跌额", 0),
                "change_percent": record.get("涨跌幅", 0),
                "volume": record.get("成交量", 0),
                "turnover": record.get("成交额", 0),
                "open": record.get("开盘价", 0),
                "high": record.get("最高价", 0),
                "low": record.get("最低价", 0),
                "pre_close": record.get("昨收", 0),
            })
        return mapped_records
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取数据失败: {str(e)}")


@app.get("/stock_sector_spot")
async def stock_sector_spot():
    """
    获取板块行情数据
    """
    try:
        df = ak.stock_board_industry_name_em()
        records = dataframe_to_records(df)
        mapped_records = []
        for record in records:
            mapped_records.append({
                "code": record.get("代码"),
                "name": record.get("名称"),
                "change_percent": record.get("涨跌幅", 0),
            })
        return mapped_records
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取数据失败: {str(e)}")


# ==================== 数据库查询接口 ====================

@app.get("/db/stocks")
async def db_stocks(
    limit: int = Query(100, description="限制条数"),
    offset: int = Query(0, description="偏移量")
):
    """从数据库获取股票列表"""
    try:
        with get_db() as db:
            stocks = db.query(Stock).offset(offset).limit(limit).all()
            return [
                {
                    "code": s.code,
                    "name": s.name,
                    "pe": s.pe,
                    "pb": s.pb,
                    "sector": s.sectorCode,
                }
                for s in stocks
            ]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@app.get("/db/klines/{symbol}")
async def db_klines(
    symbol: str,
    period: str = Query("day", description="周期: day/week/month"),
    limit: int = Query(60, description="限制条数")
):
    """从数据库获取K线数据"""
    try:
        with get_db() as db:
            klines = db.query(StockKLine).filter(
                StockKLine.stockCode == symbol,
                StockKLine.period == period
            ).order_by(StockKLine.date.desc()).limit(limit).all()
            
            return [
                {
                    "date": k.date.strftime("%Y-%m-%d"),
                    "open": k.open,
                    "high": k.high,
                    "low": k.low,
                    "close": k.close,
                    "volume": k.volume,
                }
                for k in reversed(klines)
            ]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


# ==================== 数据同步接口 ====================

@app.post("/sync/stocks")
async def api_sync_stocks(background_tasks: BackgroundTasks):
    """同步股票列表（后台任务）"""
    background_tasks.add_task(sync_all_stocks)
    return {"message": "股票列表同步已启动", "status": "running"}


@app.post("/sync/quotes")
async def api_sync_quotes(background_tasks: BackgroundTasks):
    """同步实时行情（后台任务）"""
    background_tasks.add_task(sync_realtime_quotes)
    return {"message": "实时行情同步已启动", "status": "running"}


@app.post("/sync/klines")
async def api_sync_klines(
    background_tasks: BackgroundTasks,
    days: int = Query(365, description="同步天数"),
    max_stocks: Optional[int] = Query(None, description="最大股票数量")
):
    """同步K线数据（后台任务）"""
    background_tasks.add_task(sync_all_klines, days=days, max_stocks=max_stocks)
    return {"message": f"K线同步已启动（{days}天）", "status": "running"}


@app.post("/sync/all")
async def api_sync_all(
    background_tasks: BackgroundTasks,
    quick: bool = Query(False, description="快速模式（只同步少量数据）")
):
    """执行全量同步（后台任务）"""
    background_tasks.add_task(sync_all, quick=quick)
    return {"message": "全量同步已启动", "status": "running", "quick": quick}


@app.post("/sync/daily")
async def api_sync_daily(background_tasks: BackgroundTasks):
    """执行每日增量同步（后台任务）"""
    background_tasks.add_task(sync_daily)
    return {"message": "每日增量同步已启动", "status": "running"}


@app.get("/sync/status")
async def sync_status():
    """获取同步状态"""
    db_ok = check_connection()
    
    # 获取统计数据
    stats = {}
    try:
        with get_db() as db:
            from database import StockQuote, Sector, MarketIndex
            stats["stocks"] = db.query(Stock).count()
            stats["quotes"] = db.query(StockQuote).count()
            stats["klines"] = db.query(StockKLine).count()
            stats["sectors"] = db.query(Sector).count()
            stats["indices"] = db.query(MarketIndex).count()
    except Exception as e:
        stats["error"] = str(e)
    
    return {
        "database_connected": db_ok,
        "timestamp": datetime.now().isoformat(),
        "stats": stats
    }


# ==================== 初始化 ====================

@app.on_event("startup")
async def startup_event():
    """启动时初始化数据库"""
    print("正在初始化数据库...")
    if check_connection():
        init_db()
        print("数据库初始化完成")
    else:
        print("警告：数据库连接失败")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)